import os import json import numpy as np from typing import Dict, Any, Optional, Tuple, List from src.connectome.types import ConnectomeGraph, GraphMode, ProvenanceStatus from src.brain.state import BrainState, HomeostaticDrives from src.brain.plasticity import PlasticityEngine from src.brain.eligibility import ( EligibilityState, EligibilityEngine, NeuromodulationConfig, PLASTICITY_MODES, ) from src.compute.vulkan_backend import VulkanComputeEngine from src.compute.cpu_reference import cpu_lif_step class BrainRuntime: """ Central Brain Runtime integrating biological connectome graph, persistent Vulkan GPU acceleration, LIF spiking neural dynamics, synaptic plasticity, homeostatic drives, and biological population mapping. plasticity_mode: "v1_hebbian" legacy three-factor reward rule (compat baseline). "v2_eligibility" persistent eligibility traces + versioned neuromodulatory signal (reward/novelty/prediction-error/social/goal). """ def __init__( self, graph: ConnectomeGraph, use_gpu: bool = True, enable_plasticity: bool = True, seed: int = 42, plasticity_mode: str = "v1_hebbian", neuromod: Optional[NeuromodulationConfig] = None, trace_decay: float = 0.9, prediction_influence: bool = False, prediction_gain: float = 0.5, ): if plasticity_mode not in PLASTICITY_MODES: raise ValueError(f"unknown plasticity_mode {plasticity_mode!r}; " f"expected one of {PLASTICITY_MODES}") self.graph = graph self.use_gpu = use_gpu self.enable_plasticity = enable_plasticity self.seed = seed self.plasticity_mode = plasticity_mode self.prediction_influence = bool(prediction_influence) self.prediction_gain = float(np.clip(prediction_gain, 0.0, 1.0)) self.state = BrainState.create_initial(graph.num_neurons, seed=seed) self.plasticity = PlasticityEngine() self.plasticity_updates = 0 # cumulative synapses updated (observability) self.eligibility: Optional[EligibilityState] = None self.eligibility_engine: Optional[EligibilityEngine] = None if plasticity_mode == "v2_eligibility": self.eligibility = EligibilityState(len(graph.weights), seed=seed) self.eligibility_engine = EligibilityEngine( trace_decay=trace_decay, neuromod=neuromod if neuromod is not None else NeuromodulationConfig()) self.gpu_engine: Optional[VulkanComputeEngine] = None # Lazy GPU weight sync (v4.1, measured: weight readback dominates # rewarded-step cost at scale, e.g. ~0.85ms of ~1.0ms at N=1024). # GPU weights are authoritative; the CPU mirror (self.graph.weights) # is refreshed only at explicit sync points. Telemetry reports # weights_synced so staleness is never silent. self._gpu_weights_dirty = False if self.use_gpu: try: self.gpu_engine = VulkanComputeEngine() self.gpu_engine.load_circuit( graph.row_offsets, graph.col_indices, graph.weights, self.state.membrane_potentials, self.state.spikes ) except Exception as e: print(f"[BrainRuntime] Notice: Vulkan physical initialization not available ({e}). Using verified CPU LIF reference engine.") self.gpu_engine = None self.use_gpu = False # Population mapping derived from biological connectome anatomy N = graph.num_neurons self.sensory_visual_indices = graph.get_population_indices("visual") if len(self.sensory_visual_indices) == 0: self.sensory_visual_indices = np.arange(0, min(64, N), dtype=np.int32) self.sensory_audio_indices = graph.get_population_indices("auditory") if len(self.sensory_audio_indices) == 0: self.sensory_audio_indices = np.arange(min(64, N), min(128, N), dtype=np.int32) self.sensory_olfactory_indices = graph.get_population_indices("olfactory") if len(self.sensory_olfactory_indices) == 0: self.sensory_olfactory_indices = self.sensory_audio_indices self.sensory_memory_indices = graph.get_population_indices("memory_association") if len(self.sensory_memory_indices) == 0: self.sensory_memory_indices = np.arange(min(128, N), min(192, N), dtype=np.int32) motor_all = graph.get_population_indices("motor") if len(motor_all) >= 32: self.motor_speak_indices = motor_all[:16] self.motor_act_indices = motor_all[16:32] else: self.motor_speak_indices = np.arange(max(0, N - 64), max(0, N - 32), dtype=np.int32) self.motor_act_indices = np.arange(max(0, N - 32), N, dtype=np.int32) desc_all = graph.get_population_indices("descending") if len(desc_all) >= 16: self.motor_image_indices = desc_all[:16] self.motor_remember_indices = desc_all[16:32] if len(desc_all) >= 32 else desc_all[:16] else: self.motor_image_indices = np.arange(max(0, N - 96), max(0, N - 64), dtype=np.int32) self.motor_remember_indices = self.motor_image_indices def sync_gpu_weights(self) -> bool: """Refresh the CPU weight mirror from authoritative GPU weights. Returns True if a sync was performed, False if the mirror was current (CPU path, GPU unavailable, or nothing changed). Called automatically by save_snapshot(); call explicitly before reading graph.weights / graph_hash after rewarded GPU steps (experiment manifests, validation, curriculum measurements). """ if not self._gpu_weights_dirty: return False if self.gpu_engine is None or not self.use_gpu: self._gpu_weights_dirty = False return False downloaded = self.gpu_engine.download_weights() if downloaded is not None: self.graph.weights = downloaded self.graph.graph_hash = self.graph.compute_graph_hash() self._gpu_weights_dirty = False return downloaded is not None @property def gpu_weights_dirty(self) -> bool: return self._gpu_weights_dirty def step( self, sensory_inputs: Optional[Dict[str, np.ndarray]] = None, reward: float = 0.0 ) -> Dict[str, Any]: """ Executes one full biological LIF cognitive step: 1. Injects sensory currents into biological populations 2. Dispatches leaky integration, threshold check, spike generation, and refractory clamping 3. Updates homeostatic drives and prediction errors 4. Applies synaptic plasticity if enabled 5. Computes motor readout """ N = self.graph.num_neurons ext_inputs = np.zeros(N, dtype=np.float32) # Inject sensory currents if sensory_inputs: if "visual" in sensory_inputs: v_in = np.asarray(sensory_inputs["visual"], dtype=np.float32).flatten() length = min(len(v_in), len(self.sensory_visual_indices)) ext_inputs[self.sensory_visual_indices[:length]] += v_in[:length] if "audio" in sensory_inputs: a_in = np.asarray(sensory_inputs["audio"], dtype=np.float32).flatten() length = min(len(a_in), len(self.sensory_audio_indices)) ext_inputs[self.sensory_audio_indices[:length]] += a_in[:length] if "olfactory" in sensory_inputs: o_in = np.asarray(sensory_inputs["olfactory"], dtype=np.float32).flatten() length = min(len(o_in), len(self.sensory_olfactory_indices)) ext_inputs[self.sensory_olfactory_indices[:length]] += o_in[:length] if "memory" in sensory_inputs: m_in = np.asarray(sensory_inputs["memory"], dtype=np.float32).flatten() length = min(len(m_in), len(self.sensory_memory_indices)) ext_inputs[self.sensory_memory_indices[:length]] += m_in[:length] prev_spikes = self.state.spikes.copy() pot_in = self.state.membrane_potentials.copy() ref_in = self.state.refractory_steps.copy() # Execute LIF step (Persistent GPU or CPU) if self.gpu_engine is not None and self.use_gpu: new_pot, new_spk, new_ref = self.gpu_engine.run_step_persistent( external_inputs=ext_inputs, readback=True ) # Authoritative GPU refractory counters are used directly (P4/P5: # never recompute GPU state locally). else: new_pot, new_spk, new_ref = cpu_lif_step( self.graph.row_offsets, self.graph.col_indices, self.graph.weights, prev_spikes, ext_inputs, pot_in, refractory_in=ref_in, decay=0.85, threshold=1.0, v_reset=0.0, v_rest=0.0, t_ref=2 ) # Update running activations (moving average spike rate) self.state.activations = self.state.activations * 0.9 + new_spk * 0.1 self.state.membrane_potentials = new_pot self.state.spikes = new_spk self.state.refractory_steps = new_ref active_spikes = int(np.sum(new_spk > 0.5)) self.state.total_spikes += active_spikes self.state.step_count += 1 self.state.current_reward = reward # Predictive coding & homeostasis activity_mean = float(np.mean(self.state.activations)) prediction = self.state.predicted_reward error = abs(reward - prediction) self.state.prediction_error = float(error) self.state.predicted_reward += 0.1 * (reward - prediction) novelty = float(np.std(self.state.activations)) self.state.drives.step(activity_level=activity_mean, novelty=novelty) # Prediction-influenced attention/curiosity (opt-in; off in compat mode) if self.prediction_influence and len(self.state.attention) == N: err = float(self.state.prediction_error) gate = np.float32(np.clip(1.0 + self.prediction_gain * err, 0.5, 2.0)) for idx in (self.sensory_visual_indices, self.sensory_olfactory_indices, self.sensory_memory_indices): if len(idx): self.state.attention[idx] = np.clip( self.state.attention[idx] * gate, 0.0, 1.0) self.state.attention[idx] /= max( 1e-6, float(np.mean(self.state.attention[idx]))) self.state.drives.curiosity = float(np.clip( self.state.drives.curiosity + 0.1 * self.prediction_gain * err, 0.0, 1.5)) # Synaptic Plasticity synapses_updated = 0 neuromod_signal = 0.0 if self.plasticity_mode == "v2_eligibility" and self.eligibility is not None: # v2: traces update every step; weights move when the versioned # neuromodulatory signal is nonzero (structural changes resize traces). self.eligibility.sync_size(len(self.graph.weights)) if self.enable_plasticity: res = self.eligibility_engine.step( self.graph, self.eligibility, pre_spikes=prev_spikes, post_spikes=new_spk, reward=reward, novelty=novelty, prediction_error=float(self.state.prediction_error)) neuromod_signal = res["signal"] synapses_updated = res["synapses_updated"] if abs(res["signal"]) > 1e-6 else 0 if synapses_updated: # Provenance: weights changed -> refresh graph hash. self.graph.graph_hash = self.graph.compute_graph_hash() elif self.enable_plasticity and abs(reward) > 1e-4: if self.gpu_engine is not None and self.use_gpu: # P4: the persistent step's ping-pong copy already overwrote the # prev_spikes buffer with S(t); restore true S(t-1) so the # three-factor rule sees identical pre/post on CPU and GPU, # then re-upload S(t) to leave next-step state intact. # v4.1: weights stay GPU-resident (no per-step readback); # the CPU mirror syncs lazily via sync_gpu_weights(). self.gpu_engine.upload_buffer_data("prev_spikes", prev_spikes) self.gpu_engine.run_plasticity_persistent( learning_rate=0.05, reward=reward, readback=False ) self.gpu_engine.upload_buffer_data("prev_spikes", new_spk) self._gpu_weights_dirty = True synapses_updated = len(self.graph.weights) else: synapses_updated = self.plasticity.apply_hebbian_update( self.graph, prev_activations=prev_spikes, post_activations=new_spk, reward=reward ) self.plasticity_updates += synapses_updated # Motor decoding motor_speak = float(np.mean(self.state.activations[self.motor_speak_indices])) if len(self.motor_speak_indices) else 0.0 motor_image = float(np.mean(self.state.activations[self.motor_image_indices])) if len(self.motor_image_indices) else 0.0 motor_act = float(np.mean(self.state.activations[self.motor_act_indices])) if len(self.motor_act_indices) else 0.0 motor_remember = float(np.mean(self.state.activations[self.motor_remember_indices])) if len(self.motor_remember_indices) else 0.0 self.state.tool_associations = { "speak": round(motor_speak, 3), "generate_image": round(motor_image, 3), "act_in_environment": round(motor_act, 3), "remember": round(motor_remember, 3) } # Determine dominant motor action candidate_actions = sorted(self.state.tool_associations.items(), key=lambda x: x[1], reverse=True) selected_action = candidate_actions[0][0] if candidate_actions else "explore" return { "step": self.state.step_count, "spikes": active_spikes, "mean_potential": float(np.mean(new_pot)), "mean_activation": activity_mean, "prediction_error": round(self.state.prediction_error, 4), "drives": { "energy": round(self.state.drives.energy, 3), "curiosity": round(self.state.drives.curiosity, 3), "social": round(self.state.drives.social, 3), "integrity": round(self.state.drives.integrity, 3) }, "selected_action": selected_action, "synapses_updated": synapses_updated, "weights_synced": not self._gpu_weights_dirty, "plasticity_mode": self.plasticity_mode, "neuromod_signal": round(neuromod_signal, 6), "backend": "vulkan_gpu" if (self.use_gpu and self.gpu_engine) else "cpu_reference" } def save_snapshot(self, filepath: str): """Saves full deterministic state snapshot for perfect resumption.""" self.sync_gpu_weights() os.makedirs(os.path.dirname(filepath), exist_ok=True) np.savez_compressed( filepath, step_count=self.state.step_count, total_spikes=self.state.total_spikes, membrane_potentials=self.state.membrane_potentials, spikes=self.state.spikes, refractory_steps=self.state.refractory_steps, activations=self.state.activations, attention=self.state.attention, goal_embedding=self.state.goal_embedding, active_memory_refs=np.array(self.state.active_memory_refs, dtype=str), prediction_error=self.state.prediction_error, predicted_reward=self.state.predicted_reward, current_reward=self.state.current_reward, weights=self.graph.weights, row_offsets=self.graph.row_offsets, col_indices=self.graph.col_indices, neuron_ids=self.graph.neuron_ids, energy=self.state.drives.energy, curiosity=self.state.drives.curiosity, social=self.state.drives.social, integrity=self.state.drives.integrity, graph_hash=self.graph.graph_hash, graph_mode=self.graph.mode.value, tool_associations=json.dumps(self.state.tool_associations), active_goal=self.state.active_goal, plasticity_mode=self.plasticity_mode, eligibility_traces=(self.eligibility.traces if self.eligibility is not None else np.zeros(0, dtype=np.float32)), eligibility_updates=(self.eligibility.updates if self.eligibility is not None else 0), ) def load_snapshot(self, filepath: str): """Restores complete brain state and connectome weights from snapshot.""" if not os.path.exists(filepath): raise FileNotFoundError(f"Snapshot file not found: {filepath}") data = np.load(filepath, allow_pickle=True) self.state.step_count = int(data["step_count"]) self.state.total_spikes = int(data["total_spikes"]) self.state.membrane_potentials = np.copy(data["membrane_potentials"]) self.state.spikes = np.copy(data["spikes"]) self.state.refractory_steps = np.copy(data["refractory_steps"]) self.state.activations = np.copy(data["activations"]) if "attention" in data and len(data["attention"]) == len(self.state.activations): self.state.attention = np.copy(data["attention"]) if "goal_embedding" in data: self.state.goal_embedding = np.copy(data["goal_embedding"]) if "active_memory_refs" in data: self.state.active_memory_refs = [str(x) for x in list(data["active_memory_refs"])] self.state.prediction_error = float(data["prediction_error"]) self.state.predicted_reward = float(data["predicted_reward"]) self.state.current_reward = float(data["current_reward"]) self.state.drives.energy = float(data["energy"]) self.state.drives.curiosity = float(data["curiosity"]) self.state.drives.social = float(data["social"]) self.state.drives.integrity = float(data["integrity"]) if "tool_associations" in data: self.state.tool_associations = json.loads(str(data["tool_associations"])) if "active_goal" in data: self.state.active_goal = str(data["active_goal"]) if "plasticity_mode" in data: self.plasticity_mode = str(data["plasticity_mode"]) if self.plasticity_mode == "v2_eligibility": if self.eligibility is None: self.eligibility = EligibilityState(len(self.graph.weights), seed=self.seed) self.eligibility_engine = EligibilityEngine() if "eligibility_traces" in data: self.eligibility.traces = np.copy(data["eligibility_traces"]).astype(np.float32) self.eligibility.updates = int(data["eligibility_updates"]) else: self.eligibility = None self.eligibility_engine = None # Restore graph weights and topology self.graph.weights = np.copy(data["weights"]) self.graph.row_offsets = np.copy(data["row_offsets"]) self.graph.col_indices = np.copy(data["col_indices"]) self.graph.graph_hash = self.graph.compute_graph_hash() # If GPU backend active, reload persistent circuit if self.gpu_engine is not None and self.use_gpu: self.gpu_engine.load_circuit( self.graph.row_offsets, self.graph.col_indices, self.graph.weights, self.state.membrane_potentials, self.state.spikes, self.state.refractory_steps ) def restore_snapshot(self, filepath: str): return self.load_snapshot(filepath) def cleanup(self): if self.gpu_engine: self.gpu_engine.cleanup() self.gpu_engine = None