plot_training_progress.py

"""
Visualize training progress and identify when policy stabilizes.
Plots metrics from TensorBoard logs and evaluation results.
"""
import argparse
import os
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

try:
    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
    TENSORBOARD_AVAILABLE = True
except ImportError:
    TENSORBOARD_AVAILABLE = False
    print("Warning: TensorBoard not available. Install with: pip install tensorboard")


def load_evaluation_data(eval_path):
    """Load evaluation data from npz file."""
    if not os.path.exists(eval_path):
        print(f"Evaluation file not found: {eval_path}")
        return None
    
    data = np.load(eval_path)
    return {
        'timesteps': data['timesteps'],
        'results': data['results'],
        'ep_lengths': data['ep_lengths'] if 'ep_lengths' in data else None,
    }


def load_tensorboard_scalars(log_dir):
    """Load scalar metrics from TensorBoard logs."""
    if not TENSORBOARD_AVAILABLE:
        return {}
    
    scalars = {}
    
    # Find all event files
    event_files = []
    for root, dirs, files in os.walk(log_dir):
        for file in files:
            if file.startswith('events.out.tfevents'):
                event_files.append(os.path.join(root, file))
    
    if not event_files:
        print(f"No TensorBoard event files found in {log_dir}")
        return {}
    
    # Load the most recent event file
    latest_event = max(event_files, key=os.path.getmtime)
    print(f"Loading TensorBoard data from: {latest_event}")
    
    try:
        ea = EventAccumulator(os.path.dirname(latest_event))
        ea.Reload()
        
        scalar_tags = ea.Tags()['scalars']
        for tag in scalar_tags:
            scalar_events = ea.Scalars(tag)
            steps = [e.step for e in scalar_events]
            values = [e.value for e in scalar_events]
            scalars[tag] = {'steps': steps, 'values': values}
        
        return scalars
    except Exception as e:
        print(f"Error loading TensorBoard data: {e}")
        return {}


def calculate_stability_metrics(timesteps, rewards, window_size=5):
    """Calculate stability metrics (moving average, std, trend)."""
    if len(rewards) < window_size:
        return None, None, None
    
    # Moving average
    moving_avg = []
    moving_std = []
    
    for i in range(window_size - 1, len(rewards)):
        window_rewards = rewards[i - window_size + 1:i + 1]
        moving_avg.append(np.mean(window_rewards))
        moving_std.append(np.std(window_rewards))
    
    moving_timesteps = timesteps[window_size - 1:]
    
    # Calculate trend (slope of recent points)
    if len(moving_avg) >= 10:
        recent_timesteps = np.array(moving_timesteps[-10:])
        recent_avg = np.array(moving_avg[-10:])
        trend = np.polyfit(recent_timesteps, recent_avg, 1)[0]
    else:
        trend = None
    
    return moving_timesteps, moving_avg, moving_std, trend


