surface_code_in_stem/rl_control/nested_agent.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Any, Tuple, Optional
import numpy as np

from .titans import TitansMemory
from .replay_buffer import ReplayBuffer, Experience


class BaseAgent:
    """Base class for RL agents."""
    def __init__(self):
        pass

    def act(self, state: Any) -> Any:
        raise NotImplementedError

    def update(self, *args, **kwargs):
        raise NotImplementedError

class NestedLearningAgent(BaseAgent):
    """
    Nested Learning Agent for QEC with experience replay and proper RL.

    Implements nested optimization with:
    - Inner loop: Fast updates to decoder weights using policy gradients
    - Outer loop: Slow updates to memory/hyperparameters
    - Experience replay buffer for sample efficiency
    - Actor-critic architecture for stable learning
    """

    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        hidden_dim: int = 64,
        replay_capacity: int = 10000,
        gamma: float = 0.99,
        gae_lambda: float = 0.95,
    ):
        super().__init__()
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.gae_lambda = gae_lambda

        # TITANS memory for long-term syndrome history
        self.memory_module = TitansMemory(input_dim=state_dim, hidden_dim=hidden_dim)

        # Actor (policy network)
        self.actor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )

        # Critic (value network) for advantage estimation
        self.critic = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )

        # Inner loop optimizers (fast)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=1e-3)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)

        # Outer loop optimizer (slow)
        self.outer_optimizer = torch.optim.Adam(self.memory_module.parameters(), lr=1e-4)

        # Experience replay buffer
        self.replay_buffer = ReplayBuffer(capacity=replay_capacity, prioritized=True)

        # Training statistics
        self.training_step = 0
        self.episode_rewards = []

    def act(self, state: torch.Tensor, deterministic: bool = False) -> Tuple[int, torch.Tensor]:
        """Select action using current policy.

        Returns:
            Tuple of (action_index, log_probability)
        """
        with torch.no_grad():
            if state.dim() == 1:
                state = state.unsqueeze(0)

            context = self.memory_module(state)
            logits = self.actor(context)
            probs = F.softmax(logits, dim=-1)

            if deterministic:
                action = torch.argmax(probs, dim=-1)
            else:
                action = torch.multinomial(probs, 1).squeeze(-1)

            # Compute log probability
            log_prob = F.log_softmax(logits, dim=-1)
            action_log_prob = log_prob.gather(-1, action.unsqueeze(-1)).squeeze(-1)

        return action.item(), action_log_prob

    def compute_advantages(
        self,
        rewards: torch.Tensor,
        values: torch.Tensor,
        dones: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute GAE advantages and returns.

        Args:
            rewards: Tensor of rewards (batch_size,)
            values: Tensor of value estimates (batch_size,)
            dones: Tensor of done flags (batch_size,)

        Returns:
            Tuple of (advantages, returns)
        """
        batch_size = rewards.shape[0]
        advantages = torch.zeros_like(rewards)
        returns = torch.zeros_like(rewards)

        # Compute GAE advantages
        next_advantage = 0
        for t in reversed(range(batch_size)):
            if t == batch_size - 1:
                next_value = 0
            else:
                next_value = values[t + 1]

            delta = rewards[t] + self.gamma * next_value * (1 - dones[t]) - values[t]
            advantages[t] = delta + self.gamma * self.gae_lambda * next_advantage * (1 - dones[t])
            next_advantage = advantages[t]

        returns = advantages + values

        return advantages, returns

    def inner_loop(
        self,
        batch_size: int = 32,
        clip_epsilon: float = 0.2,
        value_coef: float = 0.5,
        entropy_coef: float = 0.01,
    ) -> Tuple[float, float, float]:
        """Policy gradient update using experience replay.

        Implements PPO-style clipped surrogate objective with GAE advantages.

        Returns:
            Tuple of (actor_loss, critic_loss, entropy)
        """
        if not self.replay_buffer.is_ready(batch_size):
            return 0.0, 0.0, 0.0

        # Sample from replay buffer
        states, actions, rewards, next_states, dones, weights, indices = self.replay_buffer.sample(batch_size)

        # Prepare states for TITANS (add sequence dimension if needed)
        if states.dim() == 2:
            states = states.unsqueeze(1)
            next_states = next_states.unsqueeze(1)

        # Get memory context
        context = self.memory_module(states.squeeze(1) if states.dim() == 3 else states)

        # Compute current policy logits and values
        logits = self.actor(context)
        values = self.critic(context).squeeze(-1)

        # Compute log probabilities of actions
        log_probs = F.log_softmax(logits, dim=-1)
        action_log_probs = log_probs.gather(-1, actions.unsqueeze(-1)).squeeze(-1)

        # Compute entropy for exploration
        probs = F.softmax(logits, dim=-1)
        entropy = -(probs * log_probs).sum(dim=-1).mean()

        # Compute advantages using GAE
        advantages, returns = self.compute_advantages(rewards, values, dones)

        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # PPO-style clipped surrogate loss (without the ratio for simplicity, using vanilla PG)
        policy_loss = -(action_log_probs * advantages * weights).mean()

        # Value loss
        value_loss = F.mse_loss(values, returns, reduction='none')
        value_loss = (value_loss * weights).mean()

        # Combined loss
        actor_loss = policy_loss - entropy_coef * entropy
        critic_loss = value_coef * value_loss

        # Update actor
        self.actor_optimizer.zero_grad()
        actor_loss.backward(retain_graph=True)
        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
        self.actor_optimizer.step()

        # Update critic
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
        self.critic_optimizer.step()

        # Update replay buffer priorities (using TD errors as priorities)
        with torch.no_grad():
            td_errors = torch.abs(returns - values).cpu().numpy()
        self.replay_buffer.update_priorities(indices, td_errors)

        self.training_step += 1

        return policy_loss.item(), value_loss.item(), entropy.item()

    def outer_loop(self, states: torch.Tensor, performance_metric: float) -> float:
        """Slow update: Update memory representations and hyperparameters.

        Uses meta-learning signal (performance_metric) to adapt the memory module.
        """
        if states.dim() == 2:
            states = states.unsqueeze(1)

        self.outer_optimizer.zero_grad()

        # Get memory context
        context = self.memory_module(states.squeeze(1) if states.dim() == 3 else states)

        # Meta-learning loss: encourage representations that predict performance
        # Use auxiliary task: predict the performance metric from context
        performance_prediction = self.critic(context).squeeze(-1).mean()

        # Loss encourages accurate performance prediction
        meta_loss = F.mse_loss(performance_prediction, torch.tensor(performance_metric))

        # Add regularization to encourage diverse representations
        context_variance = context.var(dim=0).mean()
        diversity_bonus = -0.01 * context_variance  # Encourage variance

        total_loss = meta_loss + diversity_bonus
        total_loss.backward()

        # Gradient clipping for stability
        torch.nn.utils.clip_grad_norm_(self.memory_module.parameters(), 0.5)
        self.outer_optimizer.step()

        return total_loss.item()

    def store_experience(
        self,
        state: torch.Tensor,
        action: int,
        reward: float,
        next_state: torch.Tensor,
        done: bool,
        priority: Optional[float] = None,
    ) -> None:
        """Store experience in replay buffer with optional priority."""
        experience = Experience(
            state=state if state.dim() > 0 else state.unsqueeze(0),
            action=action,
            reward=reward,
            next_state=next_state if next_state.dim() > 0 else next_state.unsqueeze(0),
            done=done,
        )
        self.replay_buffer.push(experience, priority)

    def get_stats(self) -> dict:
        """Return training statistics."""
        return {
            "training_step": self.training_step,
            "buffer_size": len(self.replay_buffer),
            "episode_rewards": self.episode_rewards[-10:] if self.episode_rewards else [],
        }

    def save_checkpoint(self, path: str) -> None:
        """Save agent state to checkpoint."""
        torch.save({
            "memory_module": self.memory_module.state_dict(),
            "actor": self.actor.state_dict(),
            "critic": self.critic.state_dict(),
            "actor_optimizer": self.actor_optimizer.state_dict(),
            "critic_optimizer": self.critic_optimizer.state_dict(),
            "outer_optimizer": self.outer_optimizer.state_dict(),
            "training_step": self.training_step,
        }, path)

    def load_checkpoint(self, path: str) -> None:
        """Load agent state from checkpoint."""
        checkpoint = torch.load(path, map_location="cpu")
        self.memory_module.load_state_dict(checkpoint["memory_module"])
        self.actor.load_state_dict(checkpoint["actor"])
        self.critic.load_state_dict(checkpoint["critic"])
        self.actor_optimizer.load_state_dict(checkpoint["actor_optimizer"])
        self.critic_optimizer.load_state_dict(checkpoint["critic_optimizer"])
        self.outer_optimizer.load_state_dict(checkpoint["outer_optimizer"])
        self.training_step = checkpoint.get("training_step", 0)