corvidae.py

"""
corvidae.py -- the base architecture.

This is the foundation every other file (species_memory_bank.py, corvid_extensions.py,
corvid_aviary.py) imports and extends. It was referenced throughout those files but never
actually written until now; this fills that gap and, per the new batch of papers, upgrades
three of its core components so they reflect specific published findings rather than being
generic placeholders:

  CorvidaeMemory (generic DNC-style store)
      A simplified differentiable neural computer: content-addressed read/write over a
      persistent memory matrix, plus a `consolidate()` step that blends working memory
      into a slower long-term buffer (a crude analogue of sleep-dependent memory
      consolidation) -- called periodically by the trainer via `model.consolidate_memory()`.

  MetatoolPlanningBuffer
      Gruber et al. (2019, both Current Biology papers) showed New Caledonian crows keep
      mental representations of a functional sub-goal, a DISTRACTOR sub-goal, and a final
      goal simultaneously in mind across multi-stage, out-of-sight tool problems, and
      suppress attention to the distractor rather than just failing to notice it. This
      buffer keeps three role-tagged slots (functional sub-goal / distractor sub-goal /
      final goal) alive across the sequence and explicitly down-weights the distractor
      slot's contribution to the planning readout, rather than only encoding "the tool.

  CausalRelationModule
      Taylor et al. (2009, New Caledonian crows) found crows that solved a trap-tube
      transferred to a trap-TABLE that shared no visual features at all (different shape,
      colour, material) -- evidence they'd abstracted a causal relation (object-hole
      interaction), not a perceptually-bound rule. This module encodes a "causal" embedding
      via a bottleneck trained to predict task outcome while a gradient-reversal adversary
      actively strips out predictability of *which surface/apparatus context* produced it --
      forcing the representation to be invariant to surface appearance, which is exactly
      the property needed for trap-tube -> trap-table-style analogical transfer.

InhibitionModule reflects a recurring theme across the rook and NC crow trap-tube papers:
success correlated with willingness to switch sides / stop an already-started pull (Bird &
Emery 2009; Taylor et al. 2009 NC crows) -- i.e., inhibitory control over a prepotent
response, not just correct classification. It's implemented as a learned partial gate on
the output logits rather than a hard mask, since inhibition in these studies was graded
(some individuals showed partial hesitation), not all-or-nothing.
"""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


# ====================== GENERIC DNC-STYLE MEMORY ======================
class CorvidaeMemory(nn.Module):
    """
    A simplified differentiable neural computer. Content-addressed read/write over a
    persistent (B, memory_size, word_size) matrix, with `num_read_heads` independent
    read heads. State is lazily created on first forward() call for a given batch size
    and cleared by `reset()` -- mirroring the lazy-init/explicit-reset convention used by
    the stateful specialists in corvid_extensions.py.

    `consolidate(strength)` blends the working memory into a slower long-term buffer and
    partially blends that consolidated trace back into working memory -- a crude
    analogue of sleep-dependent replay consolidation. Call this periodically (e.g. every
    `consolidate_every` training steps), not every step.
    """
    def __init__(self, input_dim: int, memory_size: int = 32, word_size: int = 16,
                 num_read_heads: int = 2):
        super().__init__()
        self.memory_size = memory_size
        self.word_size = word_size
        self.num_read_heads = num_read_heads

        interface_size = (num_read_heads * word_size) + num_read_heads + word_size + 1 + word_size + word_size
        self.interface_proj = nn.Linear(input_dim, interface_size)

        self._memory: Optional[torch.Tensor] = None       # (B, N, W) working memory
        self._long_term: Optional[torch.Tensor] = None     # (B, N, W) consolidated trace

    def _ensure_state(self, batch_size: int, device):
        if self._memory is None or self._memory.size(0) != batch_size:
            self._memory = torch.zeros(batch_size, self.memory_size, self.word_size, device=device)
        if self._long_term is None or self._long_term.size(0) != batch_size:
            self._long_term = torch.zeros(batch_size, self.memory_size, self.word_size, device=device)

    def reset(self):
        """Clear memory. Lazily re-created on the next forward() call, whatever batch
        size that call uses -- so this needs no batch_size argument."""
        self._memory = None
        self._long_term = None

    @torch.no_grad()
    def consolidate(self, strength: float = 0.16):
        if self._memory is None or self._long_term is None:
            return
        self._long_term = (1 - strength) * self._long_term + strength * self._memory.detach()
        self._memory = (1 - strength * 0.5) * self._memory + (strength * 0.5) * self._long_term

    def forward(self, control: torch.Tensor) -> torch.Tensor:
        """
        control: (batch, seq_len, input_dim)
        Returns: (batch, seq_len, num_read_heads * word_size) concatenated read vectors.
        """
        b, t, _ = control.shape
        self._ensure_state(b, control.device)

        interface = self.interface_proj(control)
        i = 0
        read_keys = interface[..., i: i + self.num_read_heads * self.word_size]
        i += self.num_read_heads * self.word_size
        read_keys = read_keys.view(b, t, self.num_read_heads, self.word_size)
        read_strengths = F.softplus(interface[..., i: i + self.num_read_heads]); i += self.num_read_heads
        write_key = interface[..., i: i + self.word_size]; i += self.word_size
        write_strength = torch.sigmoid(interface[..., i: i + 1]); i += 1
        erase = torch.sigmoid(interface[..., i: i + self.word_size]); i += self.word_size
        add = torch.tanh(interface[..., i: i + self.word_size]); i += self.word_size

        reads = []
        for step in range(t):
            wk = write_key[:, step, :]
            ws = write_strength[:, step, :]
            er = erase[:, step, :]
            ad = add[:, step, :]

            mem_n = F.normalize(self._memory, dim=-1)
            wk_n = F.normalize(wk, dim=-1).unsqueeze(1)
            write_weight = F.softmax(torch.einsum('bnd,bod->bn', mem_n, wk_n), dim=-1) * ws

            self._memory = self._memory * (1 - write_weight.unsqueeze(-1) * er.unsqueeze(1)) \
                + write_weight.unsqueeze(-1) * ad.unsqueeze(1)

            step_reads = []
            for h in range(self.num_read_heads):
                rk = read_keys[:, step, h, :]
                rs = read_strengths[:, step, h:h + 1]
                mem_n2 = F.normalize(self._memory, dim=-1)
                rk_n = F.normalize(rk, dim=-1).unsqueeze(1)
                read_weight = F.softmax(rs * torch.einsum('bnd,bod->bn', mem_n2, rk_n), dim=-1)
                read_vec = torch.einsum('bn,bnd->bd', read_weight, self._memory)
                step_reads.append(read_vec)
            reads.append(torch.cat(step_reads, dim=-1))

        return torch.stack(reads, dim=1)


