"""AAM Diffusion LLM — MCTS Reasoning Engine AlphaZero-style tree search for reasoning about narrative arrangement from graph evidence. AAM-specific: each node = a sentence arrangement, value = narrative coherence, policy = next arrangement step. """ from __future__ import annotations import math from dataclasses import dataclass, field from typing import Optional, List, Dict, Any, Tuple import torch import torch.nn as nn import torch.nn.functional as F class MCTSConfig: def __init__( self, num_simulations: int = 64, c_puct: float = 1.5, temperature: float = 1.0, max_depth: int = 10, use_value_network: bool = True, use_progressive_widening: bool = True, max_children: int = 8, ) -> None: self.num_simulations = num_simulations self.c_puct = c_puct self.temperature = temperature self.max_depth = max_depth self.use_value_network = use_value_network self.use_progressive_widening = use_progressive_widening self.max_children = max_children if num_simulations <= 0: raise ValueError(f"num_simulations must be positive, got {num_simulations}") if c_puct <= 0: raise ValueError(f"c_puct must be positive, got {c_puct}") if max_depth <= 0: raise ValueError(f"max_depth must be positive, got {max_depth}") @dataclass class MCTSNode: state: Optional[torch.Tensor] = None parent: Optional["MCTSNode"] = None children: List["MCTSNode"] = field(default_factory=list) visit_count: int = 0 total_value: float = 0.0 prior: float = 0.0 depth: int = 0 is_expanded: bool = False action: Optional[int] = None hidden_state: Optional[torch.Tensor] = None @property def q_value(self) -> float: if self.visit_count == 0: return 0.0 return self.total_value / self.visit_count @property def is_leaf(self) -> bool: return not self.is_expanded @property def is_root(self) -> bool: """Whether this node is the root.""" return self.parent is None class ValueNetwork(nn.Module): """Evaluate narrative coherence of a state.""" def __init__(self, d_model: int, hidden_dim: int = 256) -> None: super().__init__() self.network = nn.Sequential( nn.Linear(d_model, hidden_dim, bias=False), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim // 2, bias=False), nn.SiLU(), nn.Linear(hidden_dim // 2, 1, bias=False), nn.Tanh(), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.network(x) class PolicyNetwork(nn.Module): """Suggest next arrangement step.""" def __init__(self, d_model: int, num_actions: int = 8) -> None: super().__init__() self.network = nn.Sequential( nn.Linear(d_model, d_model // 2, bias=False), nn.SiLU(), nn.Linear(d_model // 2, num_actions, bias=False), ) self.num_actions = num_actions def forward(self, x: torch.Tensor) -> torch.Tensor: return self.network(x) class MCTSReasoner(nn.Module): """MCTS Reasoning Engine for AAM sentence arrangement.""" def __init__( self, d_model: int, num_actions: int = 8, config: Optional[MCTSConfig] = None, ) -> None: super().__init__() self.d_model = d_model self.num_actions = num_actions self.config = config or MCTSConfig() if self.config.use_value_network: self.value_network = ValueNetwork(d_model) else: self.value_network = None self.policy_network = PolicyNetwork(d_model, num_actions) self.state_encoder = nn.Sequential( nn.Linear(d_model, d_model, bias=False), nn.SiLU(), nn.Linear(d_model, d_model, bias=False), ) def forward( self, x: torch.Tensor, num_simulations: Optional[int] = None, ) -> Tuple[torch.Tensor, Dict[str, Any]]: batch_size = x.shape[0] n_sims = num_simulations or self.config.num_simulations encoded_state = self.state_encoder(x) # Pool over sequence dimension if 3D input if encoded_state.dim() == 3: pooled_state = encoded_state.mean(dim=1) # (batch, d_model) else: pooled_state = encoded_state # (batch, d_model) policy_logits = self.policy_network(pooled_state) # (batch, num_actions) policy_probs = F.softmax(policy_logits / self.config.temperature, dim=-1) # (batch, num_actions) if self.value_network is not None: root_value = self.value_network(pooled_state) # (batch, 1) else: root_value = torch.zeros(batch_size, 1, device=x.device) visit_counts = torch.zeros(batch_size, self.num_actions, device=x.device, dtype=torch.float32) total_values = torch.zeros(batch_size, self.num_actions, device=x.device, dtype=torch.float32) for sim_idx in range(n_sims): ucb_scores = self._compute_ucb(visit_counts, total_values, policy_probs, n_sims) selected_actions = ucb_scores.argmax(dim=-1) if self.value_network is not None: action_onehot = F.one_hot(selected_actions, self.num_actions).float() action_proj = action_onehot @ self.policy_network.network[-1].weight padding = torch.zeros(batch_size, self.d_model - action_proj.shape[-1], device=x.device) action_embedding = torch.cat([action_proj, padding], dim=-1) state_action = pooled_state + action_embedding sim_values = self.value_network(state_action) else: sim_values = torch.rand(batch_size, 1, device=x.device) * 2 - 1 visit_counts.scatter_add_(1, selected_actions.unsqueeze(1), torch.ones(batch_size, 1, device=x.device, dtype=visit_counts.dtype)) total_values.scatter_add_(1, selected_actions.unsqueeze(1), sim_values) if self.config.temperature > 0: action_probs = F.softmax(visit_counts.log() / self.config.temperature, dim=-1) action_probs = torch.where(visit_counts > 0, action_probs, torch.zeros_like(action_probs)) row_sums = action_probs.sum(dim=-1, keepdim=True) action_probs = torch.where( row_sums > 1e-6, action_probs / (row_sums + 1e-8), torch.full_like(action_probs, 1.0 / self.num_actions), ) else: action_probs = F.one_hot(visit_counts.argmax(dim=-1), self.num_actions).float() info = { "total_simulations": n_sims, "root_value": root_value.mean().item(), "max_visit_count": visit_counts.max().item(), "entropy": -(action_probs * (action_probs + 1e-8).log()).sum(-1).mean().item(), "visit_distribution": visit_counts / (visit_counts.sum(-1, keepdim=True) + 1e-8), } return action_probs, info def _compute_ucb(self, visit_counts, total_values, priors, total_simulations): q_values = torch.where( visit_counts > 0, total_values / (visit_counts + 1e-8), torch.zeros_like(total_values), ) parent_visits = visit_counts.sum(dim=-1, keepdim=True) exploration = self.config.c_puct * priors * torch.sqrt(parent_visits + 1) / (1 + visit_counts) return q_values + exploration def compute_thinking_budget(self, complexity_score: float, base_simulations: int = 16, max_simulations: int = 256) -> int: """Compute number of MCTS simulations based on complexity. Adaptive compute budget: more complex inputs get more simulations. Args: complexity_score: Complexity score [0, 1] from ThinkingToggle. base_simulations: Minimum number of simulations. max_simulations: Maximum number of simulations. Returns: Recommended number of simulations. """ return int(base_simulations + (max_simulations - base_simulations) * (complexity_score ** 2)) def get_reasoning_summary(self, info: Dict[str, Any]) -> str: """Summary of reasoning for logging. Args: info: Dictionary from forward output. Returns: Summary string. """ return ( f"MCTS(sims={info['total_simulations']}, " f"root_val={info['root_value']:.3f}, " f"max_visits={info['max_visit_count']:.0f}, " f"entropy={info['entropy']:.3f})" )