train_fleet_phase4.py

"""
Phase 4: Ultimate Integration (PER + Vectorization + Phase 2.3 Training Strategy)
Combines the best of Phase 2.3 (stable learning) and Phase 3 (speed)
"""

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 sys
import time
from torch.amp import autocast, GradScaler
import gymnasium as gym
from gymnasium.vector import AsyncVectorEnv

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

from fleet_environment_gym import (
    FleetManagementGym, FleetConfig, UrbanAgentDQN, RuralAgentDQN
)


class PrioritizedNStepBuffer:
    """Prioritized N-step replay buffer"""
    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, urban_state, urban_actions, rural_state, rural_action,
             reward, next_urban_state, next_rural_state, done):
        """Add experience to n-step buffer"""
        self.n_step_buffer.append((
            urban_state, urban_actions, rural_state, rural_action,
            reward, next_urban_state, next_rural_state, done
        ))
        
        if len(self.n_step_buffer) < self.n_steps:
            return
        
        # Compute n-step return
        n_step_reward = 0
        for i, (_, _, _, _, r, _, _, _) in enumerate(self.n_step_buffer):
            n_step_reward += (self.gamma ** i) * r
        
        urban_s, urban_a, rural_s, rural_a, _, _, _, _ = self.n_step_buffer[0]
        _, _, _, _, _, next_urban_s, next_rural_s, done = self.n_step_buffer[-1]
        
        max_priority = self.priorities[:self.size].max() if self.size > 0 else 1.0
        
        if len(self.buffer) < self.capacity:
            self.buffer.append((
                urban_s, urban_a, rural_s, rural_a,
                n_step_reward, next_urban_s, next_rural_s, done
            ))
        else:
            self.buffer[self.position] = (
                urban_s, urban_a, rural_s, rural_a,
                n_step_reward, next_urban_s, next_rural_s, done
            )
        
        self.priorities[self.position] = max_priority
        self.position = (self.position + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)
    
    def sample(self, batch_size: int):
        """Sample with prioritization"""
        if self.size == 0:
            return None
        
        priorities = self.priorities[:self.size]
        probs = priorities ** self.alpha
        
        # NaN safety
        if np.isnan(probs).any() or probs.sum() == 0:
            probs = np.ones_like(probs) / len(probs)
        else:
            probs = probs / probs.sum()
        
        probs = np.clip(probs, 1e-8, 1.0)
        probs = probs / probs.sum()
        
        indices = np.random.choice(self.size, min(batch_size, self.size), p=probs, replace=False)
        
        samples = [self.buffer[idx] for idx in indices]
        
        # Importance sampling weights
        weights = (self.size * probs[indices]) ** (-self.beta)
        weights = weights / weights.max()
        
        self.beta = min(1.0, self.beta + self.beta_increment)
        
        return samples, indices, weights
    
    def update_priorities(self, indices, td_errors):
        """Update priorities based on TD errors"""
        for idx, error in zip(indices, td_errors):
            self.priorities[idx] = abs(error) + 1e-6
    
    def __len__(self):
        return self.size


def make_env():
    """Factory function to create environment"""
    def _init():
        return FleetManagementGym()
    return _init


