train_markov_fleet.py

"""
Markov Fleet DQN Training Script (v0.6)

Training 100-bridge fleet with:
- Unified Markov transition matrices (municipality-level)
- Urban 20 bridges (higher importance)
- Rural 80 bridges (standard importance)
- Vectorized parallel training with AsyncVectorEnv
- GPU acceleration with Mixed Precision Training (AMP)

Based on: Phase 3 Vectorized DQN + dqn_bridge_maintenance.py
"""

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from pathlib import Path
import argparse
from collections import deque
from tqdm import tqdm
import time
from torch.amp import autocast, GradScaler
import gymnasium as gym
from gymnasium.vector import AsyncVectorEnv
import yaml
import sys

sys.path.insert(0, str(Path(__file__).parent / 'src'))

from markov_fleet_environment import MarkovFleetEnvironment


# ----- DQN Network for Fleet Actions -----

class FleetDQN(nn.Module):
    """
    Dueling DQN for 100-bridge fleet maintenance.
    
    Architecture:
        - Shared network: Input -> [512, 256]
        - Value stream: [256] -> [128] -> [1]
        - Advantage stream: [256] -> [128] -> [100 * 6] (per-bridge actions)
    
    Input: 100-dim state vector (discrete states 0-2)
    Output: 600-dim Q-values (100 bridges × 6 actions)
    """
    
    def __init__(self, n_bridges: int = 100, n_actions: int = 6):
        super().__init__()
        self.n_bridges = n_bridges
        self.n_actions = n_actions
        self.total_actions = n_bridges * n_actions
        
        # Shared feature extractor
        self.shared = nn.Sequential(
            nn.Linear(n_bridges, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
        )
        
        # Value stream
        self.value_stream = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )
        
        # Advantage stream
        self.advantage_stream = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, self.total_actions)
        )
    
    def forward(self, x):
        """
        Forward pass
        
        Args:
            x: State tensor [batch_size, n_bridges]
        
        Returns:
            Q-values [batch_size, n_bridges, n_actions]
        """
        # Shared features
        features = self.shared(x)
        
        # Value and advantage
        value = self.value_stream(features)  # [B, 1]
        advantage = self.advantage_stream(features)  # [B, n_bridges * n_actions]
        
        # Reshape advantage
        advantage = advantage.view(-1, self.n_bridges, self.n_actions)  # [B, n_bridges, n_actions]
        
        # Dueling architecture: Q = V + (A - mean(A))
        # value: [B, 1] -> [B, 1, 1] for broadcasting
        value = value.unsqueeze(-1)  # [B, 1, 1]
        q_values = value + (advantage - advantage.mean(dim=2, keepdim=True))
        
        return q_values


# ----- Prioritized N-Step Replay Buffer -----