# ====================== METATOOL PLANNING BUFFER ======================
class MetatoolPlanningBuffer(nn.Module):
    """
    Gruber et al. (2019) x2: NC crows solving multi-stage metatool problems kept mental
    representations of a functional sub-goal, a distractor sub-goal, and the final goal
    active while stages were out of sight, and specifically avoided distractor
    apparatuses/tools rather than just failing to encode them -- an active suppression,
    not an absence of representation.

    Three role-tagged slots (index 0 = functional sub-goal, 1 = distractor sub-goal,
    2 = final goal), initialized from learned role embeddings and slowly updated from
    context. A learned distractor gate actively suppresses slot 1's contribution to the
    read-out (rather than that slot simply never being written), and a stage-attention
    head decides which slot is most relevant to output at each step -- a crude analogue
    of shifting attention across sub-goal/goal representations as a multi-stage plan
    unfolds.

    State persists across a trial (call `clear()` at episode start, matching the
    lazy-reinit convention of CorvidaeMemory) since the whole point is remembering
    sub-goals while the relevant apparatus is out of sight.
    """
    def __init__(self, embedding_dim: int, num_slots: int = 3):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.num_slots = num_slots  # convention: 0=functional subgoal, 1=distractor, 2=final goal
        self.role_embeddings = nn.Parameter(torch.randn(num_slots, embedding_dim) * 0.02)
        self.slot_write = nn.Linear(embedding_dim, embedding_dim)
        self.distractor_gate = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim), nn.ReLU(), nn.Linear(embedding_dim, 1)
        )
        self.stage_attn = nn.Linear(embedding_dim, num_slots)
        self.output_proj = nn.Linear(embedding_dim, embedding_dim)

        self._slot_content: Optional[torch.Tensor] = None  # (B, num_slots, D)

    def clear(self):
        self._slot_content = None

    def _ensure_state(self, batch_size: int, device):
        if self._slot_content is None or self._slot_content.size(0) != batch_size:
            self._slot_content = self.role_embeddings.unsqueeze(0).expand(batch_size, -1, -1).clone().to(device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (batch, seq_len, embedding_dim)
        Returns: (batch, seq_len, embedding_dim) plan read-out, distractor-suppressed.
        """
        b, t, d = x.shape
        self._ensure_state(b, x.device)

        outs = []
        for step in range(t):
            xt = x[:, step, :]
            update = self.slot_write(xt)
            new_slot0 = 0.9 * self._slot_content[:, 0, :] + 0.1 * update
            self._slot_content = torch.cat([
                new_slot0.unsqueeze(1), self._slot_content[:, 1:, :]
            ], dim=1)

            stage_logits = self.stage_attn(xt)  # (B, num_slots)
            distractor_suppress = torch.sigmoid(self.distractor_gate(self._slot_content[:, 1, :]))  # (B,1)
            suppressed = stage_logits.clone()
            suppressed[:, 1] = suppressed[:, 1] - 5.0 * distractor_suppress.squeeze(-1)
            weights = F.softmax(suppressed, dim=-1)

            read = torch.einsum('bk,bkd->bd', weights, self._slot_content)
            outs.append(read)

        out = torch.stack(outs, dim=1)
        return self.output_proj(out)


# ====================== CAUSAL RELATION MODULE (DOMAIN-INVARIANT) ======================
class _GradientReversal(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, lambd):
        ctx.lambd = lambd
        return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output):
        return -ctx.lambd * grad_output, None


def grad_reverse(x: torch.Tensor, lambd: float = 1.0) -> torch.Tensor:
    return _GradientReversal.apply(x, lambd)


class CausalRelationModule(nn.Module):
    """
    Taylor et al. (2009): NC crows that solved a trap-tube transferred immediately to a
    trap-TABLE sharing no visual features (different shape, colour, material) -- the
    strongest available evidence for a causal/analogical representation rather than a
    perceptually-bound rule. The design choice that would produce this kind of transfer
    is a representation that predicts task-relevant outcome while being actively
    prevented from encoding which surface/apparatus context it came from.

    Implemented via a small domain-adversarial bottleneck: a causal encoder maps context
    to a `causal_dim` embedding; an outcome head predicts the task-relevant
    approach/avoid signal from it; a gradient-reversal adversary tries to predict which
    of `num_surface_contexts` surface contexts produced the embedding, and its gradient
    is negated before reaching the causal encoder -- so the encoder is pushed to become
    WORSE at leaking surface identity even as it stays predictive of outcome. That
    invariance is what should let a decision rule learned in one surface context
    transfer to a perceptually distinct one, as in the crows' trap-table transfer.

    `forward` returns only the causal content embedding (for drop-in compatibility with
    the rest of the fusion pipeline); the outcome/surface logits needed for the
    adversarial training objective are stashed on `self.last_outcome_logit` /
    `self.last_surface_logits` after each call.
    """
    def __init__(self, embedding_dim: int, causal_dim: Optional[int] = None,
                 num_surface_contexts: int = 8, grl_lambda: float = 1.0):
        super().__init__()
        causal_dim = causal_dim or embedding_dim
        self.causal_encoder = nn.Sequential(
            nn.Linear(embedding_dim, causal_dim), nn.ReLU(), nn.Linear(causal_dim, causal_dim)
        )
        self.outcome_head = nn.Sequential(
            nn.Linear(causal_dim, causal_dim), nn.ReLU(), nn.Linear(causal_dim, 1)
        )
        self.surface_adversary = nn.Sequential(
            nn.Linear(causal_dim, causal_dim), nn.ReLU(), nn.Linear(causal_dim, num_surface_contexts)
        )
        self.output_proj = nn.Linear(causal_dim, embedding_dim)
        self.grl_lambda = grl_lambda

        self.last_outcome_logit: Optional[torch.Tensor] = None
        self.last_surface_logits: Optional[torch.Tensor] = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        causal_repr = self.causal_encoder(x)
        self.last_outcome_logit = self.outcome_head(causal_repr).squeeze(-1)
        reversed_repr = grad_reverse(causal_repr, self.grl_lambda)
        self.last_surface_logits = self.surface_adversary(reversed_repr)
        return self.output_proj(causal_repr)

    @staticmethod
    def outcome_loss(outcome_logit: torch.Tensor, outcome_labels: torch.Tensor) -> torch.Tensor:
        return F.binary_cross_entropy_with_logits(outcome_logit, outcome_labels)

    @staticmethod
    def surface_adversary_loss(surface_logits: torch.Tensor, surface_labels: torch.Tensor) -> torch.Tensor:
        """Trained normally (not reversed again) -- the reversal already happened on the
        forward pass via grad_reverse; this loss just gives the adversary a real
        objective to (try to) solve, which is what makes the reversed gradient
        meaningful signal for the causal encoder."""
        return F.cross_entropy(surface_logits.reshape(-1, surface_logits.size(-1)), surface_labels.reshape(-1))


# ====================== INHIBITORY CONTROL ======================
class InhibitionModule(nn.Module):
    """
    Bird & Emery (2009, rooks) and Taylor et al. (2009, NC crows): success on the
    trap-tube correlated with willingness to switch/inhibit an already-started response
    (side-switching), and failure sometimes looked like partial hesitation rather than
    an outright inability to represent the task (their crow Espanol repeatedly
    "momentarily halted pulling actions... before eventually pulling it into the trap").
    That's graded inhibition, not a hard veto -- so this is a soft, learned partial gate
    on the output logits rather than a mask that can zero them out entirely.
    """
    def __init__(self, embedding_dim: int, num_classes: int):
        super().__init__()
        self.gate_proj = nn.Linear(embedding_dim, num_classes)

    def forward(self, inhibition_context: torch.Tensor, logits: torch.Tensor) -> torch.Tensor:
        gate = torch.sigmoid(self.gate_proj(inhibition_context)).unsqueeze(1)  # (B,1,num_classes)
        return logits * (0.5 + 0.5 * gate)


# ====================== FULL BASE MODEL ======================
class Corvidae(nn.Module):
    """
    The base architecture. Sensory encoding -> two lateralized "hemisphere" streams
    (ncl_left / ncl_right, named for the nidopallium caudolaterale, a key executive-
    function region in corvid brains per Marzluff et al. 2012 and the Swift/Marzluff/
    Cross dead-conspecific imaging work) blended into a shared executive stream ->
    generic DNC memory + metatool planning buffer + causal relation module, fused with
    lateralized threat processing and a social-context head -> final encoder -> output
    logits, optionally passed through inhibitory gating.

    Subclasses (CorvidaeMultiSpecies, CorvidaeAviary) replace the plain `mem_proj`
    contribution to fusion with a routed bank of species-specific specialists, but this
    base class is fully runnable standalone.
    """
    def __init__(self, num_embeddings: int, embedding_dim: int, max_seq_len: int,
                 nhead: int = 4, dim_feedforward: int = 128, memory_size: int = 32,
                 memory_word_size: int = 16, num_read_heads: int = 2,
                 num_classes: Optional[int] = None, num_sensory_layers: int = 1,
                 num_ncl_layers: int = 1, num_final_layers: int = 1, dropout: float = 0.1,
                 num_tool_classes: int = 8, num_surface_contexts: int = 8):
        super().__init__()
        num_classes = num_classes or num_embeddings
        self.embedding_dim = embedding_dim
        self.max_seq_len = max_seq_len

        self.token_emb = nn.Embedding(num_embeddings, embedding_dim)
        self.pos_emb = nn.Embedding(max_seq_len, embedding_dim)
        self.emb_dropout = nn.Dropout(dropout)

        def make_encoder(num_layers):
            layer = nn.TransformerEncoderLayer(
                embedding_dim, nhead, dim_feedforward, dropout=dropout, batch_first=True
            )
            return nn.TransformerEncoder(layer, num_layers=num_layers)

        self.sensory = make_encoder(num_sensory_layers)
        self.ncl_left = make_encoder(num_ncl_layers)
        self.ncl_right = make_encoder(num_ncl_layers)
        self.final = make_encoder(num_final_layers)

        self.lateral_strength = nn.Parameter(torch.tensor(0.5))

        self.controller = nn.Sequential(nn.Linear(embedding_dim, embedding_dim), nn.ReLU())
        self.memory = CorvidaeMemory(embedding_dim, memory_size, memory_word_size, num_read_heads)
        self.mem_proj = nn.Linear(num_read_heads * memory_word_size, embedding_dim)

        self.plan_buffer = MetatoolPlanningBuffer(embedding_dim)
        self.plan_proj = nn.Linear(embedding_dim, embedding_dim)

        self.causal_relation = CausalRelationModule(embedding_dim, num_surface_contexts=num_surface_contexts)

        self.threat_left = nn.Linear(embedding_dim, embedding_dim)
        self.threat_right = nn.Linear(embedding_dim, embedding_dim)
        self.social_head = nn.Linear(embedding_dim, embedding_dim)

        fusion_in = embedding_dim * 6  # executive, species/mem, plan, causal, threat, social
        self.fusion_gate = nn.Linear(fusion_in, embedding_dim)
        self.fusion_dropout = nn.Dropout(dropout)

        self.output_proj = nn.Linear(embedding_dim, num_classes)
        self.inhibition = InhibitionModule(embedding_dim, num_classes)

        self.planning_head = nn.Linear(embedding_dim, embedding_dim)
        self.tool_predictor = nn.Linear(embedding_dim, num_tool_classes)

    def reset_memory(self):
        """No batch_size needed -- both the DNC memory and the plan buffer lazily
        re-create their state on the next forward() call, using whatever batch size
        that call provides."""
        self.memory.reset()
        self.plan_buffer.clear()

    def consolidate_memory(self, strength: float = 0.16):
        self.memory.consolidate(strength)

    def forward(self, x: torch.Tensor, return_planning: bool = False, use_inhibition: bool = True):
        seq_len = x.size(1)
        if seq_len > self.max_seq_len:
            raise ValueError(f"Sequence too long: {seq_len}")

        causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=x.device)

        emb = self.token_emb(x) + self.pos_emb(torch.arange(seq_len, device=x.device))
        emb = self.emb_dropout(emb)
        sensory = self.sensory(emb, mask=causal_mask, is_causal=True)

        left = self.ncl_left(sensory, mask=causal_mask, is_causal=True)
        right = self.ncl_right(sensory, mask=causal_mask, is_causal=True)
        lateral = torch.sigmoid(self.lateral_strength)
        executive = (1 - lateral) * left + lateral * right

        control = self.controller(executive)
        mem_read = self.memory(control)
        mem_proj = self.mem_proj(mem_read)

        plan_read = self.plan_buffer(executive)
        plan_proj = self.plan_proj(plan_read)

        causal_ctx = self.causal_relation(executive)

        threat_l = self.threat_left(executive.mean(dim=1))
        threat_r = self.threat_right(executive.mean(dim=1))
        threat_emb = (threat_l + 1.3 * threat_r) * 0.5  # right-hemisphere threat bias, Marzluff et al. 2012

        fused_cat = torch.cat([
            executive, mem_proj, plan_proj, causal_ctx,
            threat_emb.unsqueeze(1).expand(-1, seq_len, -1),
            self.social_head(executive.mean(dim=1)).unsqueeze(1).expand(-1, seq_len, -1),
        ], dim=-1)

        gate = torch.sigmoid(self.fusion_gate(fused_cat))
        gate = self.fusion_dropout(gate)
        fused = executive + gate * (mem_proj + plan_proj + causal_ctx)

        final = self.final(fused, mask=causal_mask, is_causal=True)
        logits = self.output_proj(final)

        if use_inhibition:
            inhibition_context = (causal_ctx.mean(dim=1) + plan_proj.mean(dim=1)) * 0.5
            logits = self.inhibition(inhibition_context, logits)

        if return_planning:
            plan = self.planning_head(final.mean(dim=1))
            tool_logits = self.tool_predictor(plan)
            return logits, tool_logits
        return logits


# ====================== BASE TRAINER ======================
class CorvidaeTrainer:
    """
    Minimal training loop for the base Corvidae model. Subclassed by
    SpeciesAwareTrainer / AviaryTrainer to add specialist-specific auxiliary losses.
    """
    def __init__(self, model: "Corvidae", lr: float = 3e-4, replay_weight: float = 0.35,
                 consolidate_every: int = 40):
        self.model = model
        self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        self.replay_weight = replay_weight
        self.consolidate_every = consolidate_every
        self.global_step = 0

    def train_step(self, input_ids: torch.Tensor, targets: torch.Tensor,
                    replay_targets: Optional[torch.Tensor] = None):
        self.model.train()
        self.optimizer.zero_grad()

        logits = self.model(input_ids)
        ce_loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))

        replay_loss = torch.tensor(0.0, device=logits.device)
        if replay_targets is not None:
            replay_loss = F.mse_loss(logits.mean(dim=1), replay_targets)

        total_loss = ce_loss + self.replay_weight * replay_loss
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        self.optimizer.step()

        self.global_step += 1
        if self.global_step % self.consolidate_every == 0:
            self.model.consolidate_memory(strength=0.16)

        return {"total_loss": total_loss.item(), "ce_loss": ce_loss.item(), "replay_loss": replay_loss.item()}


if __name__ == "__main__":
    torch.manual_seed(0)
    vocab = 200
    model = Corvidae(num_embeddings=vocab, embedding_dim=64, max_seq_len=32, nhead=4,
                      dim_feedforward=128, memory_size=32, memory_word_size=16,
                      num_read_heads=2, num_classes=vocab)
    trainer = CorvidaeTrainer(model, consolidate_every=2)

    x = torch.randint(0, vocab, (2, 16))
    y = torch.randint(0, vocab, (2, 16))
    for step in range(5):
        stats = trainer.train_step(x, y)
        print(step, stats)

    model.reset_memory()
    logits, tool_logits = model(x, return_planning=True)
    print("logits shape:", logits.shape, "tool_logits shape:", tool_logits.shape)
    print("smoke test passed")