def train_phase4(cfg, n_episodes, n_envs, output_dir, phase='4'):
    """
    Phase 4 Training: PER + Vectorization + Phase 2.3 Strategy
    
    Key improvements over Phase 3:
    1. Epsilon-greedy exploration (Phase 2.3 style)
    2. Aggressive training (higher batch size, more frequent updates)
    3. Budget-aware reward shaping
    """
    
    save_dir = Path(output_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    (save_dir / "models").mkdir(exist_ok=True)
    
    # Create vectorized environments
    print(f"\nCreating {n_envs} parallel environments...")
    envs = AsyncVectorEnv([make_env() for _ in range(n_envs)])
    
    # Initialize agents (same as Phase 2.3)
    urban_agent = UrbanAgentDQN(state_dim=81, n_bridges=20, n_actions=5).to(cfg.device)
    urban_target = UrbanAgentDQN(state_dim=81, n_bridges=20, n_actions=5).to(cfg.device)
    urban_target.load_state_dict(urban_agent.state_dict())
    urban_target.eval()
    
    rural_agent = RuralAgentDQN(state_dim=10, n_strategies=8).to(cfg.device)
    rural_target = RuralAgentDQN(state_dim=10, n_strategies=8).to(cfg.device)
    rural_target.load_state_dict(rural_agent.state_dict())
    rural_target.eval()
    
    # Phase 2.3 style optimizers
    urban_optimizer = optim.AdamW(urban_agent.parameters(), lr=cfg.lr, weight_decay=1e-5)
    rural_optimizer = optim.AdamW(rural_agent.parameters(), lr=cfg.lr, weight_decay=1e-5)
    
    scaler = GradScaler('cuda') if cfg.device == 'cuda' else None
    buffer = PrioritizedNStepBuffer(cfg.buffer_capacity, n_steps=3, gamma=cfg.gamma,
                                     alpha=0.6, beta=0.4, beta_increment=0.001)
    
    # Phase 4 specific: Epsilon-greedy exploration (Phase 2.3 style)
    epsilon_start = 1.0
    epsilon_end = 0.01
    epsilon_decay = 0.995
    epsilon = epsilon_start
    
    # Training tracking
    rewards_history = []
    costs_history = []
    losses_urban = []
    losses_rural = []
    total_steps = 0
    episodes_completed = 0
    start_time = time.time()
    
    # Per-environment episode tracking
    env_episode_rewards = np.zeros(n_envs)
    env_episode_costs = np.zeros(n_envs)
    
    print("\n" + "="*80)
    print(f"PHASE 4: ULTIMATE INTEGRATION ({n_envs} parallel environments)")
    print("Optimizations: AMP + Double DQN + Dueling DQN + N-step + PER + Vectorization")
    print("+ Phase 2.3 Training Strategy (epsilon-greedy, aggressive updates)")
    print("="*80)
    print(f"Target Episodes: {n_episodes}")
    print(f"Parallel Envs: {n_envs}")
    print(f"Device: {cfg.device}")
    print(f"Epsilon: {epsilon_start} → {epsilon_end} (decay={epsilon_decay})")
    print(f"Batch Size: {cfg.batch_size}")
    print("="*80 + "\n")
    
    # Reset environments
    observations, _ = envs.reset()
    urban_states = observations['urban']
    rural_states = observations['rural']
    
    pbar = tqdm(total=n_episodes, desc="Training Phase 4")
    
    while episodes_completed < n_episodes:
        # Epsilon-greedy action selection (Phase 2.3 style)
        if np.random.random() < epsilon:
            # Random exploration
            urban_actions_batch = np.random.randint(0, 5, size=(n_envs, 20))
            rural_actions_batch = np.random.randint(0, 8, size=n_envs)
        else:
            # Greedy exploitation
            with torch.no_grad():
                urban_states_t = torch.FloatTensor(urban_states).to(cfg.device)
                urban_q_values = urban_agent(urban_states_t)
                urban_actions_batch = urban_q_values.argmax(dim=2).cpu().numpy()
                
                rural_states_t = torch.FloatTensor(rural_states).to(cfg.device)
                rural_q_values = rural_agent(rural_states_t)
                rural_actions_batch = rural_q_values.argmax(dim=1).cpu().numpy()
        
        # Convert to dict format for vectorized env
        actions = {
            'urban': urban_actions_batch,
            'rural': rural_actions_batch
        }
        
        # Step all environments
        next_observations, rewards, terminateds, truncateds, infos = envs.step(actions)
        next_urban_states = next_observations['urban']
        next_rural_states = next_observations['rural']
        
        # Store transitions for all environments
        for i in range(n_envs):
            buffer.push(
                urban_states[i], urban_actions_batch[i],
                rural_states[i], rural_actions_batch[i],
                rewards[i], next_urban_states[i], next_rural_states[i],
                terminateds[i] or truncateds[i]
            )
            
            env_episode_rewards[i] += rewards[i]
            info = infos.get(i, {}) if isinstance(infos, dict) else (infos[i] if i < len(infos) else {})
            env_episode_costs[i] += info.get('total_cost', 0)
            
            # Episode完了チェック
            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
                pbar.update(1)
                
                # Decay epsilon (Phase 2.3 style)
                epsilon = max(epsilon_end, epsilon * epsilon_decay)
                
                if episodes_completed >= n_episodes:
                    break
        
        urban_states = next_urban_states
        rural_states = next_rural_states
        total_steps += n_envs
        
        # Training (Phase 2.3 style: aggressive updates)
        if len(buffer) >= cfg.batch_size and total_steps % 4 == 0:  # More frequent updates
            sample_result = buffer.sample(cfg.batch_size)
            if sample_result is None:
                continue
            
            samples, sample_indices, importance_weights = sample_result
            
            # Unpack batch
            batch_urban_states = torch.FloatTensor(np.array([s[0] for s in samples])).to(cfg.device)
            batch_urban_actions = torch.LongTensor(np.array([s[1] for s in samples])).to(cfg.device)
            batch_rural_states = torch.FloatTensor(np.array([s[2] for s in samples])).to(cfg.device)
            batch_rural_actions = torch.LongTensor(np.array([s[3] for s in samples])).to(cfg.device)
            batch_rewards = torch.FloatTensor(np.array([s[4] for s in samples])).to(cfg.device)
            batch_next_urban_states = torch.FloatTensor(np.array([s[5] for s in samples])).to(cfg.device)
            batch_next_rural_states = torch.FloatTensor(np.array([s[6] for s in samples])).to(cfg.device)
            batch_dones = torch.FloatTensor(np.array([s[7] for s in samples])).to(cfg.device)
            importance_weights_t = torch.FloatTensor(importance_weights).to(cfg.device)
            
            # Urban agent training
            urban_optimizer.zero_grad()
            if cfg.device == 'cuda' and scaler is not None:
                with autocast('cuda'):
                    urban_q_values = urban_agent(batch_urban_states)
                    batch_size, n_bridges, n_actions = urban_q_values.shape
                    urban_q_flat = urban_q_values.reshape(batch_size * n_bridges, n_actions)
                    actions_flat = batch_urban_actions.reshape(-1)
                    urban_q_selected = urban_q_flat.gather(1, actions_flat.unsqueeze(1)).squeeze()
                    
                    with torch.no_grad():
                        next_urban_actions = urban_agent(batch_next_urban_states).argmax(dim=2)
                        next_urban_q_values = urban_target(batch_next_urban_states)
                        next_urban_actions_flat = next_urban_actions.reshape(-1)
                        next_urban_q_flat = next_urban_q_values.reshape(batch_size * n_bridges, n_actions)
                        next_urban_q = next_urban_q_flat.gather(1, next_urban_actions_flat.unsqueeze(1)).squeeze()
                        next_urban_q = next_urban_q.reshape(batch_size, n_bridges).mean(dim=1)
                        
                        batch_rewards_expanded = batch_rewards.unsqueeze(1).expand(-1, n_bridges).reshape(-1)
                        batch_dones_expanded = batch_dones.unsqueeze(1).expand(-1, n_bridges).reshape(-1)
                        urban_targets = batch_rewards_expanded + (cfg.gamma ** buffer.n_steps) * next_urban_q.unsqueeze(1).expand(-1, n_bridges).reshape(-1) * (1 - batch_dones_expanded)
                    
                    urban_td_errors = torch.abs(urban_q_selected - urban_targets)
                    importance_weights_expanded = importance_weights_t.unsqueeze(1).expand(-1, n_bridges).reshape(-1)
                    urban_loss = (importance_weights_expanded * (urban_td_errors ** 2)).mean()
                
                scaler.scale(urban_loss).backward()
                scaler.unscale_(urban_optimizer)
                torch.nn.utils.clip_grad_norm_(urban_agent.parameters(), cfg.gradient_clip)
                scaler.step(urban_optimizer)
                scaler.update()
            else:
                urban_q_values = urban_agent(batch_urban_states)
                batch_size, n_bridges, n_actions = urban_q_values.shape
                urban_q_flat = urban_q_values.reshape(batch_size * n_bridges, n_actions)
                actions_flat = batch_urban_actions.reshape(-1)
                urban_q_selected = urban_q_flat.gather(1, actions_flat.unsqueeze(1)).squeeze()
                
                with torch.no_grad():
                    next_urban_actions = urban_agent(batch_next_urban_states).argmax(dim=2)
                    next_urban_q_values = urban_target(batch_next_urban_states)
                    next_urban_actions_flat = next_urban_actions.reshape(-1)
                    next_urban_q_flat = next_urban_q_values.reshape(batch_size * n_bridges, n_actions)
                    next_urban_q = next_urban_q_flat.gather(1, next_urban_actions_flat.unsqueeze(1)).squeeze()
                    next_urban_q = next_urban_q.reshape(batch_size, n_bridges).mean(dim=1)
                    
                    batch_rewards_expanded = batch_rewards.unsqueeze(1).expand(-1, n_bridges).reshape(-1)
                    batch_dones_expanded = batch_dones.unsqueeze(1).expand(-1, n_bridges).reshape(-1)
                    urban_targets = batch_rewards_expanded + (cfg.gamma ** buffer.n_steps) * next_urban_q.unsqueeze(1).expand(-1, n_bridges).reshape(-1) * (1 - batch_dones_expanded)
                
                urban_td_errors = torch.abs(urban_q_selected - urban_targets)
                importance_weights_expanded = importance_weights_t.unsqueeze(1).expand(-1, n_bridges).reshape(-1)
                urban_loss = (importance_weights_expanded * (urban_td_errors ** 2)).mean()
                urban_loss.backward()
                torch.nn.utils.clip_grad_norm_(urban_agent.parameters(), cfg.gradient_clip)
                urban_optimizer.step()
            
            losses_urban.append(urban_loss.item())
            
            # Rural agent training
            rural_optimizer.zero_grad()
            if cfg.device == 'cuda' and scaler is not None:
                with autocast('cuda'):
                    rural_q_values = rural_agent(batch_rural_states)
                    rural_q_selected = rural_q_values.gather(1, batch_rural_actions.unsqueeze(1)).squeeze()
                    
                    with torch.no_grad():
                        next_rural_actions = rural_agent(batch_next_rural_states).argmax(dim=1)
                        next_rural_q_values = rural_target(batch_next_rural_states)
                        next_rural_q = next_rural_q_values.gather(1, next_rural_actions.unsqueeze(1)).squeeze()
                        rural_targets = batch_rewards + (cfg.gamma ** buffer.n_steps) * next_rural_q * (1 - batch_dones)
                    
                    rural_td_errors = torch.abs(rural_q_selected - rural_targets)
                    rural_loss = (importance_weights_t * (rural_td_errors ** 2)).mean()
                
                scaler.scale(rural_loss).backward()
                scaler.unscale_(rural_optimizer)
                torch.nn.utils.clip_grad_norm_(rural_agent.parameters(), cfg.gradient_clip)
                scaler.step(rural_optimizer)
                scaler.update()
            else:
                rural_q_values = rural_agent(batch_rural_states)
                rural_q_selected = rural_q_values.gather(1, batch_rural_actions.unsqueeze(1)).squeeze()
                
                with torch.no_grad():
                    next_rural_actions = rural_agent(batch_next_rural_states).argmax(dim=1)
                    next_rural_q_values = rural_target(batch_next_rural_states)
                    next_rural_q = next_rural_q_values.gather(1, next_rural_actions.unsqueeze(1)).squeeze()
                    rural_targets = batch_rewards + (cfg.gamma ** buffer.n_steps) * next_rural_q * (1 - batch_dones)
                
                rural_td_errors = torch.abs(rural_q_selected - rural_targets)
                rural_loss = (importance_weights_t * (rural_td_errors ** 2)).mean()
                rural_loss.backward()
                torch.nn.utils.clip_grad_norm_(rural_agent.parameters(), cfg.gradient_clip)
                rural_optimizer.step()
            
            losses_rural.append(rural_loss.item())
            
            # Update priorities
            combined_td_errors = (urban_td_errors.reshape(batch_size, n_bridges).mean(dim=1) + rural_td_errors).detach().cpu().numpy()
            buffer.update_priorities(sample_indices, combined_td_errors)
        
        # Sync target networks
        if total_steps % cfg.target_sync_steps == 0:
            urban_target.load_state_dict(urban_agent.state_dict())
            rural_target.load_state_dict(rural_agent.state_dict())
    
    pbar.close()
    envs.close()
    
    end_time = time.time()
    total_time = end_time - start_time
    
    # Save models
    torch.save(urban_agent, save_dir / "models" / "urban_agent_final.pt")
    torch.save(rural_agent, save_dir / "models" / "rural_agent_final.pt")
    
    # Save metrics
    import json
    metrics = {
        'episode_rewards': rewards_history,
        'episode_costs': costs_history,
        'urban_losses': losses_urban,
        'rural_losses': losses_rural,
        'total_time': total_time,
        'total_steps': total_steps,
        'n_episodes': episodes_completed,
        'n_envs': n_envs,
        'final_epsilon': epsilon
    }
    with open(save_dir / 'metrics.json', 'w') as f:
        json.dump(metrics, f, indent=2)
    
    print("\n" + "="*80)
    print("PHASE 4 TRAINING COMPLETE")
    print("="*80)
    print(f"Episodes: {episodes_completed}")
    print(f"Total Time: {total_time:.2f}s ({total_time/60:.2f} min)")
    print(f"Time per Episode: {total_time/episodes_completed:.3f}s")
    print(f"Final Epsilon: {epsilon:.4f}")
    print(f"Avg Reward (last 100): {np.mean(rewards_history[-100:]):.2f}")
    print("="*80)


def main():
    parser = argparse.ArgumentParser(description="Train Fleet Management Phase 4 (Ultimate)")
    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('--device', type=str, default='cuda', help="Device (cuda/cpu)")
    parser.add_argument('--output', type=str, default='outputs_phase4', help="Output directory")
    
    args = parser.parse_args()
    
    cfg = FleetConfig()
    cfg.device = args.device if torch.cuda.is_available() else 'cpu'
    
    train_phase4(cfg, args.episodes, args.n_envs, args.output, phase='4')


if __name__ == "__main__":
    main()