scripts/train_boolean_mmep.py

#!/usr/bin/env python3
"""
Train Boolean actor network using MMEP (Min-Max Equilibrium Propagation).

License: BSL-1.1
Contact: jessica@collectivekitty.com

MMEP training loop:
  1. Free phase: run network forward to equilibrium (energy minimization)
  2. Nudge phase: clamp output toward target, re-equilibrate with nudge factor beta
  3. Local update: compute weight updates from difference in activations between phases
  4. Huntington regularization: penalize violations of Boolean algebra postulates

The Boolean perceptron actors enforce algebraic structure while MMEP provides
biologically plausible gradient-free learning.
"""

import argparse
import math
import sys
import time
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml


class BooleanPerceptronActor(nn.Module):
    """Single Boolean perceptron actor with Huntington-compatible operations."""

    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.weight = nn.Parameter(torch.randn(output_dim, input_dim) * 0.01)
        self.bias = nn.Parameter(torch.zeros(output_dim))
        # Boolean operation parameters (learnable)
        self.or_weight = nn.Parameter(torch.randn(output_dim, 2 * output_dim) * 0.01)
        self.and_weight = nn.Parameter(torch.randn(output_dim, 2 * output_dim) * 0.01)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through sigmoid activation (Boolean-compatible)."""
        return torch.sigmoid(F.linear(x, self.weight, self.bias))

    def boolean_or(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Approximate Boolean OR via learned gate."""
        combined = torch.cat([a, b], dim=-1)
        return torch.sigmoid(F.linear(combined, self.or_weight))

    def boolean_and(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Approximate Boolean AND via learned gate."""
        combined = torch.cat([a, b], dim=-1)
        return torch.sigmoid(F.linear(combined, self.and_weight))


class BooleanActorNetwork(nn.Module):
    """Network of Boolean perceptron actors trained with MMEP."""

    def __init__(self, layer_dims: list, beta: float = 0.1):
        super().__init__()
        self.beta = beta
        self.layers = nn.ModuleList()
        for i in range(len(layer_dims) - 1):
            self.layers.append(BooleanPerceptronActor(layer_dims[i], layer_dims[i + 1]))
        self.layer_dims = layer_dims

    def free_phase(self, x: torch.Tensor, num_steps: int = 10) -> list:
        """Run free phase to equilibrium. Returns activations at each layer."""
        activations = [x]
        h = x
        for layer in self.layers:
            h = layer(h)
            activations.append(h)
        # Iterate to equilibrium
        for _ in range(num_steps - 1):
            h = activations[0]
            for i, layer in enumerate(self.layers):
                h = layer(h)
                activations[i + 1] = h
        return activations

    def nudge_phase(self, x: torch.Tensor, target: torch.Tensor, num_steps: int = 10) -> list:
        """Run nudge phase with target clamping. Returns activations at each layer."""
        activations = [x]
        h = x
        for layer in self.layers:
            h = layer(h)
            activations.append(h)
        # Nudge final layer toward target
        for _ in range(num_steps - 1):
            # Clamp output toward target
            activations[-1] = (1 - self.beta) * activations[-1] + self.beta * target
            h = activations[0]
            for i, layer in enumerate(self.layers):
                h = layer(h)
                if i < len(self.layers) - 1:
                    activations[i + 1] = h
                else:
                    activations[i + 1] = (1 - self.beta) * h + self.beta * target
        return activations

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Standard forward pass for inference."""
        h = x
        for layer in self.layers:
            h = layer(h)
        return h


def huntington_regularization(network: BooleanActorNetwork, x: torch.Tensor) -> torch.Tensor:
    """Compute Huntington postulate violation as regularization loss.

    Penalizes violations of:
      - Commutativity: OR(a,b) = OR(b,a), AND(a,b) = AND(b,a)
      - Idempotence: OR(a,a) = a, AND(a,a) = a
      - Complement: OR(a, 1-a) ~ 1, AND(a, 1-a) ~ 0
    """
    loss = torch.tensor(0.0, device=x.device)
    count = 0

    for layer in network.layers:
        # Generate test activations
        a = torch.sigmoid(torch.randn(x.shape[0], layer.output_dim, device=x.device))
        b = torch.sigmoid(torch.randn(x.shape[0], layer.output_dim, device=x.device))

        # Commutativity: OR(a,b) = OR(b,a)
        or_ab = layer.boolean_or(a, b)
        or_ba = layer.boolean_or(b, a)
        loss = loss + F.mse_loss(or_ab, or_ba)

        # Commutativity: AND(a,b) = AND(b,a)
        and_ab = layer.boolean_and(a, b)
        and_ba = layer.boolean_and(b, a)
        loss = loss + F.mse_loss(and_ab, and_ba)

        # Idempotence: OR(a,a) = a
        or_aa = layer.boolean_or(a, a)
        loss = loss + F.mse_loss(or_aa, a)

        # Idempotence: AND(a,a) = a
        and_aa = layer.boolean_and(a, a)
        loss = loss + F.mse_loss(and_aa, a)

        # Complement: OR(a, 1-a) ~ 1
        comp_a = 1.0 - a
        or_comp = layer.boolean_or(a, comp_a)
        loss = loss + F.mse_loss(or_comp, torch.ones_like(or_comp))

        # Complement: AND(a, 1-a) ~ 0
        and_comp = layer.boolean_and(a, comp_a)
        loss = loss + F.mse_loss(and_comp, torch.zeros_like(and_comp))

        count += 6

    return loss / max(count, 1)


def compute_mmep_update(
    free_activations: list,
    nudge_activations: list,
    network: BooleanActorNetwork,
    lr: float,
):
    """Compute and apply MMEP local weight updates.

    The update rule is:
      dW_i = (1/beta) * (s_i^nudge * s_{i-1}^nudge^T - s_i^free * s_{i-1}^free^T)
    """
    beta = network.beta
    for i, layer in enumerate(network.layers):
        # Pre-synaptic and post-synaptic activations
        pre_free = free_activations[i]
        post_free = free_activations[i + 1]
        pre_nudge = nudge_activations[i]
        post_nudge = nudge_activations[i + 1]

        # Compute correlation differences
        # dW = (1/beta) * (post_nudge @ pre_nudge^T - post_free @ pre_free^T) / batch_size
        batch_size = pre_free.shape[0]
        corr_nudge = torch.mm(post_nudge.t(), pre_nudge) / batch_size
        corr_free = torch.mm(post_free.t(), pre_free) / batch_size

        dW = (1.0 / beta) * (corr_nudge - corr_free)
        db = (1.0 / beta) * (post_nudge - post_free).mean(dim=0)

        # Apply updates
        with torch.no_grad():
            layer.weight.add_(lr * dW)
            layer.bias.add_(lr * db)


def load_config(config_path: str) -> dict:
    """Load configuration from YAML file."""
    path = Path(config_path)
    if not path.exists():
        print(f"Warning: Config file {config_path} not found. Using defaults.")
        return {}
    with open(path, "r") as f:
        return yaml.safe_load(f) or {}


def generate_boolean_data(batch_size: int, input_dim: int, device: torch.device):
    """Generate synthetic Boolean data for training."""
    x = torch.randint(0, 2, (batch_size, input_dim), dtype=torch.float32, device=device)
    # Target: XOR of first half and second half (non-trivial Boolean function)
    half = input_dim // 2
    target_bits = ((x[:, :half] + x[:, half:2*half]) % 2)
    # Pad or truncate to match output dim
    return x, target_bits


def main():
    parser = argparse.ArgumentParser(
        description="Train Boolean actor network with MMEP (Min-Max Equilibrium Propagation)."
    )
    parser.add_argument("--config", type=str, default="config/boolean_actor.yaml",
                        help="Path to config YAML (default: config/boolean_actor.yaml)")
    parser.add_argument("--epochs", type=int, default=None, help="Number of epochs (overrides config)")
    parser.add_argument("--device", type=str, default=None, help="Device (overrides config)")
    args = parser.parse_args()

    # Load config
    config = load_config(args.config)

    # Extract parameters with defaults
    model_cfg = config.get("model", {})
    train_cfg = config.get("training", {})
    mmep_cfg = config.get("mmep", {})

    input_dim = model_cfg.get("input_dim", 256)
    hidden_dims = model_cfg.get("hidden_dims", [512, 256, 128])
    output_dim = model_cfg.get("output_dim", 128)
    layer_dims = [input_dim] + hidden_dims + [output_dim]

    epochs = args.epochs or train_cfg.get("epochs", 100)
    batch_size = train_cfg.get("batch_size", 256)
    lr = train_cfg.get("learning_rate", 0.01)
    huntington_weight = train_cfg.get("huntington_weight", 0.1)
    log_interval = train_cfg.get("log_interval", 10)

    beta = mmep_cfg.get("beta", 0.1)
    free_steps = mmep_cfg.get("free_steps", 10)
    nudge_steps = mmep_cfg.get("nudge_steps", 10)

    device_str = args.device or train_cfg.get("device", "cpu")
    device = torch.device(device_str)

    print("BURT-IMMA Boolean MMEP Training")
    print("=" * 60)
    print(f"  Config:           {args.config}")
    print(f"  Layer dims:       {layer_dims}")
    print(f"  Epochs:           {epochs}")
    print(f"  Batch size:       {batch_size}")
    print(f"  Learning rate:    {lr}")
    print(f"  Beta (nudge):     {beta}")
    print(f"  Free steps:       {free_steps}")
    print(f"  Nudge steps:      {nudge_steps}")
    print(f"  Huntington wt:    {huntington_weight}")
    print(f"  Device:           {device}")
    print("=" * 60)
    print()

    # Create network
    network = BooleanActorNetwork(layer_dims, beta=beta).to(device)
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)

    # Training loop
    metrics_history = []
    start_time = time.time()

    for epoch in range(1, epochs + 1):
        epoch_loss = 0.0
        epoch_accuracy = 0.0
        epoch_huntington = 0.0
        epoch_entropy = 0.0
        num_batches = 0

        # Generate training data
        steps_per_epoch = max(1, 1000 // batch_size)
        for step in range(steps_per_epoch):
            x, target = generate_boolean_data(batch_size, input_dim, device)

            # Ensure target matches output dim
            if target.shape[1] != output_dim:
                # Tile or truncate
                if target.shape[1] < output_dim:
                    repeats = output_dim // target.shape[1] + 1
                    target = target.repeat(1, repeats)[:, :output_dim]
                else:
                    target = target[:, :output_dim]

            # === MMEP Training ===
            # Phase 1: Free phase
            free_acts = network.free_phase(x, num_steps=free_steps)

            # Phase 2: Nudge phase
            nudge_acts = network.nudge_phase(x, target, num_steps=nudge_steps)

            # Phase 3: Local update (MMEP rule)
            compute_mmep_update(free_acts, nudge_acts, network, lr)

            # Compute task loss for logging
            with torch.no_grad():
                output = network(x)
                task_loss = F.mse_loss(output, target)

            # Phase 4: Huntington regularization (uses gradient)
            optimizer.zero_grad()
            h_loss = huntington_regularization(network, x)
            reg_loss = huntington_weight * h_loss
            reg_loss.backward()
            optimizer.step()

            # Compute metrics
            with torch.no_grad():
                output = network(x)
                predictions = (output > 0.5).float()
                accuracy = (predictions == target).float().mean().item()
                # Binary entropy of outputs
                p = torch.clamp(output, 1e-7, 1 - 1e-7)
                entropy = -(p * torch.log(p) + (1 - p) * torch.log(1 - p)).mean().item()

            epoch_loss += task_loss.item()
            epoch_accuracy += accuracy
            epoch_huntington += h_loss.item()
            epoch_entropy += entropy
            num_batches += 1

        # Average metrics
        avg_loss = epoch_loss / num_batches
        avg_accuracy = epoch_accuracy / num_batches
        avg_huntington = epoch_huntington / num_batches
        avg_entropy = epoch_entropy / num_batches

        metrics = {
            "epoch": epoch,
            "loss": avg_loss,
            "accuracy": avg_accuracy,
            "huntington_violation": avg_huntington,
            "entropy": avg_entropy,
        }
        metrics_history.append(metrics)

        if epoch % log_interval == 0 or epoch == 1:
            elapsed = time.time() - start_time
            print(
                f"Epoch {epoch:4d}/{epochs} | "
                f"loss: {avg_loss:.6f} | "
                f"acc: {avg_accuracy:.4f} | "
                f"hunt_viol: {avg_huntington:.6f} | "
                f"entropy: {avg_entropy:.4f} | "
                f"time: {elapsed:.1f}s"
            )

    # Final summary
    elapsed = time.time() - start_time
    print()
    print("=" * 60)
    print("Training Complete")
    print("=" * 60)
    print(f"  Total time:              {elapsed:.1f}s")
    print(f"  Final loss:              {metrics_history[-1]['loss']:.6f}")
    print(f"  Final accuracy:          {metrics_history[-1]['accuracy']:.4f}")
    print(f"  Final Huntington viol.:  {metrics_history[-1]['huntington_violation']:.6f}")
    print(f"  Final entropy:           {metrics_history[-1]['entropy']:.4f}")

    # Save model
    save_path = Path("checkpoints/boolean_mmep_latest.pt")
    save_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({
        "model_state_dict": network.state_dict(),
        "config": config,
        "metrics": metrics_history,
        "layer_dims": layer_dims,
    }, save_path)
    print(f"  Model saved to:          {save_path}")


if __name__ == "__main__":
    main()