ModularBrainAgent / CLEANED_modular_brain_agent (1).py

Upload CLEANED_modular_brain_agent (1).py

55bc86e verified 10 months ago

304 kB

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.data import DataLoader, TensorDataset
	from torchvision import datasets, transforms
	from torch.utils.tensorboard import SummaryWriter
	import numpy as np
	import random
	from collections import deque
	import os # For creating log directory

	# --- 0. Global Setup and Vocabulary ---

	# Simulated language data
	sample_sentences = [
	"I love AI",
	"Deep learning is fun",
	"Spiking neurons are cool",
	"Brain inspired models rock",
	"Replay buffer helps learning"
	]


	def simple_tokenize(text):
	"""Basic tokenizer for sample sentences."""
	return text.lower().split()[:10]

	# Build vocab manually
	vocab = {"<unk>": 0}
	word_counter = 1
	for sentence in sample_sentences:
	for word in simple_tokenize(sentence):
	if word not in vocab:
	vocab[word] = word_counter
	word_counter += 1

	# Task IDs (using an Enum-like class for clarity)
	class Task:
	REGRESSION = 3
	BINARY = 4
	VISION = 5

	# --- 1. Base Classes and Spiking Neuron Implementations ---

	class Module(nn.Module):
	"""Base class for brain-inspired modules."""

	def __init__(self):
	super().__init__()

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	raise NotImplementedError

	class SurrogateLIF(torch.autograd.Function):
	"""
	Surrogate gradient function for LIF neurons.
	Allows backpropagation through the non-differentiable spiking
	function.
	"""
	@staticmethod
	def forward(
	ctx, input: torch.Tensor) -> torch.Tensor:
	ctx.save_for_backward(input)
	return (input > 0).float()

	@staticmethod
	def backward(
	ctx, grad_output: torch.Tensor) -> torch.Tensor:
	input, = ctx.saved_tensors
	grad_input = grad_output.clone()
	# Approximate derivative: a
	# constant value where
	# input is
	close to 0
	grad_input[input.abs(
	) < 1] = 1.0
	return grad_input

	class SpikingNeuron(
	Module):
	"""
	Leaky Integrate-and-Fire (LIF) neuron model.
	Resets membrane potential to zero after spiking.
	"""

	def __init__(
	self, threshold: float = 1.0, decay: float = 0.95):
	super().__init__()
	self.threshold = threshold
	self.decay = decay
	# Register buffer
	# for membrane
	# potential,
	# persists across
	forward calls
	# type:
	# torch.Tensor
	self.register_buffer(
	'mem', None)

	def forward(
	self, x: torch.Tensor) -> torch.Tensor:
	# Initialize or
	# re-initialize
	# membrane
	# potential if
	# shape
	changes
	if self.mem is None or self.mem.shape != x.shape:
	self.mem = torch.zeros_like(
	x)

	# Update
	# membrane
	# potential
	self.mem = self.decay * self.mem + x
	# Generate
	# spike
	spike = (
	self.mem >= self.threshold).float()

	# Reset
	# membrane
	# potential
	# for
	# spiking
	# neurons
	self.mem = torch.where(spike.bool(),
	torch.zeros_like(self.mem), self.mem)

	# Apply
	# surrogate
	# gradient
	# for
	# backpropagation
	return SurrogateLIF.apply(
	spike)

	class AdaptiveLIFNeuron(
	SpikingNeuron):
	"""
	LIF neuron with an adaptive threshold based on recent firing rate.
	"""

	def __init__(
	self, threshold: float = 1.0, decay: float = 0.95):
	super().__init__(threshold=threshold, decay=decay)
	# Register
	# buffer
	# for
	# adaptive
	# threshold
	# state
	self.register_buffer('threshold_state',
	torch.tensor(threshold))

	def forward(
	self, x: torch.Tensor) -> torch.Tensor:
	if self.mem is None or self.mem.shape != x.shape:
	self.mem = torch.zeros_like(
	x)

	self.mem = self.decay * self.mem + x
	spike = (
	self.mem >= self.threshold_state).float()
	self.mem = torch.where(spike.bool(),
	torch.zeros_like(self.mem), self.mem)

	# Adapt
	# threshold
	# based
	# on
	# scalar
	# mean
	# spike
	# rate
	# (no_grad
	# for
	threshold update)
	with torch.no_grad():
	spike_rate = spike.mean().item() # Scalar rate
	# Simple
	# moving
	# average
	# for
	# threshold
	# adaptation
	self.threshold_state = self.threshold_state * 0.99 +
	spike_rate * 0.01

	return SurrogateLIF.apply(
	spike)

	# --- 2. Encoder Modules (Task-Specific Input Processing) ---

	class SharedEncoder(
	nn.Module):
	"""Generic encoder for numerical data."""
	def __init__(
	self, input_size: int, hidden_size: int = 4):
	super().__init__()
	self.encoder = nn.Sequential(
	nn.Linear(
	input_size, hidden_size),
	nn.LayerNorm(
	hidden_size),
	nn.ReLU(),
	nn.Dropout(
	0.3)

	)

	def forward(
	self, x: torch.Tensor) -> torch.Tensor:
	assert x.dim() == 2, f"Expected 2D input for SharedEncoder,
	got {x.shape}"
	return self.encoder(x.float())


	class CNNVision(nn.Module):
	"""CNN encoder for image data (e.g., MNIST)."""
	def __init__(self, output_features: int = 4):
	super().__init__()
	self.conv = nn.Sequential(
	nn.Conv2d(1, 4, kernel_size=5),
	nn.ReLU(),
	nn.MaxPool2d(2),
	nn.Conv2d(4, 8, kernel_size=5),
	nn.ReLU(),
	nn.AdaptiveAvgPool2d((1, 1)), # Outputs [batch, 8, 1, 1]
	nn.Flatten(), # Outputs [batch, 8]
	nn.Linear(8, output_features) # Maps to desired output
	size
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.conv(x)


	class GRULanguage(nn.Module):
	"""GRU encoder for sequential language data."""
	def __init__(self, vocab_size: int, embedding_dim: int = 4,
	hidden_dim: int = 4):
	super().__init__()
	self.embed = nn.Embedding(vocab_size, embedding_dim)
	self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	if x.dim() == 1:
	x = x.unsqueeze(0) # Add batch dimension if missing
	emb = self.embed(x.long()) # Ensure input is long type for
	embedding
	out, _ = self.gru(emb)
	return out[:, -1, :] # Use last hidden state as sequence
	representation

	# --- 3. Brain-Inspired Modules with Specific Neuron Counts ---

	class SensoryProcessor(Module):

	"""Processes initial sensory input, often with low-level
	features."""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	self.dropout = nn.Dropout(0.3)
	self.neuron = SpikingNeuron()

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	z = self.dropout(self.norm(self.linear(x)))
	return self.neuron(z)


	class RelayLayer(Module):
	"""
	Routes information, potentially performing attention-like
	operations.
	Simulates a small sequence for MultiheadAttention.
	"""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	# MultiheadAttention requires embed_dim to be divisible by
	num_heads
	self.attn = nn.MultiheadAttention(embed_dim=output_dim,
	num_heads=2, batch_first=True)
	self.dropout = nn.Dropout(0.3)
	self.neuron = SpikingNeuron()

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	z = self.dropout(self.norm(self.linear(x)))

	# Create fake time dimension by repeating the input for
	attention
	seq_len = 4
	z_seq = z.unsqueeze(1).repeat(1, seq_len, 1)

	# Apply attention across the simulated time steps
	z_out, _ = self.attn(z_seq, z_seq, z_seq)

	# Pool over the time dimension to get a single representation
	routed = z_out.mean(dim=1)

	return self.neuron(routed)



	class InterneuronLogic(Module):
	"""Core processing unit, potentially for decision making or
	high-level abstraction."""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	self.dropout = nn.Dropout(0.3)
	self.neuron = AdaptiveLIFNeuron() # Using adaptive neuron here

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	z = self.dropout(self.norm(self.linear(x)))
	return self.neuron(z)


	class NeuroendocrineModulator(Module):
	"""
	Modulates signals, potentially for gain control or emotional
	states.
	Applies a sigmoid-based gain to its input.
	"""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	self.dropout = nn.Dropout(0.3)
	self.gain_control = nn.Linear(output_dim, output_dim) # Maps
	output to gain factor

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	z = self.dropout(self.norm(self.linear(x)))
	gain = torch.sigmoid(self.gain_control(z)) # Gain factor
	between 0 and 1
	return z * gain


	class AutonomicProcessor(Module):
	"""
	Manages internal states, often involves recurrent processing.
	Uses a GRU for sequential processing.
	"""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	self.recurrent = nn.GRU(output_dim, output_dim,
	batch_first=True)
	self.feedback_gain = 0.9


	def forward(self, x: torch.Tensor) -> torch.Tensor:
	z = self.norm(self.linear(x))
	# GRU expects (batch, seq_len, features) -> seq_len=1 for
	single step
	h, _ = self.recurrent(z.unsqueeze(1))
	# Squeeze the sequence dimension back out
	return self.feedback_gain * h.squeeze(1)


	class MirrorComparator(Module):
	"""
	Compares current state to a reference, potentially for self-other
	distinction or goal comparison.
	Outputs spikes and an optional similarity score.
	"""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	self.comparison_layer = nn.CosineSimilarity(dim=1) # Cosine
	similarity for comparison
	self.neuron = SpikingNeuron()

	def compare(self, a: torch.Tensor, b: torch.Tensor) ->
	torch.Tensor:
	"""Compare pre-spike representations for similarity."""
	a_proj = self.linear(a)
	b_proj = self.linear(b)
	return self.comparison_layer(a_proj, b_proj)

	def forward(self, x: torch.Tensor, reference: torch.Tensor = None)
	-> (torch.Tensor, torch.Tensor):
	z = self.norm(self.linear(x))
	spike = self.neuron(z)
	similarity = None
	if reference is not None:
	# Note: The 'reference' input would need to be processed
	similarly before comparison
	# For simplicity, assuming 'reference' is already in the
	correct feature space for comparison_layer
	similarity = self.compare(z, reference) # Pass processed
	'z' for comparison
	return spike, similarity


	class PlaceGridMemory(Module):
	"""

	Spatial memory system using population codes and LSTM for
	sequential memory.
	"""
	def __init__(self, input_dim: int, output_dim: int):
	super().__init__()
	self.linear = nn.Linear(input_dim, output_dim)
	self.norm = nn.LayerNorm(output_dim)
	self.positional_encoder = self.population_code # Reference to
	internal method
	self.memory_cell = nn.LSTM(output_dim, output_dim,
	batch_first=True)

	def population_code(self, x: torch.Tensor, pop_size: int) ->
	torch.Tensor:
	"""Expands a scalar value into a population-coded (one-hot
	like) vector."""
	# Calculate a scalar representation (e.g., mean)
	x_scalar = x.mean(dim=1)
	x_normalized = torch.sigmoid(x_scalar) # Normalize to [0, 1]
	# Map normalized value to an index within the population size
	idx = (x_normalized * (pop_size - 1)).long().clamp(0, pop_size
	- 1)

	encoded = torch.zeros(x.size(0), pop_size, device=x.device)
	# Set the corresponding index to 1.0
	encoded.scatter_(1, idx.unsqueeze(1), 1.0)
	return encoded

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	encoded = self.norm(self.linear(x))
	# Apply population coding based on configured output_dim
	pos_encoded = self.positional_encoder(encoded,
	pop_size=self.linear.out_features)
	# LSTM expects (batch, seq_len, features) -> seq_len=1 for
	single step
	memory_out, _ = self.memory_cell(pos_encoded.unsqueeze(1))
	# Squeeze the sequence dimension back out
	return memory_out.squeeze(1)


	# --- 4. Task Heads (Task-Specific Output Layers) ---

	class TaskHead(nn.Module):
	"""Generic task head for different output types."""
	def __init__(self, input_dim: int, output_dim: int, task_type:
	str):
	super().__init__()
	self.task_type = task_type

	self.head = nn.Sequential(
	nn.Linear(input_dim, output_dim),
	nn.Dropout(0.3)
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	# Regression and Binary will be squeezed to [batch]
	# Vision will remain [batch, num_classes]
	return self.head(x).squeeze(-1) if self.task_type in
	['binary', 'regression'] else self.head(x)


	# --- 5. Replay Buffer for Continual Learning ---

	class TaskReplayBuffer:
	"""
	A replay buffer that stores experiences tagged with their task ID.
	Samples are drawn proportionally or randomly across tasks.
	"""
	def __init__(self, buffer_size: int = 1000, device: str = "cpu"):
	self.buffer_size = buffer_size
	self.task_buffers = {task_id: deque(maxlen=buffer_size) for
	task_id in [Task.REGRESSION, Task.BINARY, Task.VISION]}
	self.device = device

	def add(self, task_id: int, states: torch.Tensor, labels:
	torch.Tensor):
	"""Adds experiences (state-label pairs) to the buffer for a
	specific task."""
	if task_id not in self.task_buffers:
	# Should not happen if self.task_buffers is
	pre-initialized with all Task IDs
	self.task_buffers[task_id] =
	deque(maxlen=self.buffer_size)

	for state, label in zip(states, labels):
	# Detach from graph and move to CPU to prevent memory
	leaks
	self.task_buffers[task_id].append((state.detach().cpu(),
	label.detach().cpu()))

	def sample(self, batch_size: int = 32) -> (torch.Tensor,
	torch.Tensor, list):
	"""
	Samples a batch of experiences from the buffer, ensuring mixed
	tasks.
	Returns: (states, labels, list of corresponding task_ids)
	"""

	active_tasks = [tid for tid, buffer in
	self.task_buffers.items() if len(buffer) > 0]
	if not active_tasks:
	return None, None, None

	# Sample at least two distinct tasks if possible
	num_tasks_to_sample = min(2, len(active_tasks))
	sampled_tasks_ids = random.sample(active_tasks,
	num_tasks_to_sample)

	batch_states, batch_labels, batch_task_ids = [], [], []

	# Distribute batch_size across sampled tasks
	per_task_samples_base = batch_size // num_tasks_to_sample
	remainder = batch_size % num_tasks_to_sample

	for i, task_id in enumerate(sampled_tasks_ids):
	task_list = list(self.task_buffers[task_id]) # Convert
	deque to list for random.sample/choices
	if not task_list:
	continue

	# Add remainder to first few tasks
	k = per_task_samples_base + (1 if i < remainder else 0)
	k = min(k, len(task_list)) # Ensure we don't sample more
	than available

	if k == 0: continue

	try:
	samples = random.sample(task_list, k)
	except ValueError:
	# If k is larger than task_list length (should be
	caught by min(k, len(task_list))
	# but good fallback for robustness)
	samples = random.choices(task_list, k=k)

	for state, label in samples:
	batch_states.append(state)
	batch_labels.append(label)
	batch_task_ids.append(task_id)

	if not batch_states:
	return None, None, None

	# Stack tensors and return task IDs for individual processing
	return (
	torch.stack(batch_states),

	torch.stack(batch_labels),
	batch_task_ids # Return as list to preserve individual
	task identities
	)


	# --- 6. Full Agent Model ---

	class ModularBrainAgent(nn.Module):
	"""
	A comprehensive modular neural agent inspired by brain
	architecture,
	integrating spiking neurons, attention, and recurrent memory for
	multi-task learning.
	"""
	def __init__(self, neuron_counts: dict = None):
	super().__init__()

	# --- Parameterization of Neuron Counts ---
	if neuron_counts is None:
	# Default neuron counts for each brain region
	neuron_counts = {
	'sensory': 4,
	'relay': 12,
	'interneurons': 2,
	'neuroendocrine': 8,
	'autonomic': 10,
	'mirror': 14,
	'place_grid': 16
	}
	self.neuron_counts = neuron_counts

	# --- Encoders (Input Modalities) ---
	self.encoders = nn.ModuleDict({
	'regression': SharedEncoder(2,
	self.neuron_counts['sensory']),
	'language': GRULanguage(len(vocab),
	embedding_dim=self.neuron_counts['sensory'],
	hidden_dim=self.neuron_counts['sensory']),
	'vision':
	CNNVision(output_features=self.neuron_counts['sensory'])
	})

	# --- Brain Modules with Realistic Neuron Counts ---
	self.sensory = SensoryProcessor(
	input_dim=self.neuron_counts['sensory'],
	output_dim=self.neuron_counts['sensory']
	)

	self.relay = RelayLayer(
	input_dim=self.neuron_counts['sensory'], # Input from
	SensoryProcessor
	output_dim=self.neuron_counts['relay']
	)
	self.interneurons = InterneuronLogic(
	input_dim=self.neuron_counts['relay'], # Input from
	RelayLayer
	output_dim=self.neuron_counts['interneurons']
	)
	self.neuroendocrine = NeuroendocrineModulator(
	input_dim=self.neuron_counts['interneurons'],
	output_dim=self.neuron_counts['neuroendocrine']
	)
	self.autonomic = AutonomicProcessor(
	input_dim=self.neuron_counts['neuroendocrine'],
	output_dim=self.neuron_counts['autonomic']
	)
	self.mirror = MirrorComparator(
	input_dim=self.neuron_counts['autonomic'],
	output_dim=self.neuron_counts['mirror']
	)
	self.place_grid = PlaceGridMemory(
	input_dim=self.neuron_counts['mirror'],
	output_dim=self.neuron_counts['place_grid']
	)

	# --- Interconnection Weights (between brain modules) ---
	# Ensure input/output dimensions match the neuron_counts
	self.connect_sensory_to_relay =
	nn.Linear(self.neuron_counts['sensory'], self.neuron_counts['relay'])
	self.connect_relay_to_inter =
	nn.Linear(self.neuron_counts['relay'],
	self.neuron_counts['interneurons'])
	self.connect_inter_to_modulators =
	nn.Linear(self.neuron_counts['interneurons'],
	self.neuron_counts['neuroendocrine'])
	self.connect_modulators_to_auto =
	nn.Linear(self.neuron_counts['neuroendocrine'],
	self.neuron_counts['autonomic'])
	self.connect_auto_to_mirror =
	nn.Linear(self.neuron_counts['autonomic'],
	self.neuron_counts['mirror'])
	self.connect_mirror_to_place =
	nn.Linear(self.neuron_counts['mirror'],
	self.neuron_counts['place_grid'])
	# Recurrent loop from place_grid back to relay
	self.connect_place_to_relay =

	nn.Linear(self.neuron_counts['place_grid'],
	self.neuron_counts['relay'])

	# --- Optional Feedback Connections ---
	self.feedback_relay_to_sensory =
	nn.Linear(self.neuron_counts['relay'], self.neuron_counts['sensory'])
	self.feedback_inter_to_relay =
	nn.Linear(self.neuron_counts['interneurons'],
	self.neuron_counts['relay'])

	# --- Task Heads (Outputs for specific tasks) ---
	self.task_heads = nn.ModuleDict({
	'binary': TaskHead(self.neuron_counts['place_grid'], 1,
	'binary'),
	'vision': TaskHead(self.neuron_counts['place_grid'], 10,
	'vision'), # MNIST has 10 classes
	'regression': TaskHead(self.neuron_counts['place_grid'],
	1, 'regression')
	})

	def route_modules(self, x: torch.Tensor) -> tuple[torch.Tensor,
	...]:
	"""
	Defines the forward pass through the interconnected brain
	modules.
	Returns intermediate activations for potential monitoring or
	later use.
	"""
	# Forward pass through brain circuits
	h_sensory = self.sensory(x)
	h_relay = self.connect_sensory_to_relay(h_sensory)
	h_relay = self.relay(h_relay)

	h_inter = self.connect_relay_to_inter(h_relay)
	h_inter = self.interneurons(h_inter)

	h_modulate = self.connect_inter_to_modulators(h_inter)
	h_modulate = self.neuroendocrine(h_modulate)

	h_auto = self.connect_modulators_to_auto(h_modulate)
	h_auto = self.autonomic(h_auto)

	h_mirror_result, _ =
	self.mirror(self.connect_auto_to_mirror(h_auto)) # Mirror returns
	tuple
	# if isinstance(mirror_result, tuple): # No longer needed as
	forward always returns tuple
	# h_mirror, _ = mirror_result

	# else:
	# h_mirror = mirror_result
	h_mirror = h_mirror_result # Renamed for clarity after tuple
	unpack

	h_place = self.connect_mirror_to_place(h_mirror)
	h_place = self.place_grid(h_place)

	# Recurrent loop from hippocampus (place_grid) back to
	thalamus (relay)
	# Note: Added ReLU here to prevent negative values from
	feeding back into spiking neurons potentially
	h_relay = h_relay +
	torch.relu(self.connect_place_to_relay(h_place))

	# Optional feedback connections
	h_sensory = h_sensory +
	torch.relu(self.feedback_relay_to_sensory(h_relay))
	h_relay = h_relay +
	torch.relu(self.feedback_inter_to_relay(h_inter))

	return h_relay, h_place, h_mirror, h_auto, h_modulate


	def encode(self, x: torch.Tensor, task_id: int) -> torch.Tensor:
	"""Selects and applies the appropriate encoder based on task
	ID and input shape."""
	# Check task ID and input shape for specific encoders
	if task_id == Task.REGRESSION and x.dim() == 2 and x.size(1)
	== 2:
	return self.encoders['regression'](x.float())
	elif task_id == Task.BINARY and x.dim() == 2 and x.size(1) ==
	10:
	return self.encoders['language'](x.long())
	elif task_id == Task.VISION:
	if x.dim() == 4: # Already 4D (batch, 1, H, W)
	return self.encoders['vision'](x.float())
	elif x.dim() == 2 and x.size(1) == 784: # Flattened 2D
	(batch, 784)
	return self.encoders['vision'](x.view(x.size(0), 1,
	28, 28))
	else:
	raise ValueError(f"Unexpected vision input shape:
	{x.shape}")
	else:
	# Fallback for unexpected task_id/shape combination, or if
	x is already "encoded"
	# In a real system, you might want a specific error or

	default encoder here.
	# For now, assumes x is already in the 'sensory' input
	dimension if it falls through.
	if x.size(-1) != self.neuron_counts['sensory']:
	raise ValueError(f"Input {x.shape} does not match
	sensory input dim {self.neuron_counts['sensory']} and no specific
	encoder found for task {task_id}.")
	return x.float()


	def forward(self, x: torch.Tensor, task_id: int) -> torch.Tensor:
	"""
	Main forward pass of the Modular Brain Agent.
	Encodes input, routes through brain modules, and selects task
	head.
	"""
	if x.dim() == 1:
	x = x.unsqueeze(0) # Ensure batch dimension

	encoded = self.encode(x, task_id)
	# Route through brain modules, focusing on the outputs used
	for task heads
	_, h_place, _, _, _ = self.route_modules(encoded) # Only
	h_place is used for task heads

	# Map task IDs to the appropriate task head name
	head_name = {
	Task.REGRESSION: 'regression',
	Task.BINARY: 'binary',
	Task.VISION: 'vision'
	}.get(task_id, 'regression') # Default to regression head if
	task_id is unexpected

	out = self.task_heads[head_name](h_place) # All tasks use the
	PlaceGridMemory output
	return out


	# --- 7. Data Loaders (Synthetic and Real) ---

	def get_mnist_loader() -> (DataLoader, int):
	"""Returns MNIST DataLoader and its Task ID."""
	transform = transforms.Compose([
	transforms.ToTensor(),
	transforms.Normalize((0.5,), (0.5,))
	])
	dataset = datasets.MNIST(root='./data', train=True, download=True,
	transform=transform)

	return DataLoader(dataset, batch_size=4, shuffle=True),
	Task.VISION


	def get_imdb_loader() -> (DataLoader, int):
	"""Returns a synthetic IMDB-like DataLoader and its Task ID."""
	texts = []
	labels = []

	for i in range(1000): # Create 1000 synthetic samples
	sentence = random.choice(sample_sentences)
	tokens = [vocab.get(word, vocab["<unk>"]) for word in
	simple_tokenize(sentence)]

	# Pad or truncate tokens to a fixed length
	while len(tokens) < 10:
	tokens.append(vocab["<unk>"])
	if len(tokens) > 10:
	tokens = tokens[:10]

	texts.append(tokens)
	labels.append(i % 2) # Binary sentiment: 0 or 1

	X = torch.tensor(texts).long()
	Y = torch.tensor(labels).float()
	dataset = TensorDataset(X, Y)
	return DataLoader(dataset, batch_size=4, shuffle=True),
	Task.BINARY


	def get_regression_loader() -> (DataLoader, int):
	"""Returns a synthetic regression DataLoader and its Task ID."""
	X = torch.randn(1000, 2) # 1000 samples, 2 features
	Y = ((X[:, 0] + X[:, 1]) / 2).float() # Simple linear relationship
	dataset = TensorDataset(X, Y)
	return DataLoader(dataset, batch_size=4, shuffle=True),
	Task.REGRESSION


	# --- 8. Helper Function: Shape Inspector ---

	def inspect_shapes(agent: ModularBrainAgent):
	"""
	Prints the shapes of tensors at various points in the agent's
	forward pass
	to help verify architectural correctness.
	"""
	print("\n=== SHAPE INSPECTOR ===")

	# Use CPU for inspection to avoid CUDA memory issues during
	debugging
	agent.cpu()
	agent.eval() # Set to eval mode for consistent behavior (e.g.,
	dropout off)

	try:
	# Test regression path
	sample_input_reg = torch.randn(4, 2)
	encoded_reg = agent.encode(sample_input_reg, Task.REGRESSION)
	print(f"Encoded (Regression): {encoded_reg.shape}")
	h_relay_reg, h_place_reg, h_mirror_reg, h_auto_reg,
	h_modulate_reg = agent.route_modules(encoded_reg)
	out_reg = agent.task_heads['regression'](h_place_reg)
	print(f"Regression Path:
	Sensory({agent.sensory(encoded_reg).shape}) ->
	Relay({h_relay_reg.shape}) -> Place({h_place_reg.shape}) ->
	Output({out_reg.shape})")

	# Test language path
	# Using fixed input_dim for language to map to sensory input
	dim
	sample_input_lang = torch.randint(0, len(vocab), (4, 10))
	encoded_lang = agent.encode(sample_input_lang, Task.BINARY) #
	Use BINARY task for language
	print(f"Encoded (Language): {encoded_lang.shape}")
	h_relay_lang, h_place_lang, h_mirror_lang, h_auto_lang,
	h_modulate_lang = agent.route_modules(encoded_lang)
	out_lang = agent.task_heads['binary'](h_place_lang)
	print(f"Language Path:
	Sensory({agent.sensory(encoded_lang).shape}) ->
	Relay({h_relay_lang.shape}) -> Place({h_place_lang.shape}) ->
	Output({out_lang.shape})")

	# Test vision path
	sample_input_vis = torch.randn(4, 1, 28, 28)
	encoded_vis = agent.encode(sample_input_vis, Task.VISION)
	print(f"Encoded (Vision): {encoded_vis.shape}")
	h_relay_vis, h_place_vis, h_mirror_vis, h_auto_vis,
	h_modulate_vis = agent.route_modules(encoded_vis)
	out_vis = agent.task_heads['vision'](h_place_vis)
	print(f"Vision Path:
	Sensory({agent.sensory(encoded_vis).shape}) ->
	Relay({h_relay_vis.shape}) -> Place({h_place_vis.shape}) ->
	Output({out_vis.shape})")

	except Exception as e:
	print(f"Shape inspection failed: {e}")

	# Re-raise to ensure the error is not ignored in main
	execution
	raise e
	finally:
	agent.train() # Set back to train mode
	print("=========================\n")


	# --- 9. Training Function ---

	def train(agent: ModularBrainAgent, episodes: int = 14400,
	buffer_size: int = 1000, replay_freq: int = 5):
	"""
	Trains the ModularBrainAgent using a curriculum learning strategy
	and experience replay.
	"""
	device = torch.device("cuda" if torch.cuda.is_available() else
	"cpu")
	agent.to(device)
	print(f"Using device: {device}")

	print("🧠 Running shape inspector...")
	inspect_shapes(agent)

	replay_buffer = TaskReplayBuffer(buffer_size=buffer_size,
	device=device)

	optimizer = torch.optim.AdamW(agent.parameters(), lr=0.001,
	weight_decay=1e-5)
	scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
	patience=100, factor=0.5, verbose=True)

	# Initialize weights
	def init_weights(m):
	if isinstance(m, nn.Linear):
	torch.nn.init.xavier_uniform_(m.weight)
	if m.bias is not None:
	torch.nn.init.zeros_(m.bias)
	elif isinstance(m, nn.Conv2d):
	torch.nn.init.kaiming_uniform_(m.weight,
	nonlinearity='relu')
	if m.bias is not None:
	torch.nn.init.zeros_(m.bias)
	agent.apply(init_weights)

	# Setup TensorBoard writer
	log_dir = "runs/modular_brain_agent"
	os.makedirs(log_dir, exist_ok=True)

	writer = SummaryWriter(log_dir=log_dir)
	print(f"TensorBoard logs are being saved to: {log_dir}")


	best_loss = float('inf')
	no_improvement = 0
	save_path = "best_model.pth"
	task_history = {i: [] for i in [Task.REGRESSION, Task.BINARY,
	Task.VISION]}
	curriculum_stages = {
	0: [Task.REGRESSION],
	300: [Task.REGRESSION, Task.BINARY], # Start binary after 300
	episodes
	600: [Task.REGRESSION, Task.BINARY, Task.VISION] # Start
	vision after 600 episodes
	}

	raw_loaders = {
	Task.REGRESSION: get_regression_loader(),
	Task.BINARY: get_imdb_loader(),
	Task.VISION: get_mnist_loader()
	}

	loaders = {k: v[0] for k, v in raw_loaders.items()}
	# task_ids are fixed as keys for loaders in this setup

	# Initialize persistent iterators for each task's DataLoader
	loaders_iter = {
	Task.REGRESSION: iter(loaders[Task.REGRESSION]),
	Task.BINARY: iter(loaders[Task.BINARY]),
	Task.VISION: iter(loaders[Task.VISION])
	}

	# Helper for loss computation
	def compute_loss(out: torch.Tensor, Y: torch.Tensor, task_id: int)
	-> torch.Tensor:
	"""Computes loss based on task type, handling shape
	alignments."""
	if task_id == Task.VISION:
	Y = Y.long() # Target labels for CrossEntropy must be long
	if out.dim() != 2 or Y.dim() != 1:
	raise ValueError(f"Vision task expects out [batch,
	num_classes], Y [batch]. Got {out.shape} vs {Y.shape}")
	return F.cross_entropy(out, Y)

	elif task_id == Task.BINARY:
	# Ensure output and target have matching shapes for
	BCEWithLogitsLoss

	if out.shape != Y.shape:
	if out.numel() == Y.numel():
	out = out.view_as(Y) # Reshape output to match
	target if num elements are same
	else:
	raise ValueError(f"Binary task shape mismatch: out
	{out.shape} vs Y {Y.shape}")
	return F.binary_cross_entropy_with_logits(out, Y.float())

	else: # REGRESSION
	if out.shape != Y.shape:
	if out.numel() == Y.numel():
	out = out.view_as(Y) # Reshape output to match
	target if num elements are same
	else:
	raise ValueError(f"Regression task shape mismatch:
	out {out.shape} vs Y {Y.shape}")
	return F.smooth_l1_loss(out, Y.float())

	loss_weights = {Task.REGRESSION: 1.5, Task.BINARY: 1.2,
	Task.VISION: 1.2}
	global_step = 0 # For TensorBoard logging

	for ep in range(episodes):
	agent.train() # Ensure model is in training mode

	# Determine current curriculum stage
	current_stage_episodes = 0
	for stage_start_ep, tasks in
	sorted(curriculum_stages.items()):
	if ep >= stage_start_ep:
	current_stage_episodes = stage_start_ep
	current_tasks = curriculum_stages[current_stage_episodes]

	# Randomly select a task from the active curriculum stage
	task_id = np.random.choice(current_tasks)

	# Fetch data using persistent iterator
	try:
	X, Y = next(loaders_iter[task_id])
	except StopIteration:
	# Reset iterator when exhausted for that task
	loaders_iter[task_id] = iter(loaders[task_id])
	X, Y = next(loaders_iter[task_id])

	X, Y = X.to(device), Y.to(device)

	# --- Primary Training Step ---

	out = agent(X, task_id)
	loss = compute_loss(out, Y, task_id) *
	loss_weights.get(task_id, 1.0)

	optimizer.zero_grad()
	loss.backward()
	torch.nn.utils.clip_grad_norm_(agent.parameters(),
	max_norm=1.0) # Clip gradients
	optimizer.step()

	# Store current experience in replay buffer
	replay_buffer.add(task_id, X, Y)

	# --- Accuracy Metric Calculation ---
	acc = 0.0
	with torch.no_grad(): # Don't track gradients for accuracy
	calculation
	if task_id == Task.VISION:
	pred = out.argmax(dim=1)
	acc = (pred == Y).float().mean().item()
	elif task_id == Task.BINARY:
	pred = (torch.sigmoid(out) > 0.5).float()
	acc = (pred == Y).float().mean().item()
	else: # REGRESSION
	acc = ((out - Y).abs() < 0.2).float().mean().item() #
	within 0.2 error margin

	task_history[task_id].append((loss.item(), acc))

	# --- Replay Phase (Fixed Logic) ---
	if ep % replay_freq == 0:
	replay_X_mixed, replay_Y_mixed, replay_task_ids_list =
	replay_buffer.sample(batch_size=4)

	if replay_X_mixed is not None:
	replay_X_mixed = replay_X_mixed.to(device)
	replay_Y_mixed = replay_Y_mixed.to(device)

	unique_replay_tasks = list(set(replay_task_ids_list))
	total_replay_loss = 0.0
	num_replay_samples = 0

	# Iterate through each unique task in the sampled
	batch
	for current_replay_task_id in unique_replay_tasks:
	# Filter samples belonging to this specific task
	indices_for_this_task = [i for i, tid in
	enumerate(replay_task_ids_list) if tid == current_replay_task_id]


	if not indices_for_this_task:
	continue

	task_replay_X =
	replay_X_mixed[indices_for_this_task]
	task_replay_Y =
	replay_Y_mixed[indices_for_this_task]

	# Process these samples using the correct encoder
	and task head
	replay_out_for_task = agent(task_replay_X,
	current_replay_task_id)

	# Compute loss for this task's samples
	replay_loss_for_task =
	compute_loss(replay_out_for_task, task_replay_Y,
	current_replay_task_id)
	total_replay_loss += replay_loss_for_task *
	len(indices_for_this_task) # Weighted sum
	num_replay_samples += len(indices_for_this_task)

	if num_replay_samples > 0:
	total_replay_loss /= num_replay_samples # Average
	over all replayed samples
	optimizer.zero_grad()
	total_replay_loss.backward()
	torch.nn.utils.clip_grad_norm_(agent.parameters(),
	max_norm=1.0)
	optimizer.step()
	# Log replay loss
	writer.add_scalar('Loss/Replay_Loss',
	total_replay_loss.item(), global_step)


	# --- Logging and Early Stopping ---
	scheduler.step(loss.item()) # Step scheduler based on current
	task loss

	# TensorBoard Logging
	writer.add_scalar(f'Loss/Current_Task_{task_id}', loss.item(),
	global_step)
	writer.add_scalar(f'Accuracy/Current_Task_{task_id}', acc,
	global_step)
	writer.add_scalar('LearningRate',
	optimizer.param_groups[0]['lr'], global_step)

	global_step += 1


	# Check for best model and early stopping
	if loss.item() < best_loss:
	best_loss = loss.item()
	no_improvement = 0
	torch.save(agent.state_dict(), save_path)
	# print(f"New best model saved at episode {ep} with loss:
	{best_loss:.4f}")
	else:
	no_improvement += 1

	# Print summary periodically
	if ep % 200 == 0:
	print(f"\n--- Episode {ep} (Curriculum Stage:
	{current_stage_episodes}) ---")
	for t_id in sorted(task_history.keys()):
	if task_history[t_id]:
	# Get recent history, default to empty if not
	enough entries
	recent_losses = [x[0] for x in
	task_history[t_id][-200:]]
	recent_accs = [x[1] for x in
	task_history[t_id][-200:]]

	avg_loss = np.mean(recent_losses) if recent_losses
	else 0.0
	avg_acc = np.mean(recent_accs) if recent_accs else
	0.0

	# Map task ID to a readable name
	task_name = {Task.REGRESSION: "Regression",
	Task.BINARY: "Binary", Task.VISION: "Vision"}.get(t_id,
	f"Unknown_{t_id}")
	print(f"Task {task_name} \| Avg Loss:
	{avg_loss:.3f} \| Avg Acc: {avg_acc:.2f} ({len(recent_losses)}
	samples)")
	print(f"Overall Current Loss: {loss.item():.4f} \| Best
	Loss: {best_loss:.4f} \| No Improvement: {no_improvement}")
	print("--------------------\n")

	if no_improvement >= 1000: # Early stopping threshold
	print(f"Early stopping at episode {ep} due to no
	improvement for {no_improvement} steps.")
	break

	writer.close()
	print("✅ Training finished.")
	return best_loss



	# --- 10. Main Execution Block ---

	if __name__ == "__main__":
	print("🚀 Initializing Modular Brain Agent...")

	# Optional: Define custom neuron counts here
	# custom_neuron_config = {
	# 'sensory': 8,
	# 'relay': 24,
	# 'interneurons': 4,
	# 'neuroendocrine': 16,
	# 'autonomic': 20,
	# 'mirror': 28,
	# 'place_grid': 32
	# }
	# agent = ModularBrainAgent(neuron_counts=custom_neuron_config)

	agent = ModularBrainAgent() # Using default neuron counts for now

	# Print model summary
	total_params = sum(p.numel() for p in agent.parameters())
	trainable_params = sum(p.numel() for p in agent.parameters() if
	p.requires_grad)

	print(f"📊 Model Summary:")
	print(f" Total parameters: {total_params:,}")
	print(f" Trainable parameters: {trainable_params:,}")
	print(f" Model size: ~{total_params * 4 / (1024**2):.2f} MB
	(approx. float32)")

	# Test forward pass on each task to verify shapes
	print("\n🔍 Testing forward passes...")
	try:
	# Test regression task
	test_reg = torch.randn(2, 2)
	out_reg = agent(test_reg, Task.REGRESSION)
	print(f"✅ Regression test passed: {test_reg.shape} ->
	{out_reg.shape}")

	# Test binary classification task (language task)
	test_bin = torch.randint(0, len(vocab), (2, 10))
	out_bin = agent(test_bin, Task.BINARY)
	print(f"✅ Binary classification test passed: {test_bin.shape}
	-> {out_bin.shape}")

	# Test vision task

	test_vis = torch.randn(2, 1, 28, 28)
	out_vis = agent(test_vis, Task.VISION)
	print(f"✅ Vision test passed: {test_vis.shape} ->
	{out_vis.shape}")

	except Exception as e:
	print(f"❌ Forward pass test failed: {e}")
	# If forward pass fails, there's a fundamental issue, so exit
	import sys
	sys.exit(1)

	print("\n🎯 Starting training...")

	# Train the agent
	try:
	final_best_loss = train(
	agent=agent,
	episodes=14400, # Total training episodes
	buffer_size=1000, # Replay buffer size
	replay_freq=5 # Replay every N episodes
	)

	print(f"\n🎉 Training completed successfully!")
	print(f"📈 Best loss achieved during training:
	{final_best_loss:.4f}")
	print(f"💾 Best model saved to: best_model.pth")

	except KeyboardInterrupt:
	print("\n⏹ Training interrupted by user.")
	except Exception as e:
	print(f"\n❌ Training failed with error: {e}")
	raise e

	print("\n🧠 Modular Brain Agent execution completed!")