Add 3-phase training pipeline, data gen, evaluation

leworld_training.py

@@ -0,0 +1,820 @@
+"""
+LeWorld Training System
+=======================
+3-Phase training procedure:
+  Phase 1: Pre-train components separately
+  Phase 2: End-to-end joint training  
+  Phase 3: Cooperative refinement with info-request loop
+
+Plus: Memory population strategies, data generation, evaluation.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import math
+import random
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+from leworld_architecture import (
+    LeWorldSystem, MemoryConfig, SLMConfig, BLMConfig,
+    ArtificialMemory, SmallLeWorldModel, BigLeWorldModel,
+    count_params
+)
+
+
+# =============================================================================
+# Training Configuration
+# =============================================================================
+
+@dataclass
+class TrainingConfig:
+    """Full training configuration."""
+    # Phase 1: Pre-training
+    phase1_lr: float = 1e-3
+    phase1_epochs: int = 50
+    phase1_batch_size: int = 32
+    
+    # Phase 2: Joint training
+    phase2_lr: float = 3e-4
+    phase2_epochs: int = 100
+    phase2_batch_size: int = 16
+    phase2_warmup_steps: int = 500
+    
+    # Phase 3: Refinement
+    phase3_lr: float = 1e-4
+    phase3_epochs: int = 50
+    phase3_batch_size: int = 16
+    
+    # General
+    weight_decay: float = 0.01
+    grad_clip: float = 1.0
+    state_dim: int = 64
+    char_dim: int = 32
+    sequence_length: int = 20  # timesteps per sequence
+    
+    # Loss weights
+    lambda_balance: float = 0.01      # routing balance
+    lambda_diversity: float = 0.001   # address diversity
+    lambda_entropy: float = 0.01      # routing entropy
+    lambda_info_util: float = 0.1     # info request utility
+    
+    # Temperature annealing
+    temp_anneal_rate: float = 3e-5
+    temp_min: float = 0.1
+
+
+# =============================================================================
+# Synthetic Data Generation
+# =============================================================================
+
+class StateTransitionDataset(Dataset):
+    """
+    Generates synthetic state transition sequences for training.
+    
+    Each sequence has:
+    - States that evolve according to learnable dynamics
+    - Characteristics that stay fixed per sequence
+    - Ground-truth "useful memory" labels (for Phase 1 SLM pre-training)
+    
+    The key insight: we embed patterns into memory, and the state transitions
+    DEPEND on what's in specific memory regions. This creates a genuine need
+    for memory retrieval — the model can't predict next state without reading
+    the right memory.
+    """
+    
+    def __init__(
+        self,
+        num_sequences: int,
+        seq_length: int,
+        state_dim: int,
+        char_dim: int,
+        memory: ArtificialMemory,
+        difficulty: str = "easy",  # easy, medium, hard
+    ):
+        self.num_sequences = num_sequences
+        self.seq_length = seq_length
+        self.state_dim = state_dim
+        self.char_dim = char_dim
+        self.memory = memory
+        
+        # Generate all sequences upfront
+        self.data = self._generate_sequences(difficulty)
+    
+    def _generate_sequences(self, difficulty: str) -> List[Dict]:
+        """Generate synthetic state-transition sequences."""
+        data = []
+        mem_size = self.memory.config.num_words
+        
+        for _ in range(self.num_sequences):
+            # Static characteristics for this sequence
+            characteristics = torch.randn(self.char_dim)
+            
+            # Choose "relevant" memory regions (ground truth for SLM training)
+            if difficulty == "easy":
+                n_relevant = 1  # only one memory region matters
+            elif difficulty == "medium":
+                n_relevant = 2
+            else:
+                n_relevant = 3
+            
+            relevant_addrs = []
+            for _ in range(n_relevant):
+                start = random.randint(0, mem_size - 256)
+                length = random.randint(16, 128)
+                relevant_addrs.append((start, start + length))
+            
+            # Generate state sequence where transitions depend on memory content
+            states = torch.zeros(self.seq_length, self.state_dim)
+            states[0] = torch.randn(self.state_dim)
+            
+            # The transition rule: next_state = f(current_state, memory_content)
+            # We use a simple linear rule seeded by the memory content
+            with torch.no_grad():
+                for addr_start, addr_end in relevant_addrs:
+                    mem_bits = self.memory.memory[addr_start:addr_end].mean(dim=0)
+                    # Memory content influences the transition dynamics
+                    # Pad/tile mem_bits to state_dim
+                    transition_seed_raw = mem_bits * 2 - 1  # map 0,1 → -1,1
+                    transition_seed = transition_seed_raw.repeat(
+                        math.ceil(self.state_dim / len(transition_seed_raw))
+                    )[:self.state_dim]
+                    
+                    # Pad/tile characteristics to state_dim
+                    char_padded = characteristics.repeat(
+                        math.ceil(self.state_dim / len(characteristics))
+                    )[:self.state_dim]
+                    
+                    for t in range(1, self.seq_length):
+                        noise = torch.randn(self.state_dim) * 0.1
+                        # State evolves based on current state + memory influence
+                        states[t] = (
+                            0.8 * states[t-1] 
+                            + 0.15 * transition_seed 
+                            + 0.05 * char_padded
+                            + noise
+                        )
+            
+            data.append({
+                'states': states,                    # (seq_length, state_dim)
+                'characteristics': characteristics,  # (char_dim,)
+                'relevant_addrs': relevant_addrs,    # list of (start, end) tuples
+                'n_relevant': n_relevant,
+            })
+        
+        return data
+    
+    def __len__(self):
+        return self.num_sequences
+    
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        
+        # Pad relevant addresses to fixed length (3 = max n_slms)
+        padded_starts = torch.zeros(3, dtype=torch.long)
+        padded_ends = torch.zeros(3, dtype=torch.long)
+        for i, (s, e) in enumerate(item['relevant_addrs']):
+            padded_starts[i] = s
+            padded_ends[i] = e
+        
+        return {
+            'states': item['states'],
+            'characteristics': item['characteristics'],
+            'relevant_starts': padded_starts,
+            'relevant_ends': padded_ends,
+            'n_relevant': item['n_relevant'],
+        }
+
+
+# =============================================================================
+# Phase 1: Pre-training (Components Separately)
+# =============================================================================
+
+class Phase1Trainer:
+    """
+    Pre-train SLMs and BLM separately.
+    
+    SLMs: Given (past_state, current_state, characteristics), learn to output
+          address ranges that point to "relevant" memory regions.
+          Loss: distance between predicted address range and ground-truth relevant region.
+    
+    BLM: Given perfect memory reads, learn to predict next state.
+         Loss: MSE between predicted and actual next state.
+    """
+    
+    def __init__(self, system: LeWorldSystem, config: TrainingConfig):
+        self.system = system
+        self.config = config
+        
+        # Separate optimizers for SLMs and BLM
+        self.slm_optimizer = optim.AdamW(
+            system.slms.parameters(),
+            lr=config.phase1_lr,
+            weight_decay=config.weight_decay
+        )
+        self.blm_optimizer = optim.AdamW(
+            list(system.blm.parameters()) + list(system.memory.parameters()),
+            lr=config.phase1_lr,
+            weight_decay=config.weight_decay
+        )
+    
+    def train_slms_step(self, batch: Dict) -> Dict:
+        """
+        Train SLMs to find relevant memory regions.
+        
+        Loss: |predicted_addr - target_addr| normalized by address space.
+        """
+        self.slm_optimizer.zero_grad()
+        
+        states = batch['states']  # (B, T, state_dim)
+        chars = batch['characteristics']  # (B, char_dim)
+        target_starts = batch['relevant_starts']  # (B, 3)
+        target_ends = batch['relevant_ends']  # (B, 3)
+        
+        total_loss = None
+        
+        # For each SLM, train to find the corresponding relevant region
+        for i, slm in enumerate(self.system.slms):
+            # Use first two timesteps as past/current
+            past_state = states[:, 0, :]
+            current_state = states[:, 1, :]
+            
+            output = slm(past_state, current_state, chars)
+            
+            # Use logits (differentiable) instead of hard addresses
+            # Target: which high/low byte corresponds to the target address
+            tgt_start = target_starts[:, i].long()
+            
+            half_space = slm.address_head.half_space  # 256
+            tgt_high = tgt_start // half_space  # high byte
+            tgt_low = tgt_start % half_space     # low byte
+            
+            # Cross-entropy over address components (differentiable!)
+            addr_loss = (
+                F.cross_entropy(output['start_logits_high'], tgt_high) +
+                F.cross_entropy(output['start_logits_low'], tgt_low)
+            )
+            
+            # Range length loss
+            tgt_range = (target_ends[:, i] - target_starts[:, i]).clamp(1, self.system.memory.config.max_read_range) - 1
+            range_loss = F.cross_entropy(output['range_logits'], tgt_range.long())
+            
+            slm_loss = addr_loss + 0.5 * range_loss
+            
+            if total_loss is None:
+                total_loss = slm_loss
+            else:
+                total_loss = total_loss + slm_loss
+        
+        total_loss = total_loss / len(self.system.slms)
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.system.slms.parameters(), self.config.grad_clip)
+        self.slm_optimizer.step()
+        
+        return {'slm_loss': total_loss.item()}
+    
+    def train_blm_step(self, batch: Dict) -> Dict:
+        """
+        Train BLM to predict next state given oracle memory reads.
+        
+        Oracle: we read from the KNOWN relevant memory regions (ground truth).
+        """
+        self.blm_optimizer.zero_grad()
+        
+        states = batch['states']
+        chars = batch['characteristics']
+        target_starts = batch['relevant_starts']
+        target_ends = batch['relevant_ends']
+        
+        batch_size = states.shape[0]
+        
+        # Read oracle memory
+        oracle_reads = []
+        slm_fake_outputs = []
+        for i in range(3):
+            _, encoded, _ = self.system.memory.read(
+                target_starts[:, i], target_ends[:, i]
+            )
+            oracle_reads.append(encoded)
+            # Create fake SLM output (just need hidden state)
+            fake_hidden = torch.zeros(batch_size, 128)  # SLM d_model = 128
+            slm_fake_outputs.append({
+                'hidden': fake_hidden,
+                'start_addr': target_starts[:, i],
+                'end_addr': target_ends[:, i],
+                'confidence': torch.ones(batch_size),
+            })
+        
+        # BLM forward with oracle reads
+        total_loss = None
+        for t in range(states.shape[1] - 1):
+            past_state = states[:, max(0, t-1), :]
+            current_state = states[:, t, :]
+            next_state = states[:, t+1, :]
+            
+            blm_out = self.system.blm(
+                past_state, current_state,
+                slm_fake_outputs, oracle_reads
+            )
+            
+            loss = F.mse_loss(blm_out['next_state'], next_state)
+            if total_loss is None:
+                total_loss = loss
+            else:
+                total_loss = total_loss + loss
+        
+        total_loss = total_loss / (states.shape[1] - 1)
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(
+            list(self.system.blm.parameters()) + list(self.system.memory.parameters()),
+            self.config.grad_clip
+        )
+        self.blm_optimizer.step()
+        
+        return {'blm_loss': total_loss.item()}
+
+
+# =============================================================================
+# Phase 2: End-to-End Joint Training
+# =============================================================================
+
+class Phase2Trainer:
+    """
+    Joint training of the entire system end-to-end.
+    
+    The full pipeline runs: SLMs → Memory Read → BLM → Next State
+    
+    Key challenge: gradient flow through discrete decisions
+    - SLM address selection: use soft attention + hard address (ST trick)
+    - BLM routing: use straight-through sigmoid
+    
+    Losses:
+    1. next_state_loss: primary prediction accuracy
+    2. balance_loss: balanced SLM routing  
+    3. diversity_loss: SLMs read different memory regions
+    4. info_utility_loss: BLM's info request improves future predictions
+    """
+    
+    def __init__(self, system: LeWorldSystem, config: TrainingConfig):
+        self.system = system
+        self.config = config
+        
+        # Single optimizer for everything
+        self.optimizer = optim.AdamW(
+            system.parameters(),
+            lr=config.phase2_lr,
+            weight_decay=config.weight_decay
+        )
+        
+        # Learning rate scheduler
+        self.scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            self.optimizer, T_0=config.phase2_epochs // 3, T_mult=2
+        )
+        
+        self.global_step = 0
+    
+    def train_step(self, batch: Dict) -> Dict:
+        """Full end-to-end training step."""
+        self.optimizer.zero_grad()
+        
+        states = batch['states']
+        chars = batch['characteristics']
+        
+        # Multi-step forward
+        output = self.system.multi_step_forward(states, chars)
+        
+        loss = output['total_loss']
+        loss.backward()
+        
+        # Gradient clipping
+        torch.nn.utils.clip_grad_norm_(
+            self.system.parameters(), self.config.grad_clip
+        )
+        
+        self.optimizer.step()
+        
+        # Temperature annealing for router
+        self.global_step += 1
+        self.system.blm.router.anneal_temperature(
+            self.global_step, 
+            self.config.temp_anneal_rate,
+            self.config.temp_min
+        )
+        
+        return {
+            'total_loss': loss.item(),
+            'temperature': self.system.blm.router.temperature.item(),
+            'step': self.global_step,
+        }
+
+
+# =============================================================================
+# Phase 3: Cooperative Refinement with Info-Request Loop
+# =============================================================================
+
+class Phase3Trainer:
+    """
+    Refinement phase: train the info-request mechanism.
+    
+    The BLM learns to generate useful "what info do I need?" queries that
+    improve the SLMs' memory retrieval in the NEXT timestep.
+    
+    Training signal: compare prediction quality WITH vs WITHOUT info-request
+    modulation. If info-request helped → reward; if not → penalize.
+    
+    This is inspired by ProactAgent (arxiv:2604.20572) paired-branch reward.
+    """
+    
+    def __init__(self, system: LeWorldSystem, config: TrainingConfig):
+        self.system = system
+        self.config = config
+        
+        # Optimizer: higher LR for info-request modules, lower for rest
+        info_params = set(id(p) for p in system.blm.info_request.parameters())
+        info_params.update(id(p) for p in system.info_to_slm.parameters())
+        
+        other_blm_params = [p for p in system.blm.parameters() if id(p) not in info_params]
+        
+        self.optimizer = optim.AdamW([
+            {'params': list(system.blm.info_request.parameters()) + list(system.info_to_slm.parameters()), 'lr': config.phase3_lr},
+            {'params': list(system.slms.parameters()), 'lr': config.phase3_lr * 0.1},
+            {'params': other_blm_params, 'lr': config.phase3_lr * 0.1},
+            {'params': list(system.memory.parameters()), 'lr': config.phase3_lr * 0.01},
+        ], weight_decay=config.weight_decay)
+    
+    def train_step(self, batch: Dict) -> Dict:
+        """
+        Paired-branch training:
+        Branch A: Run with info-request modulation (full system)
+        Branch B: Run WITHOUT info-request (baseline)
+        Reward = improvement of A over B
+        """
+        self.optimizer.zero_grad()
+        
+        states = batch['states']
+        chars = batch['characteristics']
+        
+        # Branch A: with info-request loop
+        output_with = self.system.multi_step_forward(states, chars)
+        loss_with = output_with['total_loss']
+        
+        # Branch B: without info-request (set info_query to None at each step)
+        # We do this by running forward without passing info_query between steps
+        batch_size, T, state_dim = states.shape
+        loss_without = None
+        
+        for t in range(T - 1):
+            past_state = states[:, max(0, t-1), :]
+            current_state = states[:, t, :]
+            next_state = states[:, t+1, :]
+            
+            output = self.system(
+                past_state, current_state, chars,
+                next_state, info_query_prev=None  # NO info request
+            )
+            if output['losses']:
+                if loss_without is None:
+                    loss_without = output['losses']['next_state_loss']
+                else:
+                    loss_without = loss_without + output['losses']['next_state_loss']
+        
+        if loss_without is None:
+            loss_without = torch.tensor(0.0)
+        else:
+            loss_without = loss_without / max(1, T - 1)
+        
+        # Info utility: reward if info-request helps, penalize if not
+        improvement = (loss_without - loss_with).detach()  # positive = info helped
+        
+        # Total loss: prediction loss + info utility bonus
+        total_loss = loss_with - self.config.lambda_info_util * improvement
+        
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.system.parameters(), self.config.grad_clip)
+        self.optimizer.step()
+        
+        return {
+            'loss_with_info': loss_with.item(),
+            'loss_without_info': loss_without.item(),
+            'improvement': improvement.item(),
+            'total_loss': total_loss.item(),
+        }
+
+
+# =============================================================================
+# Memory Population Strategies
+# =============================================================================
+
+class MemoryPopulator:
+    """
+    Strategies for populating the artificial memory with meaningful content.
+    
+    In a real application, memory would be populated by experience / observations.
+    Here we provide several strategies for initial content.
+    """
+    
+    @staticmethod
+    def random_bits(memory: ArtificialMemory):
+        """Fill with random bits (baseline)."""
+        memory.memory.uniform_(0, 1).round_()
+    
+    @staticmethod
+    def structured_patterns(memory: ArtificialMemory):
+        """
+        Fill with structured patterns that encode different "knowledge types."
+        
+        Memory layout:
+        - [0x0000 - 0x3FFF]: Dynamics patterns (state transition rules)
+        - [0x4000 - 0x7FFF]: Context patterns (characteristic-dependent info)
+        - [0x8000 - 0xBFFF]: History patterns (temporal sequences)
+        - [0xC000 - 0xFFFF]: Association patterns (cross-references)
+        """
+        N = memory.config.num_words
+        W = memory.config.word_size
+        quarter = N // 4
+        
+        with torch.no_grad():
+            # Region 1: Dynamics — repeating patterns (easy to learn)
+            for i in range(quarter):
+                pattern = torch.zeros(W)
+                pattern[i % W] = 1.0  # cyclic single-bit pattern
+                memory.memory[i] = pattern
+            
+            # Region 2: Context — characteristic-dependent
+            for i in range(quarter, 2 * quarter):
+                seed = i - quarter
+                torch.manual_seed(seed)
+                memory.memory[i] = torch.randint(0, 2, (W,)).float()
+            
+            # Region 3: History — sequential counting in binary
+            for i in range(2 * quarter, 3 * quarter):
+                binary = torch.zeros(W)
+                val = i - 2 * quarter
+                for bit in range(min(W, 16)):
+                    binary[bit] = float((val >> bit) & 1)
+                memory.memory[i] = binary
+            
+            # Region 4: Associations — XOR patterns
+            for i in range(3 * quarter, N):
+                a = memory.memory[i % quarter]  # reference region 1
+                b = memory.memory[quarter + (i % quarter)]  # reference region 2
+                memory.memory[i] = ((a + b) % 2)  # XOR
+    
+    @staticmethod
+    def from_experience(memory: ArtificialMemory, experiences: torch.Tensor):
+        """
+        Populate memory from observed data.
+        
+        Args:
+            experiences: (N, feature_dim) tensor of observed features
+                        Each feature vector gets encoded to bits and stored
+        """
+        with torch.no_grad():
+            N = min(experiences.shape[0], memory.config.num_words)
+            W = memory.config.word_size
+            
+            # Simple quantization: threshold at median
+            for i in range(N):
+                feat = experiences[i]
+                # Truncate/pad to word_size
+                if len(feat) >= W:
+                    bits = (feat[:W] > feat[:W].median()).float()
+                else:
+                    bits = torch.zeros(W)
+                    bits[:len(feat)] = (feat > feat.median()).float()
+                memory.memory[i] = bits
+
+
+# =============================================================================
+# Evaluation
+# =============================================================================
+
+class Evaluator:
+    """Evaluation metrics for the LeWorld system."""
+    
+    @staticmethod
+    def prediction_accuracy(
+        system: LeWorldSystem,
+        dataloader: DataLoader,
+        n_steps: int = 5
+    ) -> Dict:
+        """
+        Evaluate next-state prediction accuracy.
+        
+        Metrics:
+        - MSE: mean squared error of state predictions
+        - MAE: mean absolute error
+        - Multi-step MSE: prediction error at different horizons
+        - Routing diversity: how varied the SLM selections are
+        """
+        system.eval()
+        total_mse = 0.0
+        total_mae = 0.0
+        step_mses = [0.0] * n_steps
+        all_masks = []
+        n_batches = 0
+        
+        with torch.no_grad():
+            for batch in dataloader:
+                states = batch['states']
+                chars = batch['characteristics']
+                
+                output = system.multi_step_forward(states, chars, n_steps)
+                
+                # Ground truth future states
+                gt_future = states[:, 1:n_steps+1, :]
+                pred_future = output['predictions'][:, :n_steps, :]
+                
+                actual_steps = min(n_steps, pred_future.shape[1])
+                
+                mse = F.mse_loss(pred_future[:, :actual_steps], gt_future[:, :actual_steps])
+                mae = F.l1_loss(pred_future[:, :actual_steps], gt_future[:, :actual_steps])
+                
+                total_mse += mse.item()
+                total_mae += mae.item()
+                
+                # Per-step MSE
+                for t in range(actual_steps):
+                    step_mse = F.mse_loss(pred_future[:, t], gt_future[:, t])
+                    step_mses[t] += step_mse.item()
+                
+                # Collect routing masks
+                all_masks.append(output['masks'])
+                n_batches += 1
+        
+        # Routing diversity: entropy of SLM usage
+        all_masks = torch.cat(all_masks, dim=0)  # (total, T, n_slms)
+        usage_per_slm = all_masks.mean(dim=(0, 1))  # (n_slms,)
+        routing_entropy = -(usage_per_slm * torch.log(usage_per_slm + 1e-8)).sum().item()
+        
+        system.train()
+        
+        return {
+            'mse': total_mse / max(1, n_batches),
+            'mae': total_mae / max(1, n_batches),
+            'step_mses': [m / max(1, n_batches) for m in step_mses],
+            'routing_entropy': routing_entropy,
+            'slm_usage': usage_per_slm.tolist(),
+        }
+
+
+# =============================================================================
+# Full Training Pipeline
+# =============================================================================
+
+def run_training(
+    system: LeWorldSystem,
+    train_config: TrainingConfig,
+    num_train_sequences: int = 1000,
+    num_val_sequences: int = 200,
+):
+    """Execute the full 3-phase training pipeline."""
+    
+    print("=" * 70)
+    print("LeWorld Training Pipeline")
+    print("=" * 70)
+    
+    # Populate memory with structured patterns
+    print("\n[Setup] Populating artificial memory...")
+    MemoryPopulator.structured_patterns(system.memory)
+    
+    # Create datasets
+    print("[Setup] Generating training data...")
+    train_dataset = StateTransitionDataset(
+        num_sequences=num_train_sequences,
+        seq_length=train_config.sequence_length,
+        state_dim=train_config.state_dim,
+        char_dim=train_config.char_dim,
+        memory=system.memory,
+        difficulty="medium",
+    )
+    
+    val_dataset = StateTransitionDataset(
+        num_sequences=num_val_sequences,
+        seq_length=train_config.sequence_length,
+        state_dim=train_config.state_dim,
+        char_dim=train_config.char_dim,
+        memory=system.memory,
+        difficulty="medium",
+    )
+    
+    train_loader = DataLoader(train_dataset, batch_size=train_config.phase1_batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=train_config.phase1_batch_size)
+    
+    evaluator = Evaluator()
+    
+    # ===== Phase 1: Pre-training =====
+    print(f"\n{'='*70}")
+    print("Phase 1: Pre-training (SLMs + BLM separately)")
+    print(f"{'='*70}")
+    
+    phase1 = Phase1Trainer(system, train_config)
+    
+    for epoch in range(min(3, train_config.phase1_epochs)):  # shortened for demo
+        epoch_slm_loss = 0
+        epoch_blm_loss = 0
+        n_batches = 0
+        
+        for batch in train_loader:
+            slm_metrics = phase1.train_slms_step(batch)
+            blm_metrics = phase1.train_blm_step(batch)
+            
+            epoch_slm_loss += slm_metrics['slm_loss']
+            epoch_blm_loss += blm_metrics['blm_loss']
+            n_batches += 1
+        
+        print(f"  Epoch {epoch+1}: SLM loss={epoch_slm_loss/n_batches:.4f}, "
+              f"BLM loss={epoch_blm_loss/n_batches:.4f}")
+    
+    # Evaluate after Phase 1
+    val_metrics = evaluator.prediction_accuracy(system, val_loader, n_steps=5)
+    print(f"  Phase 1 eval: MSE={val_metrics['mse']:.4f}, "
+          f"Routing entropy={val_metrics['routing_entropy']:.4f}")
+    
+    # ===== Phase 2: Joint Training =====
+    print(f"\n{'='*70}")
+    print("Phase 2: End-to-End Joint Training")
+    print(f"{'='*70}")
+    
+    phase2 = Phase2Trainer(system, train_config)
+    train_loader2 = DataLoader(train_dataset, batch_size=train_config.phase2_batch_size, shuffle=True)
+    val_loader2 = DataLoader(val_dataset, batch_size=train_config.phase2_batch_size)
+    
+    for epoch in range(min(5, train_config.phase2_epochs)):  # shortened for demo
+        epoch_loss = 0
+        n_batches = 0
+        
+        for batch in train_loader2:
+            metrics = phase2.train_step(batch)
+            epoch_loss += metrics['total_loss']
+            n_batches += 1
+        
+        print(f"  Epoch {epoch+1}: loss={epoch_loss/n_batches:.4f}, "
+              f"temp={metrics['temperature']:.4f}")
+    
+    val_metrics = evaluator.prediction_accuracy(system, val_loader2, n_steps=5)
+    print(f"  Phase 2 eval: MSE={val_metrics['mse']:.4f}, "
+          f"Routing entropy={val_metrics['routing_entropy']:.4f}, "
+          f"SLM usage={[f'{u:.2f}' for u in val_metrics['slm_usage']]}")
+    
+    # ===== Phase 3: Info-Request Refinement =====
+    print(f"\n{'='*70}")
+    print("Phase 3: Info-Request Cooperative Refinement")
+    print(f"{'='*70}")
+    
+    phase3 = Phase3Trainer(system, train_config)
+    
+    for epoch in range(min(3, train_config.phase3_epochs)):  # shortened for demo
+        epoch_loss = 0
+        epoch_improvement = 0
+        n_batches = 0
+        
+        for batch in train_loader2:
+            metrics = phase3.train_step(batch)
+            epoch_loss += metrics['total_loss']
+            epoch_improvement += metrics['improvement']
+            n_batches += 1
+        
+        print(f"  Epoch {epoch+1}: loss={epoch_loss/n_batches:.4f}, "
+              f"info improvement={epoch_improvement/n_batches:.4f}")
+    
+    # Final evaluation
+    print(f"\n{'='*70}")
+    print("Final Evaluation")
+    print(f"{'='*70}")
+    
+    final_metrics = evaluator.prediction_accuracy(system, val_loader2, n_steps=5)
+    print(f"  Final MSE: {final_metrics['mse']:.4f}")
+    print(f"  Final MAE: {final_metrics['mae']:.4f}")
+    print(f"  Per-step MSE: {[f'{m:.4f}' for m in final_metrics['step_mses']]}")
+    print(f"  Routing entropy: {final_metrics['routing_entropy']:.4f}")
+    print(f"  SLM usage: {[f'{u:.2f}' for u in final_metrics['slm_usage']]}")
+    
+    return final_metrics
+
+
+# =============================================================================
+# Entry Point
+# =============================================================================
+
+if __name__ == "__main__":
+    # Build system
+    mem_config = MemoryConfig()
+    slm_config = SLMConfig()
+    blm_config = BLMConfig()
+    train_config = TrainingConfig(sequence_length=10)  # shorter for demo
+    
+    system = LeWorldSystem(mem_config, slm_config, blm_config)
+    count_params(system, "Full LeWorld System")
+    
+    # Run training
+    metrics = run_training(
+        system, train_config,
+        num_train_sequences=100,  # small for demo
+        num_val_sequences=30,
+    )
+    
+    print("\n✅ Training pipeline complete!")