train_decoder.py

"""Training script for quantum syndrome decoder."""
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import json
from pathlib import Path
import argparse

from decoder_model import QuantumSyndromeDecoder, SynergyExtractor, RLPolicyNetwork


class SyndromeDataset(Dataset):
    """Dataset for syndrome-to-error mapping."""
    def __init__(self, syndromes, errors):
        self.syndromes = torch.from_numpy(syndromes).float()
        self.errors = torch.from_numpy(errors).float()
    
    def __len__(self):
        return len(self.syndromes)
    
    def __getitem__(self, idx):
        # Add time dimension (T=1 for single round)
        syndrome = self.syndromes[idx].permute(2, 0, 1).unsqueeze(-1)
        error = self.errors[idx].permute(2, 0, 1).unsqueeze(-1)
        return syndrome, error


def train_supervised(model, train_loader, val_loader, epochs=50, lr=5e-4, 
                     device='cuda', save_path='decoder.pt'):
    """Train decoder with supervised learning."""
    model = model.to(device)
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=3e-2)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
    criterion = nn.BCELoss()
    
    best_val_loss = float('inf')
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0
        train_correct = 0
        train_total = 0
        
        for syndromes, errors in train_loader:
            syndromes = syndromes.to(device)
            errors = errors.to(device)
            
            optimizer.zero_grad()
            predictions = model(syndromes)
            
            # Compute loss for all 4 output channels
            loss = criterion(predictions, errors.repeat(1, 2, 1, 1, 1))
            
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            
            # Accuracy for binary classification
            pred_binary = (predictions > 0.5).float()
            train_correct += (pred_binary == errors.repeat(1, 2, 1, 1, 1)).sum().item()
            train_total += errors.numel() * 2
        
        # Validation
        model.eval()
        val_loss = 0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for syndromes, errors in val_loader:
                syndromes = syndromes.to(device)
                errors = errors.to(device)
                
                predictions = model(syndromes)
                loss = criterion(predictions, errors.repeat(1, 2, 1, 1, 1))
                
                val_loss += loss.item()
                pred_binary = (predictions > 0.5).float()
                val_correct += (pred_binary == errors.repeat(1, 2, 1, 1, 1)).sum().item()
                val_total += errors.numel() * 2
        
        train_loss /= len(train_loader)
        val_loss /= len(val_loader)
        train_acc = train_correct / train_total
        val_acc = val_correct / val_total
        
        print(f"Epoch {epoch+1}/{epochs}: "
              f"train_loss={train_loss:.4f}, train_acc={train_acc:.4f}, "
              f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}")
        
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_loss': val_loss,
            }, save_path)
            print(f"  -> Saved best model (val_loss={val_loss:.4f})")
        
        scheduler.step()
    
    return model


def train_rl_policy(policy_net, env, episodes=1000, lr=3e-4, gamma=0.99,
                    device='cuda', save_path='policy.pt'):
    """Train policy network with REINFORCE."""
    policy_net = policy_net.to(device)
    optimizer = optim.Adam(policy_net.parameters(), lr=lr)
    
    for episode in range(episodes):
        # Generate episode
        state = env.reset()
        log_probs = []
        values = []
        rewards = []
        
        done = False
        while not done:
            state_tensor = torch.from_numpy(state).float().unsqueeze(0).to(device)
            policy_logits, value = policy_net(state_tensor)
            
            # Sample action
            probs = torch.softmax(policy_logits.view(-1, 3), dim=-1)
            dist = torch.distributions.Categorical(probs)
            action = dist.sample()
            log_prob = dist.log_prob(action)
            
            # Execute action
            next_state, reward, done, _ = env.step(action.cpu().numpy())
            
            log_probs.append(log_prob)
            values.append(value)
            rewards.append(reward)
            state = next_state
        
        # Compute returns
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        
        returns = torch.tensor(returns, device=device)
        values = torch.cat(values)
        log_probs = torch.cat(log_probs)
        
        # Policy gradient loss
        advantage = returns - values.squeeze()
        policy_loss = -(log_probs * advantage.detach()).mean()
        value_loss = advantage.pow(2).mean()
        
        loss = policy_loss + 0.5 * value_loss
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if episode % 100 == 0:
            print(f"Episode {episode}: reward={sum(rewards):.2f}, "
                  f"policy_loss={policy_loss.item():.4f}")
    
    torch.save(policy_net.state_dict(), save_path)
    return policy_net


def evaluate_decoder(model, test_loader, device='cuda'):
    """Evaluate decoder performance."""
    model.eval()
    model = model.to(device)
    
    total_errors = 0
    corrected_errors = 0
    false_positives = 0
    
    with torch.no_grad():
        for syndromes, errors in test_loader:
            syndromes = syndromes.to(device)
            errors = errors.to(device)
            
            predictions = model(syndromes)
            pred_errors = (predictions > 0.5).float()
            
            true_errors = errors.repeat(1, 2, 1, 1, 1)
            
            total_errors += true_errors.sum().item()
            corrected_errors += (pred_errors * true_errors).sum().item()
            false_positives += (pred_errors * (1 - true_errors)).sum().item()
    
    precision = corrected_errors / (corrected_errors + false_positives) if (corrected_errors + false_positives) > 0 else 0
    recall = corrected_errors / total_errors if total_errors > 0 else 0
    
    print(f"Evaluation Results:")
    print(f"  Total errors: {total_errors}")
    print(f"  Corrected: {corrected_errors}")
    print(f"  False positives: {false_positives}")
    print(f"  Precision: {precision:.4f}")
    print(f"  Recall: {recall:.4f}")
    
    return {
        'precision': precision,
        'recall': recall,
        'total_errors': total_errors,
        'corrected': corrected_errors,
        'false_positives': false_positives,
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', choices=['supervised', 'rl'], default='supervised')
    parser.add_argument('--distance', type=int, default=5)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--lr', type=float, default=5e-4)
    parser.add_argument('--channels', type=int, default=128)
    parser.add_argument('--layers', type=int, default=4)
    parser.add_argument('--data_dir', default='data')
    parser.add_argument('--output_dir', default='outputs')
    args = parser.parse_args()
    
    Path(args.output_dir).mkdir(exist_ok=True)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    # Load data
    syndromes = np.load(f"{args.data_dir}/syndromes.npy")
    errors = np.load(f"{args.data_dir}/errors.npy")
    
    # Split
    n = len(syndromes)
    n_train = int(0.8 * n)
    n_val = int(0.1 * n)
    
    train_dataset = SyndromeDataset(syndromes[:n_train], errors[:n_train])
    val_dataset = SyndromeDataset(syndromes[n_train:n_train+n_val], errors[n_train:n_train+n_val])
    test_dataset = SyndromeDataset(syndromes[n_train+n_val:], errors[n_train+n_val:])
    
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size)
    
    if args.mode == 'supervised':
        model = QuantumSyndromeDecoder(
            distance=args.distance,
            channels=args.channels,
            num_layers=args.layers,
        )
        
        print(f"Training supervised decoder: {sum(p.numel() for p in model.parameters())} parameters")
        model = train_supervised(
            model, train_loader, val_loader,
            epochs=args.epochs, lr=args.lr, device=device,
            save_path=f"{args.output_dir}/decoder_best.pt"
        )
        
        # Final evaluation
        model.load_state_dict(torch.load(f"{args.output_dir}/decoder_best.pt")['model_state_dict'])
        results = evaluate_decoder(model, test_loader, device)
        
        with open(f"{args.output_dir}/results.json", "w") as f:
            json.dump(results, f, indent=2)
    
    else:
        print("RL mode requires environment implementation")