class PrioritizedNStepBuffer:
    """Prioritized N-step replay buffer for fleet transitions"""
    
    def __init__(self, capacity: int, n_steps: int = 3, gamma: float = 0.99,
                 alpha: float = 0.6, beta: float = 0.4, beta_increment: float = 0.001):
        self.capacity = capacity
        self.n_steps = n_steps
        self.gamma = gamma
        self.alpha = alpha
        self.beta = beta
        self.beta_increment = beta_increment
        self.n_step_buffer = deque(maxlen=n_steps)
        
        self.buffer = []
        self.priorities = np.zeros(capacity, dtype=np.float32)
        self.position = 0
        self.size = 0
    
    def push(self, state, actions, reward, next_state, done):
        """Add experience to n-step buffer"""
        self.n_step_buffer.append((state, actions, reward, next_state, done))
        
        if len(self.n_step_buffer) == self.n_steps or done:
            # Calculate n-step return
            n_step_reward = 0.0
            for i, (_, _, r, _, _) in enumerate(self.n_step_buffer):
                n_step_reward += (self.gamma ** i) * r
            
            # Get first and last transitions
            s0, a0 = self.n_step_buffer[0][:2]
            sn, done_n = self.n_step_buffer[-1][3:5]
            
            experience = (s0, a0, n_step_reward, sn, done_n)
            
            # Assign max priority for new experience
            max_priority = self.priorities[:self.size].max() if self.size > 0 else 1.0
            
            if len(self.buffer) < self.capacity:
                self.buffer.append(experience)
            else:
                self.buffer[self.position] = experience
            
            self.priorities[self.position] = max_priority
            self.position = (self.position + 1) % self.capacity
            self.size = len(self.buffer)
            
            if done:
                self.n_step_buffer.clear()
    
    def sample(self, batch_size: int):
        """Sample batch with prioritized sampling"""
        priorities = self.priorities[:self.size]
        probs = priorities ** self.alpha
        probs_sum = probs.sum()
        
        if probs_sum == 0 or np.isnan(probs_sum) or np.isinf(probs_sum):
            probs = np.ones(self.size) / self.size
        else:
            probs /= probs_sum
            probs = np.clip(probs, 1e-8, 1.0)
            probs /= probs.sum()
        
        indices = np.random.choice(self.size, batch_size, p=probs, replace=False)
        
        # Importance sampling weights
        weights = (self.size * probs[indices]) ** (-self.beta)
        weights /= weights.max()
        
        self.beta = min(1.0, self.beta + self.beta_increment)
        
        batch = [self.buffer[i] for i in indices]
        
        states = np.array([b[0] for b in batch], dtype=np.float32)
        actions = np.array([b[1] for b in batch], dtype=np.int64)
        rewards = np.array([b[2] for b in batch], dtype=np.float32)
        next_states = np.array([b[3] for b in batch], dtype=np.float32)
        dones = np.array([b[4] for b in batch], dtype=np.float32)
        
        return states, actions, rewards, next_states, dones, indices, weights
    
    def update_priorities(self, indices, td_errors):
        """Update priorities based on TD-errors"""
        for idx, td_error in zip(indices, td_errors):
            self.priorities[idx] = abs(td_error) + 1e-6
    
    def __len__(self):
        return len(self.buffer)


# ----- Vectorized Environment Wrapper -----

def make_env(n_urban: int, n_rural: int, horizon_years: int, cost_lambda: float, seed: int):
    """Create a single environment instance"""
    def _init():
        env = MarkovFleetEnvironment(
            n_urban=n_urban,
            n_rural=n_rural,
            horizon_years=horizon_years,
            cost_lambda=cost_lambda,
            use_onehot=False,
            seed=seed
        )
        return env
    return _init


# ----- Training Function -----

