train.py

"""Train RL agents for portfolio optimization"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from data_generator import generate_synthetic_data
from portfolio_env import PortfolioEnv
from ppo_agent import PPOAgent
from sac_agent import SACAgent

def train_ppo(env, n_episodes=200, steps_per_episode=None, update_freq=2048):
    """Train PPO agent"""
    print("\n" + "="*60)
    print("TRAINING PPO AGENT")
    print("="*60)
    
    agent = PPOAgent(
        obs_dim=env.observation_dim,
        action_dim=env.action_dim,
        lr=3e-4,
        gamma=0.99,
        epsilon=0.2
    )
    
    episode_rewards = []
    episode_returns = []
    
    total_steps = 0
    
    for episode in range(n_episodes):
        obs = env.reset(seed=episode)
        episode_reward = 0
        done = False
        
        while not done:
            # Get action
            action, log_prob = agent.get_action(obs)
            value = agent.get_value(obs)
            
            # Take step
            next_obs, reward, done, truncated, info = env.step(action)
            
            # Store transition
            agent.store_transition(obs, action, reward, value, log_prob)
            
            episode_reward += reward
            total_steps += 1
            obs = next_obs
            
            # Update policy
            if total_steps % update_freq == 0:
                losses = agent.update(n_epochs=10, batch_size=64)
                if losses:
                    print(f"  Step {total_steps}: Policy Loss={losses['policy_loss']:.4f}, Value Loss={losses['value_loss']:.4f}")
        
        # Episode done
        metrics = env.get_performance_metrics()
        episode_rewards.append(episode_reward)
        episode_returns.append(metrics['total_return'])
        
        if (episode + 1) % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            avg_return = np.mean(episode_returns[-10:])
            print(f"Episode {episode+1}/{n_episodes}: Avg Reward={avg_reward:.2f}, Avg Return={avg_return:.2%}, Sharpe={metrics['sharpe_ratio']:.2f}")
    
    # Save model
    agent.save('ppo_model.npz')
    print("\nPPO training complete. Model saved to ppo_model.npz")
    
    return agent, episode_rewards, episode_returns

def train_sac(env, n_episodes=200, update_freq=1):
    """Train SAC agent"""
    print("\n" + "="*60)
    print("TRAINING SAC AGENT")
    print("="*60)
    
    agent = SACAgent(
        obs_dim=env.observation_dim,
        action_dim=env.action_dim,
        lr=3e-4,
        gamma=0.99,
        tau=0.005,
        alpha=0.2,
        buffer_size=10000
    )
    
    episode_rewards = []
    episode_returns = []
    
    for episode in range(n_episodes):
        obs = env.reset(seed=episode)
        episode_reward = 0
        done = False
        step = 0
        
        while not done:
            # Get action
            action = agent.get_action(obs, deterministic=False)
            
            # Take step
            next_obs, reward, done, truncated, info = env.step(action)
            
            # Store transition
            agent.store_transition(obs, action, reward, next_obs, done)
            
            # Update policy
            if step % update_freq == 0 and len(agent.buffer) >= 64:
                losses = agent.update(batch_size=64)
                if losses and step % 100 == 0:
                    print(f"  Episode {episode+1}, Step {step}: Q1 Loss={losses['q1_loss']:.4f}, Actor Loss={losses['actor_loss']:.4f}")
            
            episode_reward += reward
            step += 1
            obs = next_obs
        
        # Episode done
        metrics = env.get_performance_metrics()
        episode_rewards.append(episode_reward)
        episode_returns.append(metrics['total_return'])
        
        if (episode + 1) % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            avg_return = np.mean(episode_returns[-10:])
            print(f"Episode {episode+1}/{n_episodes}: Avg Reward={avg_reward:.2f}, Avg Return={avg_return:.2%}, Sharpe={metrics['sharpe_ratio']:.2f}")
    
    # Save model
    agent.save('sac_model.npz')
    print("\nSAC training complete. Model saved to sac_model.npz")
    
    return agent, episode_rewards, episode_returns

def evaluate_agent(agent, env, agent_type='PPO', deterministic=True):
    """Evaluate trained agent"""
    print(f"\n{'='*60}")
    print(f"EVALUATING {agent_type} AGENT")
    print("="*60)
    
    obs = env.reset(seed=999)
    done = False
    
    while not done:
        if agent_type == 'PPO':
            action = agent.get_action(obs, deterministic=deterministic)
            if isinstance(action, tuple):
                action = action[0]
        else:  # SAC
            action = agent.get_action(obs, deterministic=deterministic)
        
        obs, reward, done, truncated, info = env.step(action)
    
    # Get performance metrics
    metrics = env.get_performance_metrics()
    
    print(f"\nPerformance Metrics:")
    print(f"  Total Return: {metrics['total_return']:.2%}")
    print(f"  Sharpe Ratio: {metrics['sharpe_ratio']:.2f}")
    print(f"  Max Drawdown: {metrics['max_drawdown']:.2%}")
    print(f"  Avg Turnover: {metrics['avg_turnover']:.2%}")
    print(f"  Final Value: ${metrics['final_value']:,.2f}")
    print(f"  Number of Trades: {metrics['n_trades']}")
    
    return env, metrics

def plot_results(env, metrics, agent_type='PPO'):
    """Plot portfolio performance"""
    fig, axes = plt.subplots(3, 1, figsize=(12, 10))
    
    # Portfolio value
    axes[0].plot(env.portfolio_history, label=f'{agent_type} Portfolio')
    axes[0].axhline(y=env.initial_capital, color='gray', linestyle='--', label='Initial Capital')
    axes[0].set_title(f'{agent_type} Portfolio Value Over Time')
    axes[0].set_ylabel('Portfolio Value ($)')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Portfolio weights
    weights_array = np.array(env.weights_history)
    for i, asset in enumerate(env.asset_names):
        axes[1].plot(weights_array[:, i], label=asset, alpha=0.7)
    axes[1].set_title('Portfolio Weights Over Time')
    axes[1].set_ylabel('Weight')
    axes[1].set_ylim([0, 1])
    axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    axes[1].grid(True, alpha=0.3)
    
    # Drawdown
    portfolio_values = np.array(env.portfolio_history)
    cummax = np.maximum.accumulate(portfolio_values)
    drawdown = (portfolio_values - cummax) / cummax * 100
    axes[2].fill_between(range(len(drawdown)), drawdown, 0, alpha=0.3, color='red')
    axes[2].set_title('Drawdown (%)')
    axes[2].set_ylabel('Drawdown (%)')
    axes[2].set_xlabel('Time Steps')
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(f'{agent_type.lower()}_performance.png', dpi=150, bbox_inches='tight')
    print(f"\nPlot saved to {agent_type.lower()}_performance.png")
    plt.close()

def main():
    """Main training and evaluation pipeline"""
    print("="*60)
    print("RL PORTFOLIO OPTIMIZATION SYSTEM")
    print("="*60)
    
    # 1. Generate or load data
    print("\nGenerating synthetic market data...")
    price_data = generate_synthetic_data(n_days=2000, seed=42)
    print(f"Generated {len(price_data)} days of data for {len(price_data.columns)} assets")
    
    # 2. Split data: train/test
    split_idx = int(len(price_data) * 0.7)
    train_data = price_data.iloc[:split_idx]
    test_data = price_data.iloc[split_idx:]
    
    print(f"Training data: {len(train_data)} days")
    print(f"Testing data: {len(test_data)} days")
    
    # 3. Create environments
    train_env = PortfolioEnv(
        train_data,
        window_size=20,
        initial_capital=100000,
        transaction_cost=0.001,
        reward_scaling=1.0
    )
    
    test_env = PortfolioEnv(
        test_data,
        window_size=20,
        initial_capital=100000,
        transaction_cost=0.001,
        reward_scaling=1.0
    )
    
    print(f"Observation dim: {train_env.observation_dim}")
    print(f"Action dim: {train_env.action_dim}")
    
    # 4. Train PPO
    ppo_agent, ppo_rewards, ppo_returns = train_ppo(
        train_env, 
        n_episodes=100,
        update_freq=512
    )
    
    # 5. Evaluate PPO
    test_env_ppo, ppo_metrics = evaluate_agent(ppo_agent, test_env, agent_type='PPO')
    plot_results(test_env_ppo, ppo_metrics, agent_type='PPO')
    
    # 6. Train SAC
    sac_agent, sac_rewards, sac_returns = train_sac(
        train_env,
        n_episodes=100,
        update_freq=1
    )
    
    # 7. Evaluate SAC
    test_env_sac, sac_metrics = evaluate_agent(sac_agent, test_env, agent_type='SAC')
    plot_results(test_env_sac, sac_metrics, agent_type='SAC')
    
    # 8. Compare with equal-weight benchmark
    print("\n" + "="*60)
    print("EQUAL-WEIGHT BENCHMARK")
    print("="*60)
    benchmark_env = PortfolioEnv(
        test_data,
        window_size=20,
        initial_capital=100000,
        transaction_cost=0.001
    )
    obs = benchmark_env.reset(seed=999)
    done = False
    equal_weights = np.ones(benchmark_env.n_assets) / benchmark_env.n_assets
    while not done:
        obs, reward, done, truncated, info = benchmark_env.step(equal_weights)
    
    benchmark_metrics = benchmark_env.get_performance_metrics()
    print(f"\nBenchmark Metrics:")
    print(f"  Total Return: {benchmark_metrics['total_return']:.2%}")
    print(f"  Sharpe Ratio: {benchmark_metrics['sharpe_ratio']:.2f}")
    print(f"  Max Drawdown: {benchmark_metrics['max_drawdown']:.2%}")
    
    # 9. Summary comparison
    print("\n" + "="*60)
    print("FINAL COMPARISON")
    print("="*60)
    print(f"\n{'Metric':<20} {'PPO':<15} {'SAC':<15} {'Benchmark':<15}")
    print("-" * 65)
    print(f"{'Total Return':<20} {ppo_metrics['total_return']:>14.2%} {sac_metrics['total_return']:>14.2%} {benchmark_metrics['total_return']:>14.2%}")
    print(f"{'Sharpe Ratio':<20} {ppo_metrics['sharpe_ratio']:>14.2f} {sac_metrics['sharpe_ratio']:>14.2f} {benchmark_metrics['sharpe_ratio']:>14.2f}")
    print(f"{'Max Drawdown':<20} {ppo_metrics['max_drawdown']:>14.2%} {sac_metrics['max_drawdown']:>14.2%} {benchmark_metrics['max_drawdown']:>14.2%}")
    print(f"{'Avg Turnover':<20} {ppo_metrics['avg_turnover']:>14.2%} {sac_metrics['avg_turnover']:>14.2%} {benchmark_metrics['avg_turnover']:>14.2%}")
    
    print("\n" + "="*60)
    print("TRAINING COMPLETE")
    print("="*60)
    print("\nModels saved:")
    print("  - ppo_model.npz")
    print("  - sac_model.npz")
    print("\nPlots saved:")
    print("  - ppo_performance.png")
    print("  - sac_performance.png")

if __name__ == '__main__':
    main()