SynCo_modular_brain_agent_with_spikes_and_plasticity.py

# MIT License
#
# Copyright (c) 2025 ALMUSAWIY Halliru
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# === V3 Modular Brain Agent with Plasticity - Block 1 ===

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random
from torch.utils.data import DataLoader, Dataset
from collections import deque
from torchvision import datasets, transforms

# === Plastic Synapse Mechanisms ===
class PlasticLinear(nn.Module):
    def __init__(self, in_features, out_features, plasticity_type="hebbian", learning_rate=0.01):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.1)
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.plasticity_type = plasticity_type
        self.eta = learning_rate
        self.trace = torch.zeros(out_features, in_features)
        self.register_buffer('prev_y', torch.zeros(out_features))

    def forward(self, x):
        y = F.linear(x, self.weight, self.bias)
        if self.training:
            x_detached = x.detach()
            y_detached = y.detach()
            if self.plasticity_type == "hebbian":
                hebb = torch.einsum('bi,bj->ij', y_detached, x_detached) / x.size(0)
                self.trace = (1 - self.eta) * self.trace + self.eta * hebb
                with torch.no_grad():
                    self.weight += self.trace
            elif self.plasticity_type == "stdp":
                dy = y_detached - self.prev_y
                stdp = torch.einsum('bi,bj->ij', dy, x_detached) / x.size(0)
                self.trace = (1 - self.eta) * self.trace + self.eta * stdp
                with torch.no_grad():
                    self.weight += self.trace
                self.prev_y = y_detached.clone()
        return y

# === Spiking Surrogate Functions and Base Neurons ===
class SpikeFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return (input > 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * (abs(input) < 1).float()

spike_fn = SpikeFunction.apply

class LIFNeuron(nn.Module):
    def __init__(self, tau=2.0):
        super().__init__()
        self.tau = tau
        self.mem = 0

    def forward(self, x):
        decay = torch.exp(torch.tensor(-1.0 / self.tau))
        self.mem = self.mem * decay + x
        out = spike_fn(self.mem - 1.0)
        self.mem = self.mem * (1.0 - out.detach())
        return out

# === Adaptive LIF Neuron ===
class AdaptiveLIF(nn.Module):
    def __init__(self, size, tau=2.0, beta=0.2):
        super().__init__()
        self.size = size
        self.tau = tau
        self.beta = beta
        self.mem = torch.zeros(size)
        self.thresh = torch.ones(size)

    def forward(self, x):
        decay = torch.exp(torch.tensor(-1.0 / self.tau))
        self.mem = self.mem * decay + x
        out = spike_fn(self.mem - self.thresh)
        self.thresh = self.thresh + self.beta * out
        self.mem = self.mem * (1.0 - out.detach())
        return out

# === Relay Layer with Attention ===
class RelayLayer(nn.Module):
    def __init__(self, dim, heads=4):
        super().__init__()
        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=heads, batch_first=True)
        self.lif = LIFNeuron()

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        return self.lif(attn_out)

# === Working Memory ===
class WorkingMemory(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)

    def forward(self, x):
        out, _ = self.lstm(x)
        return out[:, -1]

# === Place Cell Grid ===
class PlaceGrid(nn.Module):
    def __init__(self, grid_size=10, embedding_dim=64):
        super().__init__()
        self.embedding = nn.Embedding(grid_size**2, embedding_dim)

    def forward(self, index):
        return self.embedding(index)

# === Mirror Comparator ===
class MirrorComparator(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.cos = nn.CosineSimilarity(dim=1)

    def forward(self, x1, x2):
        return self.cos(x1, x2).unsqueeze(1)

# === Neuroendocrine Module ===
class NeuroendocrineModulator(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)

    def forward(self, x):
        out, _ = self.lstm(x)
        return out[:, -1]

# === Autonomic Feedback Module ===
class AutonomicFeedback(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.feedback = nn.Linear(input_dim, input_dim)

    def forward(self, x):
        return torch.tanh(self.feedback(x))

# === Replay Buffer ===
class ReplayBuffer:
    def __init__(self, capacity=1000):
        self.buffer = deque(maxlen=capacity)

    def add(self, inputs, labels, task):
        self.buffer.append((inputs, labels, task))

    def sample(self, batch_size):
        indices = random.sample(range(len(self.buffer)), batch_size)
        batch = [self.buffer[i] for i in indices]
        inputs, labels, tasks = zip(*batch)
        return inputs, labels, tasks

# === Full Modular Brain Agent with Plasticity ===
class ModularBrainAgent(nn.Module):
    def __init__(self, input_dims, hidden_dim, output_dims):
        super().__init__()
        self.vision_encoder = nn.Linear(input_dims['vision'], hidden_dim)
        self.language_encoder = nn.Linear(input_dims['language'], hidden_dim)
        self.numeric_encoder = nn.Linear(input_dims['numeric'], hidden_dim)

        # Plastic synapses (Hebbian and STDP)
        self.connect_sensory_to_relay = PlasticLinear(hidden_dim * 3, hidden_dim, plasticity_type='hebbian')
        self.relay_layer = RelayLayer(hidden_dim)
        self.connect_relay_to_inter = PlasticLinear(hidden_dim, hidden_dim, plasticity_type='stdp')

        self.interneuron = AdaptiveLIF(hidden_dim)
        self.memory = WorkingMemory(hidden_dim, hidden_dim)
        self.place = PlaceGrid(grid_size=10, embedding_dim=hidden_dim)
        self.comparator = MirrorComparator(hidden_dim)
        self.emotion = NeuroendocrineModulator(hidden_dim, hidden_dim)
        self.feedback = AutonomicFeedback(hidden_dim)

        self.task_heads = nn.ModuleDict({
            task: nn.Linear(hidden_dim, out_dim)
            for task, out_dim in output_dims.items()
        })

        self.replay = ReplayBuffer()

    def forward(self, inputs, task, position_idx=None):
        v = self.vision_encoder(inputs['vision'])
        l = self.language_encoder(inputs['language'])
        n = self.numeric_encoder(inputs['numeric'])

        sensory_cat = torch.cat([v, l, n], dim=-1)
        z = self.connect_sensory_to_relay(sensory_cat)

        z = self.relay_layer(z.unsqueeze(1)).squeeze(1)
        z = self.connect_relay_to_inter(z)
        z = self.interneuron(z)

        m = self.memory(z.unsqueeze(1))
        p = self.place(position_idx if position_idx is not None else torch.tensor([0]))
        e = self.emotion(z.unsqueeze(1))
        f = self.feedback(z)

        combined = z + m + p + e + f
        out = self.task_heads[task](combined)
        return out

    def remember(self, inputs, labels, task):
        self.replay.add(inputs, labels, task)

# === Main Test Block ===
if __name__ == "__main__":
    input_dims = {'vision': 32, 'language': 16, 'numeric': 8}
    output_dims = {'classification': 5, 'regression': 1, 'binary': 1}
    agent = ModularBrainAgent(input_dims, hidden_dim=64, output_dims=output_dims)

    tasks = list(output_dims.keys())

    for step in range(250):
        task = random.choice(tasks)
        inputs = {
            'vision': torch.randn(1, 32),
            'language': torch.randn(1, 16),
            'numeric': torch.randn(1, 8)
        }
        labels = torch.randint(0, output_dims[task], (1,)) if task == 'classification' else torch.randn(1, output_dims[task])
        output = agent(inputs, task)
        loss = F.cross_entropy(output, labels) if task == 'classification' else F.mse_loss(output, labels)
        print(f"Step {step:02d} | Task: {task:13s} | Loss: {loss.item():.4f}")