def train_markov_fleet(
    n_episodes: int = 1000,
    n_envs: int = 4,
    n_urban: int = 20,
    n_rural: int = 80,
    horizon_years: int = 30,
    cost_lambda: float = 1e-3,
    gamma: float = 0.95,
    lr: float = 1.5e-3,
    eps_start: float = 1.0,
    eps_end: float = 0.1,
    eps_decay_episodes: int = 2000,
    buffer_capacity: int = 10000,
    batch_size: int = 64,
    target_sync_steps: int = 500,
    device: str = 'cuda',
    seed: int = 42,
    output_dir: str = 'outputs_markov',
    verbose: bool = True,
):
    """
    Train DQN agent for Markov fleet maintenance.
    
    Args:
        n_episodes: Number of training episodes
        n_envs: Number of parallel environments
        n_urban: Number of urban bridges
        n_rural: Number of rural bridges
        horizon_years: Episode length (years)
        cost_lambda: Cost penalty scaling
        gamma: Discount factor
        lr: Learning rate
        eps_start: Initial exploration rate
        eps_end: Final exploration rate
        eps_decay_episodes: Epsilon decay schedule
        buffer_capacity: Replay buffer size
        batch_size: Batch size for training
        target_sync_steps: Target network update frequency
        device: 'cuda' or 'cpu'
        seed: Random seed
        output_dir: Output directory
        verbose: Print progress
    
    Returns:
        agent: Trained DQN agent
        rewards_history: Episode rewards
        costs_history: Episode costs
    """
    # Setup
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    (output_path / "models").mkdir(exist_ok=True)
    
    n_bridges = n_urban + n_rural
    
    if verbose:
        print("\n" + "="*80)
        print("MARKOV FLEET DQN TRAINING (v0.6)")
        print("="*80)
        print(f"Fleet Configuration:")
        print(f"  Total Bridges: {n_bridges} (Urban: {n_urban}, Rural: {n_rural})")
        print(f"  Episode Horizon: {horizon_years} years")
        print(f"  Cost Lambda: {cost_lambda}")
        print(f"\nTraining Configuration:")
        print(f"  Episodes: {n_episodes}")
        print(f"  Parallel Envs: {n_envs}")
        print(f"  Device: {device}")
        print(f"  Gamma: {gamma}, LR: {lr}")
        print(f"  Buffer: {buffer_capacity}, Batch: {batch_size}")
        print(f"  Target Sync: {target_sync_steps} steps")
        print(f"\nOptimizations:")
        print(f"  ✓ Mixed Precision Training (AMP)")
        print(f"  ✓ Double DQN")
        print(f"  ✓ Dueling DQN")
        print(f"  ✓ N-step Learning (n=3)")
        print(f"  ✓ Prioritized Experience Replay (PER)")
        print(f"  ✓ AsyncVectorEnv ({n_envs}x speedup)")
        print("="*80 + "\n")
    
    # Create vectorized environments
    env_fns = [make_env(n_urban, n_rural, horizon_years, cost_lambda, seed + i) 
               for i in range(n_envs)]
    envs = AsyncVectorEnv(env_fns)
    
    # Initialize networks
    agent = FleetDQN(n_bridges=n_bridges, n_actions=6).to(device)
    target_net = FleetDQN(n_bridges=n_bridges, n_actions=6).to(device)
    target_net.load_state_dict(agent.state_dict())
    target_net.eval()
    
    optimizer = optim.AdamW(agent.parameters(), lr=lr, weight_decay=1e-5)
    scaler = GradScaler('cuda') if device == 'cuda' else None
    buffer = PrioritizedNStepBuffer(
        buffer_capacity, n_steps=3, gamma=gamma,
        alpha=0.6, beta=0.4, beta_increment=0.001
    )
    
    # Training tracking
    rewards_history = []
    costs_history = []
    losses_history = []
    total_steps = 0
    episodes_completed = 0
    start_time = time.time()
    
    # Reset environments
    observations, infos = envs.reset()
    states = observations.astype(np.float32)
    
    # Episode tracking per environment
    env_episode_rewards = np.zeros(n_envs)
    env_episode_costs = np.zeros(n_envs)
    
    pbar = tqdm(total=n_episodes, desc="Training Markov Fleet DQN") if verbose else None
    
    while episodes_completed < n_episodes:
        # Epsilon decay
        eps = eps_end + (eps_start - eps_end) * max(0, 1 - episodes_completed / eps_decay_episodes)
        
        # Select actions for all environments
        actions_batch = []
        
        for i in range(n_envs):
            if np.random.random() < eps:
                # Random actions
                actions = np.random.randint(0, 6, size=n_bridges)
            else:
                # Greedy actions
                with torch.no_grad():
                    state_t = torch.FloatTensor(states[i]).unsqueeze(0).to(device)
                    q_values = agent(state_t)[0]  # [n_bridges, n_actions]
                    actions = q_values.argmax(dim=1).cpu().numpy()
            
            actions_batch.append(actions)
        
        actions_batch = np.array(actions_batch)
        
        # Step all environments
        next_observations, rewards, terminateds, truncateds, infos = envs.step(actions_batch)
        next_states = next_observations.astype(np.float32)
        
        # Calculate costs directly from actions (since AsyncVectorEnv doesn't return step info reliably)
        from src.markov_fleet_environment import ACTION_COST_KUSD
        
        # Store transitions and accumulate costs
        for i in range(n_envs):
            buffer.push(
                states[i], actions_batch[i], rewards[i],
                next_states[i], terminateds[i] or truncateds[i]
            )
            
            env_episode_rewards[i] += rewards[i]
            
            # Calculate cost directly from actions taken
            step_cost = np.sum(ACTION_COST_KUSD[actions_batch[i]])
            env_episode_costs[i] += step_cost
            
            # Episode completion
            if terminateds[i] or truncateds[i]:
                rewards_history.append(env_episode_rewards[i])
                costs_history.append(env_episode_costs[i])
                
                env_episode_rewards[i] = 0
                env_episode_costs[i] = 0
                
                episodes_completed += 1
                if pbar:
                    pbar.update(1)
                
                if episodes_completed >= n_episodes:
                    break
        
        states = next_states
        total_steps += n_envs
        
        # Optimization step
        if len(buffer) >= batch_size:
            s_b, a_b, r_b, sn_b, d_b, indices, weights = buffer.sample(batch_size)
            
            # Convert to tensors
            s_b_t = torch.FloatTensor(s_b).to(device)  # [B, n_bridges]
            a_b_t = torch.LongTensor(a_b).to(device)  # [B, n_bridges]
            r_b_t = torch.FloatTensor(r_b).to(device)  # [B]
            sn_b_t = torch.FloatTensor(sn_b).to(device)  # [B, n_bridges]
            d_b_t = torch.FloatTensor(d_b).to(device)  # [B]
            w_b_t = torch.FloatTensor(weights).to(device)  # [B]
            
            # Mixed precision training
            if scaler:
                with autocast('cuda'):
                    # Q(s, a) for taken actions
                    q_values = agent(s_b_t)  # [B, n_bridges, n_actions]
                    
                    # Ensure a_b_t has correct dimensions for gather
                    # q_values: [B, n_bridges, n_actions], need index: [B, n_bridges, 1]
                    if a_b_t.dim() == 2:  # [B, n_bridges]
                        a_b_t_expanded = a_b_t.unsqueeze(2)  # [B, n_bridges, 1]
                    else:
                        a_b_t_expanded = a_b_t
                    
                    q_values_taken = q_values.gather(2, a_b_t_expanded).squeeze(2)  # [B, n_bridges]
                    q_values_taken = q_values_taken.mean(dim=1)  # [B]
                    
                    # Double DQN: a* = argmax_a Q(s', a), target = r + γ * Q_target(s', a*)
                    with torch.no_grad():
                        next_q_values = agent(sn_b_t)  # [B, n_bridges, n_actions]
                        next_actions = next_q_values.argmax(dim=2)  # [B, n_bridges]
                        
                        target_next_q_values = target_net(sn_b_t)  # [B, n_bridges, n_actions]
                        next_actions_expanded = next_actions.unsqueeze(2)  # [B, n_bridges, 1]
                        target_next_q_values = target_next_q_values.gather(
                            2, next_actions_expanded
                        ).squeeze(2)  # [B, n_bridges]
                        target_next_q_values = target_next_q_values.mean(dim=1)  # [B]
                        
                        targets = r_b_t + gamma * target_next_q_values * (1 - d_b_t)
                    
                    # TD-error with importance sampling weights
                    td_errors = (q_values_taken - targets).abs()
                    loss = (w_b_t * (q_values_taken - targets).pow(2)).mean()
                
                optimizer.zero_grad()
                scaler.scale(loss).backward()
                torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm=1.0)
                scaler.step(optimizer)
                scaler.update()
            else:
                # Standard training
                q_values = agent(s_b_t)  # [B, n_bridges, n_actions]
                a_b_t_expanded = a_b_t.unsqueeze(2)  # [B, n_bridges, 1]
                q_values_taken = q_values.gather(2, a_b_t_expanded).squeeze(2).mean(dim=1)  # [B]
                
                with torch.no_grad():
                    next_q_values = agent(sn_b_t)  # [B, n_bridges, n_actions]
                    next_actions = next_q_values.argmax(dim=2)  # [B, n_bridges]
                    target_next_q_values = target_net(sn_b_t)  # [B, n_bridges, n_actions]
                    next_actions_expanded = next_actions.unsqueeze(2)  # [B, n_bridges, 1]
                    target_next_q_values = target_next_q_values.gather(
                        2, next_actions_expanded
                    ).squeeze(2).mean(dim=1)  # [B]
                    targets = r_b_t + gamma * target_next_q_values * (1 - d_b_t)
                
                td_errors = (q_values_taken - targets).abs()
                loss = (w_b_t * (q_values_taken - targets).pow(2)).mean()
                
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm=1.0)
                optimizer.step()
            
            # Update priorities
            buffer.update_priorities(indices, td_errors.detach().cpu().numpy())
            losses_history.append(loss.item())
        
        # Sync target network
        if total_steps % target_sync_steps == 0:
            target_net.load_state_dict(agent.state_dict())
    
    if pbar:
        pbar.close()
    
    envs.close()
    
    # Training summary
    elapsed_time = time.time() - start_time
    
    if verbose:
        print("\n" + "="*80)
        print("TRAINING COMPLETE")
        print("="*80)
        print(f"Total Episodes: {episodes_completed}")
        print(f"Total Time: {elapsed_time:.2f} sec ({elapsed_time/60:.2f} min)")
        print(f"Time per Episode: {elapsed_time/episodes_completed:.3f} sec")
        print(f"Final Reward (last 100): {np.mean(rewards_history[-100:]):.2f}")
        print(f"Final Cost (last 100): {np.mean(costs_history[-100:]):.2f}k USD")
        print("="*80 + "\n")
    
    # Save model
    model_path = output_path / "models" / f"markov_fleet_dqn_final_{n_episodes}ep.pt"
    torch.save({
        'agent_state_dict': agent.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'episodes': episodes_completed,
        'rewards_history': rewards_history,
        'costs_history': costs_history,
        'losses_history': losses_history,
        'config': {
            'n_urban': n_urban,
            'n_rural': n_rural,
            'horizon_years': horizon_years,
            'cost_lambda': cost_lambda,
            'gamma': gamma,
            'lr': lr,
        }
    }, model_path)
    
    if verbose:
        print(f"Model saved to: {model_path}")
    
    return agent, rewards_history, costs_history