def plot_training_progress(eval_data=None, tb_scalars=None, output_path=None, 
                          stability_window=5, show_plot=True):
    """Plot training progress and stability metrics."""
    
    fig = plt.figure(figsize=(16, 10))
    
    # Plot 1: Evaluation Rewards
    ax1 = plt.subplot(2, 3, 1)
    if eval_data and len(eval_data['timesteps']) > 0:
        timesteps = eval_data['timesteps']
        rewards = eval_data['results'].mean(axis=1) if eval_data['results'].ndim > 1 else eval_data['results']
        reward_stds = eval_data['results'].std(axis=1) if eval_data['results'].ndim > 1 else np.zeros_like(rewards)
        
        ax1.plot(timesteps, rewards, 'b-', label='Mean Reward', linewidth=2)
        ax1.fill_between(timesteps, rewards - reward_stds, rewards + reward_stds, 
                         alpha=0.3, label='±1 std')
        
        # Calculate and plot moving average
        mov_ts, mov_avg, mov_std, trend = calculate_stability_metrics(timesteps, rewards, stability_window)
        if mov_avg:
            ax1.plot(mov_ts, mov_avg, 'r--', label=f'Moving Avg ({stability_window} points)', linewidth=2)
            ax1.fill_between(mov_ts, 
                            np.array(mov_avg) - np.array(mov_std),
                            np.array(mov_avg) + np.array(mov_std),
                            alpha=0.2, color='red', label='Moving Std')
            
            # Highlight stabilization point (when std becomes small and trend is flat)
            if trend is not None:
                # Find where std is low and trend is near zero
                recent_std = np.array(mov_std[-10:])
                if len(recent_std) > 0:
                    min_std_idx = len(mov_std) - len(recent_std) + np.argmin(recent_std)
                    if abs(trend) < 0.1:  # Trend is nearly flat
                        ax1.axvline(mov_ts[min_std_idx], color='green', linestyle=':', 
                                   linewidth=2, label='Stabilization Point')
        
        ax1.set_xlabel('Training Timesteps')
        ax1.set_ylabel('Mean Episode Reward')
        ax1.set_title('Evaluation Rewards Over Time')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
    else:
        ax1.text(0.5, 0.5, 'No evaluation data available', 
                ha='center', va='center', transform=ax1.transAxes)
        ax1.set_title('Evaluation Rewards Over Time')
    
    # Plot 2: Episode Lengths
    ax2 = plt.subplot(2, 3, 2)
    if eval_data and eval_data['ep_lengths'] is not None and len(eval_data['ep_lengths']) > 0:
        timesteps = eval_data['timesteps']
        ep_lengths = eval_data['ep_lengths'].mean(axis=1) if eval_data['ep_lengths'].ndim > 1 else eval_data['ep_lengths']
        ep_stds = eval_data['ep_lengths'].std(axis=1) if eval_data['ep_lengths'].ndim > 1 else np.zeros_like(ep_lengths)
        
        ax2.plot(timesteps, ep_lengths, 'g-', label='Mean Episode Length', linewidth=2)
        ax2.fill_between(timesteps, ep_lengths - ep_stds, ep_lengths + ep_stds, 
                         alpha=0.3, label='±1 std')
        
        mov_ts, mov_avg, mov_std, _ = calculate_stability_metrics(timesteps, ep_lengths, stability_window)
        if mov_avg:
            ax2.plot(mov_ts, mov_avg, 'r--', label=f'Moving Avg ({stability_window} points)', linewidth=2)
        
        ax2.set_xlabel('Training Timesteps')
        ax2.set_ylabel('Mean Episode Length')
        ax2.set_title('Episode Lengths Over Time')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
    else:
        ax2.text(0.5, 0.5, 'No episode length data available', 
                ha='center', va='center', transform=ax2.transAxes)
        ax2.set_title('Episode Lengths Over Time')
    
    # Plot 3: Reward Variance (Stability Indicator)
    ax3 = plt.subplot(2, 3, 3)
    if eval_data and len(eval_data['timesteps']) > 0:
        timesteps = eval_data['timesteps']
        rewards = eval_data['results'].mean(axis=1) if eval_data['results'].ndim > 1 else eval_data['results']
        reward_stds = eval_data['results'].std(axis=1) if eval_data['results'].ndim > 1 else np.zeros_like(rewards)
        
        # Moving standard deviation
        mov_ts, _, mov_std, _ = calculate_stability_metrics(timesteps, rewards, stability_window)
        if mov_std:
            ax3.plot(mov_ts, mov_std, 'purple', label='Moving Std Dev', linewidth=2)
            ax3.axhline(np.mean(mov_std[-10:]), color='orange', linestyle='--', 
                        label=f'Recent Avg Std: {np.mean(mov_std[-10:]):.2f}')
            ax3.set_xlabel('Training Timesteps')
            ax3.set_ylabel('Reward Standard Deviation')
            ax3.set_title('Reward Stability (Lower = More Stable)')
            ax3.legend()
            ax3.grid(True, alpha=0.3)
    else:
        ax3.text(0.5, 0.5, 'No data available', 
                ha='center', va='center', transform=ax3.transAxes)
        ax3.set_title('Reward Stability')
    
    # Plot 4: Training Loss (from TensorBoard)
    ax4 = plt.subplot(2, 3, 4)
    if tb_scalars:
        plotted = False
        for tag in ['train/loss', 'train/policy_loss', 'train/value_loss']:
            if tag in tb_scalars:
                steps = tb_scalars[tag]['steps']
                values = tb_scalars[tag]['values']
                ax4.plot(steps, values, label=tag.split('/')[-1], linewidth=2)
                plotted = True
        if plotted:
            ax4.set_xlabel('Training Timesteps')
            ax4.set_ylabel('Loss')
            ax4.set_title('Training Losses')
            ax4.legend()
            ax4.grid(True, alpha=0.3)
            ax4.set_yscale('log')
        else:
            ax4.text(0.5, 0.5, 'No loss data in TensorBoard logs', 
                    ha='center', va='center', transform=ax4.transAxes)
            ax4.set_title('Training Losses')
    else:
        ax4.text(0.5, 0.5, 'No TensorBoard data available', 
                ha='center', va='center', transform=ax4.transAxes)
        ax4.set_title('Training Losses')
    
    # Plot 5: Learning Rate (from TensorBoard)
    ax5 = plt.subplot(2, 3, 5)
    if tb_scalars and 'train/learning_rate' in tb_scalars:
        steps = tb_scalars['train/learning_rate']['steps']
        values = tb_scalars['train/learning_rate']['values']
        ax5.plot(steps, values, 'orange', linewidth=2)
        ax5.set_xlabel('Training Timesteps')
        ax5.set_ylabel('Learning Rate')
        ax5.set_title('Learning Rate Schedule')
        ax5.grid(True, alpha=0.3)
        ax5.set_yscale('log')
    else:
        ax5.text(0.5, 0.5, 'No learning rate data available', 
                ha='center', va='center', transform=ax5.transAxes)
        ax5.set_title('Learning Rate Schedule')
    
    # Plot 6: Summary Statistics
    ax6 = plt.subplot(2, 3, 6)
    ax6.axis('off')
    
    summary_text = "Training Progress Summary\n" + "="*40 + "\n\n"
    
    if eval_data and len(eval_data['timesteps']) > 0:
        timesteps = eval_data['timesteps']
        rewards = eval_data['results'].mean(axis=1) if eval_data['results'].ndim > 1 else eval_data['results']
        
        summary_text += f"Total Evaluations: {len(timesteps)}\n"
        summary_text += f"Latest Timestep: {timesteps[-1]:,}\n"
        summary_text += f"Latest Reward: {rewards[-1]:.2f}\n"
        summary_text += f"Best Reward: {np.max(rewards):.2f}\n"
        summary_text += f"Mean Reward: {np.mean(rewards):.2f}\n\n"
        
        # Stability analysis
        mov_ts, mov_avg, mov_std, trend = calculate_stability_metrics(timesteps, rewards, stability_window)
        if mov_avg and len(mov_avg) > 0:
            summary_text += "Stability Analysis:\n"
            summary_text += f"  Recent Avg Reward: {mov_avg[-1]:.2f}\n"
            summary_text += f"  Recent Std Dev: {mov_std[-1]:.2f}\n"
            if trend is not None:
                summary_text += f"  Recent Trend: {trend*1e6:.2f} per 1M steps\n"
                # Check both trend and recent variance
                recent_std = mov_std[-1] if len(mov_std) > 0 else 0
                is_stable = abs(trend) < 0.1 and recent_std < 15
                if is_stable:
                    summary_text += "  Status: STABLE ✓\n"
                else:
                    status_parts = []
                    if abs(trend) >= 0.1:
                        status_parts.append(f"Trend: {trend*1e6:.1f}/1M steps")
                    if recent_std >= 15:
                        status_parts.append(f"High variance: {recent_std:.1f}")
                    summary_text += f"  Status: Still Learning ({', '.join(status_parts)})\n"
    
    if tb_scalars:
        summary_text += "\nTensorBoard Metrics Available:\n"
        for tag in sorted(tb_scalars.keys()):
            summary_text += f"  - {tag}\n"
    
    ax6.text(0.1, 0.9, summary_text, transform=ax6.transAxes, 
            fontsize=10, verticalalignment='top', family='monospace')
    
    plt.tight_layout()
    
    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        print(f"\nPlot saved to: {output_path}")
    
    if show_plot:
        plt.show()
    else:
        plt.close()


