hrm.py

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

class InputModule(nn.Module):
    """
    Transforms any input into internal representation for the HRM
    """
    def __init__(self, input_dim: int, hidden_dim: int):
        super().__init__()
        self.projection = nn.Linear(input_dim, hidden_dim)
        self.layer_norm = nn.LayerNorm(hidden_dim)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.projection(x)
        return self.layer_norm(x)

class OutputModule(nn.Module):
    """
    Transforms internal representation to final output
    """
    def __init__(self, hidden_dim: int, output_dim: int):
        super().__init__()
        self.projection = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.projection(x)

class LowLevelModule(nn.Module):
    """
    The 'engineer' module - handles detailed execution and computation
    Operates continuously at higher frequency
    """
    def __init__(self, hidden_dim: int):
        super().__init__()
        self.hidden_dim = hidden_dim
        
        # Core processing layers
        self.update_layer = nn.Linear(hidden_dim * 2, hidden_dim)  # state + high_level_input
        self.state_layer = nn.Linear(hidden_dim, hidden_dim)
        self.layer_norm = nn.LayerNorm(hidden_dim)
        self.activation = nn.ReLU()
        
    def forward(self, state: torch.Tensor, high_level_signal: torch.Tensor) -> torch.Tensor:
        """
        Execute one timestep of low-level processing
        
        Args:
            state: Current state of the problem/construction
            high_level_signal: Strategy signal from high-level module
        
        Returns:
            Updated state after one processing step
        """
        # Combine current state with high-level guidance
        combined = torch.cat([state, high_level_signal], dim=-1)
        
        # Process the combined information
        update = self.update_layer(combined)
        update = self.activation(update)
        
        # Update state (residual connection)
        new_state = state + self.state_layer(update)
        return self.layer_norm(new_state)

class HighLevelModule(nn.Module):
    """
    The 'architect' module - handles strategy and planning
    Operates periodically to review and update strategy
    """
    def __init__(self, hidden_dim: int):
        super().__init__()
        self.hidden_dim = hidden_dim
        
        # Strategy processing layers
        self.strategy_layer = nn.Linear(hidden_dim * 2, hidden_dim)  # prev_state + current_state
        self.planning_layer = nn.Linear(hidden_dim, hidden_dim)
        self.layer_norm = nn.LayerNorm(hidden_dim)
        self.activation = nn.Tanh()  # Tanh for more controlled strategy updates
        
    def forward(self, prev_state: torch.Tensor, current_state: torch.Tensor) -> torch.Tensor:
        """
        Generate strategy update based on progress analysis
        
        Args:
            prev_state: State before the last cycle
            current_state: Current state after low-level processing
            
        Returns:
            Updated strategy signal for low-level module
        """
        # Analyze progress by comparing states
        progress_analysis = torch.cat([prev_state, current_state], dim=-1)
        
        # Generate new strategy
        strategy = self.strategy_layer(progress_analysis)
        strategy = self.activation(strategy)
        strategy = self.planning_layer(strategy)
        
        return self.layer_norm(strategy)