# ----- Main Entry Point -----

def main():
    parser = argparse.ArgumentParser(description="Train Markov Fleet DQN")
    parser.add_argument('--episodes', type=int, default=1000, help='Number of episodes')
    parser.add_argument('--n-envs', type=int, default=4, help='Number of parallel environments')
    parser.add_argument('--n-urban', type=int, default=20, help='Number of urban bridges')
    parser.add_argument('--n-rural', type=int, default=80, help='Number of rural bridges')
    parser.add_argument('--horizon', type=int, default=30, help='Episode horizon (years)')
    parser.add_argument('--cost-lambda', type=float, default=1e-3, help='Cost penalty scaling')
    parser.add_argument('--gamma', type=float, default=0.95, help='Discount factor')
    parser.add_argument('--lr', type=float, default=1.5e-3, help='Learning rate')
    parser.add_argument('--eps-start', type=float, default=1.0, help='Initial exploration rate')
    parser.add_argument('--eps-end', type=float, default=0.1, help='Final exploration rate')
    parser.add_argument('--eps-decay', type=int, default=2000, help='Epsilon decay episodes')
    parser.add_argument('--buffer-size', type=int, default=10000, help='Replay buffer size')
    parser.add_argument('--batch-size', type=int, default=64, help='Batch size')
    parser.add_argument('--target-sync', type=int, default=500, help='Target network sync steps')
    parser.add_argument('--device', type=str, default='cuda', help='Device (cuda/cpu)')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    parser.add_argument('--output', type=str, default='outputs_markov', help='Output directory')
    
    args = parser.parse_args()
    
    # Train
    agent, rewards, costs = train_markov_fleet(
        n_episodes=args.episodes,
        n_envs=args.n_envs,
        n_urban=args.n_urban,
        n_rural=args.n_rural,
        horizon_years=args.horizon,
        cost_lambda=args.cost_lambda,
        gamma=args.gamma,
        lr=args.lr,
        eps_start=args.eps_start,
        eps_end=args.eps_end,
        eps_decay_episodes=args.eps_decay,
        buffer_capacity=args.buffer_size,
        batch_size=args.batch_size,
        target_sync_steps=args.target_sync,
        device=args.device,
        seed=args.seed,
        output_dir=args.output,
        verbose=True,
    )
    
    print("\n✓ Training complete!")


if __name__ == "__main__":
    main()