def main():
    parser = argparse.ArgumentParser(description="Visualize training progress and stability")
    parser.add_argument("--log_dir", type=str, default="./logs",
                       help="Directory containing TensorBoard logs")
    parser.add_argument("--eval_file", type=str, default="./logs/evaluations/evaluations.npz",
                       help="Path to evaluation results file")
    parser.add_argument("--output", type=str, default=None,
                       help="Path to save plot (e.g., training_progress.png)")
    parser.add_argument("--stability_window", type=int, default=5,
                       help="Window size for moving average (default: 5)")
    parser.add_argument("--no-show", action="store_true",
                       help="Don't display plot (useful when saving)")
    
    args = parser.parse_args()
    
    print("Loading training data...")
    
    # Load evaluation data
    eval_data = load_evaluation_data(args.eval_file)
    
    # Load TensorBoard data
    tb_scalars = {}
    if os.path.exists(args.log_dir):
        tb_scalars = load_tensorboard_scalars(args.log_dir)
    
    if not eval_data and not tb_scalars:
        print("Error: No training data found!")
        print(f"  Checked evaluation file: {args.eval_file}")
        print(f"  Checked log directory: {args.log_dir}")
        return
    
    print("\nGenerating plots...")
    plot_training_progress(
        eval_data=eval_data,
        tb_scalars=tb_scalars,
        output_path=args.output,
        stability_window=args.stability_window,
        show_plot=not args.no_show
    )


if __name__ == "__main__":
    main()