class AdaptiveComputeModule(nn.Module):
    """
    Q-head module that decides when to stop reasoning
    Implements adaptive compute time mechanism
    """
    def __init__(self, hidden_dim: int):
        super().__init__()
        self.halt_layer = nn.Linear(hidden_dim, 1)
        self.confidence_layer = nn.Linear(hidden_dim, 1)
        
    def forward(self, high_level_state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Decide whether to halt or continue reasoning
        
        Args:
            high_level_state: Current state from high-level module
            
        Returns:
            halt_prob: Probability of stopping (0-1)
            confidence: Confidence in current solution
        """
        halt_logit = self.halt_layer(high_level_state)
        halt_prob = torch.sigmoid(halt_logit)
        
        confidence = torch.sigmoid(self.confidence_layer(high_level_state))
        
        return halt_prob, confidence

class HierarchicalReasoningModel(nn.Module):
    """
    Main HRM implementation with all four modules
    """
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        output_dim: int,
        max_cycles: int = 10,
        timesteps_per_cycle: int = 10
    ):
        super().__init__()
        
        self.hidden_dim = hidden_dim
        self.max_cycles = max_cycles
        self.timesteps_per_cycle = timesteps_per_cycle
        
        # Four main modules
        self.input_module = InputModule(input_dim, hidden_dim)
        self.low_level_module = LowLevelModule(hidden_dim)
        self.high_level_module = HighLevelModule(hidden_dim)
        self.output_module = OutputModule(hidden_dim, output_dim)
        self.adaptive_compute = AdaptiveComputeModule(hidden_dim)
        
        # Initialize strategy
        self.initial_strategy = nn.Parameter(torch.randn(1, hidden_dim) * 0.1)
        
    def forward(
        self, 
        x: torch.Tensor, 
        target: Optional[torch.Tensor] = None,
        return_reasoning_trace: bool = False
    ) -> Dict[str, Any]:
        """
        Forward pass with hierarchical reasoning
        
        Args:
            x: Input tensor
            target: Target for training (optional)
            return_reasoning_trace: Whether to return intermediate states
            
        Returns:
            Dictionary containing output, halt decisions, and optionally reasoning trace
        """
        batch_size = x.size(0)
        
        # Transform input to internal representation
        state = self.input_module(x)
        
        # Initialize strategy (broadcast to batch size)
        strategy = self.initial_strategy.expand(batch_size, -1)
        
        # Storage for adaptive compute training
        halt_probs = []
        confidences = []
        states_trace = [] if return_reasoning_trace else None
        
        prev_state = state.clone()
        
        for cycle in range(self.max_cycles):
            # Store state before cycle for progress analysis
            cycle_start_state = state.clone()
            
            # Low-level processing (engineer working)
            for timestep in range(self.timesteps_per_cycle):
                state = self.low_level_module(state, strategy)
                
                if return_reasoning_trace:
                    states_trace.append(state.clone())
            
            # High-level review and strategy update (architect reviewing)
            strategy = self.high_level_module(prev_state, state)
            
            # Adaptive compute decision
            halt_prob, confidence = self.adaptive_compute(strategy)
            halt_probs.append(halt_prob)
            confidences.append(confidence)
            
            # During inference, actually halt if probability is high
            if not self.training and halt_prob.mean() > 0.5:
                break
                
            prev_state = cycle_start_state
        
        # Generate final output
        output = self.output_module(state)
        
        results = {
            'output': output,
            'halt_probs': torch.stack(halt_probs),
            'confidences': torch.stack(confidences),
            'final_state': state,
            'num_cycles_used': cycle + 1
        }
        
        if return_reasoning_trace:
            results['reasoning_trace'] = states_trace
            
        return results

class HRMLoss(nn.Module):
    """
    Combined loss function for HRM training
    Optimizes both task accuracy and adaptive compute efficiency
    """
    def __init__(self, task_weight: float = 1.0, efficiency_weight: float = 0.1):
        super().__init__()
        self.task_weight = task_weight
        self.efficiency_weight = efficiency_weight
        self.task_loss_fn = nn.MSELoss()  # or CrossEntropyLoss for classification
        
    def forward(
        self, 
        predictions: torch.Tensor,
        targets: torch.Tensor,
        halt_probs: torch.Tensor,
        optimal_halt_step: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        """
        Compute combined loss
        
        Args:
            predictions: Model predictions
            targets: Ground truth targets
            halt_probs: Halt probabilities for each cycle
            optimal_halt_step: When the model should have halted (for training)
        """
        # Task accuracy loss
        task_loss = self.task_loss_fn(predictions, targets)
        
        # Efficiency loss (when to halt)
        if optimal_halt_step is not None:
            # Create target halt signal
            num_cycles = halt_probs.size(0)
            batch_size = halt_probs.size(1)
            
            halt_targets = torch.zeros_like(halt_probs)
            for b in range(batch_size):
                if optimal_halt_step[b] < num_cycles:
                    halt_targets[optimal_halt_step[b], b] = 1.0
            
            efficiency_loss = F.binary_cross_entropy(halt_probs, halt_targets)
        else:
            # Encourage efficiency by penalizing late halting
            efficiency_loss = halt_probs.mean()
        
        total_loss = (self.task_weight * task_loss + 
                     self.efficiency_weight * efficiency_loss)
        
        return {
            'total_loss': total_loss,
            'task_loss': task_loss,
            'efficiency_loss': efficiency_loss
        }

# Example usage and demonstration
def create_maze_solving_hrm():
    """
    Create an HRM configured for maze solving
    """
    # Maze representation: flattened 2D grid
    maze_size = 10
    input_dim = maze_size * maze_size  # Flattened maze
    hidden_dim = 128
    output_dim = maze_size * maze_size  # Solution path
    
    model = HierarchicalReasoningModel(
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        output_dim=output_dim,
        max_cycles=15,
        timesteps_per_cycle=8
    )
    
    return model

def create_arithmetic_hrm():
    """
    Create an HRM configured for arithmetic reasoning
    """
    # For problems like complex multiplication
    input_dim = 64  # Encoded arithmetic problem
    hidden_dim = 256
    output_dim = 32  # Encoded result
    
    model = HierarchicalReasoningModel(
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        output_dim=output_dim,
        max_cycles=20,
        timesteps_per_cycle=5
    )
    
    return model

def demonstrate_hrm():
    """
    Demonstrate HRM usage
    """
    print("Creating Hierarchical Reasoning Model...")
    
    # Create model for arithmetic tasks
    model = create_arithmetic_hrm()
    model.eval()
    
    # Example input (batch_size=2, input_dim=64)
    x = torch.randn(2, 64)
    
    print("Running forward pass...")
    with torch.no_grad():
        results = model(x, return_reasoning_trace=True)
    
    print(f"Output shape: {results['output'].shape}")
    print(f"Number of cycles used: {results['num_cycles_used']}")
    print(f"Halt probabilities shape: {results['halt_probs'].shape}")
    print(f"Final halt probabilities: {results['halt_probs'][-1].squeeze()}")
    
    # Show reasoning progression
    trace = results['reasoning_trace']
    print(f"Reasoning trace length: {len(trace)} timesteps")
    
    return model, results

if __name__ == "__main__":
    # Demonstrate the model
    model, results = demonstrate_hrm()
    
    # Show parameter count
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\nTotal parameters: {total_params:,}")
    print("This is much smaller than O3's 200+ billion parameters!")
    
    # Training example
    print("\nTraining example:")
    model.train()
    
    # Dummy training data
    batch_size = 4
    x = torch.randn(batch_size, 64)
    targets = torch.randn(batch_size, 32)
    optimal_halt = torch.randint(5, 15, (batch_size,))  # Random optimal halt steps
    
    # Forward pass
    results = model(x)
    
    # Compute loss
    loss_fn = HRMLoss()
    losses = loss_fn(
        results['output'], 
        targets, 
        results['halt_probs'],
        optimal_halt
    )
    
    print(f"Total loss: {losses['total_loss'].item():.4f}")
    print(f"Task loss: {losses['task_loss'].item():.4f}")
    print(f"Efficiency loss: {losses['efficiency_loss'].item():.4f}")