#!/usr/bin/env python3 """ Comprehensive RL Training + Backtesting + Checkpoint System for Intraday Trading Features: - Model checkpoint management with best model tracking - Training progress resumption from checkpoints - Parameter sweep with progress persistence - Detailed backtest metrics (Sharpe, max drawdown, profit factor, etc.) - Performance comparison across configurations - Tensorboard integration - Detailed logging to files """ from __future__ import annotations import argparse import json import os import shutil import sys from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd from stable_baselines3 import DQN, PPO sys.path.insert(0, str(Path(__file__).parent.parent)) from src.prediction.intraday_rl.backtest import backtest_model from src.prediction.intraday_rl.environment import IntradayEnvConfig, IntradayTradingEnv from src.prediction.intraday_rl.features import ( MinuteDataAugmenter, build_intraday_features, load_ohlcv_csv, resample_ohlcv, split_sessions, ) from src.prediction.intraday_rl.trainer import ( DQNTrainingConfig, PPOTrainingConfig, describe_training_setup, prepare_training_sessions, train_dqn_baseline, train_ppo, ) from stable_baselines3.common.vec_env import DummyVecEnv def resolve_model_path(model_path: str | Path) -> Path: candidate = Path(model_path) tried: List[Path] = [] if candidate.is_dir(): for name in ("latest_model.zip", "best_model.zip"): path = candidate / name tried.append(path) if path.exists(): return path tried.append(candidate) if candidate.exists(): return candidate zip_path = candidate.with_suffix(".zip") tried.append(zip_path) if zip_path.exists(): return zip_path tried_paths = ", ".join(str(path) for path in tried) raise FileNotFoundError(f"Model not found. Tried: {tried_paths}") # ============================================================ # CHECKPOINT MANAGER # ============================================================ class CheckpointManager: """Manages model checkpoints, best models, and metadata.""" def __init__(self, checkpoint_dir: str | Path): self.checkpoint_dir = Path(checkpoint_dir) self.checkpoint_dir.mkdir(parents=True, exist_ok=True) self.metadata_file = self.checkpoint_dir / "checkpoint_metadata.json" self.best_models_file = self.checkpoint_dir / "best_models.json" def get_metadata(self) -> Dict[str, Any]: """Load checkpoint metadata.""" if self.metadata_file.exists(): with open(self.metadata_file) as f: return json.load(f) return {"checkpoints": {}, "last_updated": None} def save_metadata(self, metadata: Dict[str, Any]) -> None: """Save checkpoint metadata.""" metadata["last_updated"] = datetime.now().isoformat() with open(self.metadata_file, "w") as f: json.dump(metadata, f, indent=2) def save_checkpoint( self, model_name: str, model_path: str | Path, metrics: Dict[str, float], config: Dict[str, Any], epoch: int, ) -> None: """Save model checkpoint with metrics.""" metadata = self.get_metadata() # Create checkpoint entry checkpoint_key = f"{model_name}_epoch_{epoch}" metadata["checkpoints"][checkpoint_key] = { "model_path": str(model_path), "epoch": epoch, "metrics": metrics, "config": config, "timestamp": datetime.now().isoformat(), } self.save_metadata(metadata) def get_best_model(self, model_name: str, metric: str = "total_return") -> Optional[Dict[str, Any]]: """Get best checkpoint by metric.""" metadata = self.get_metadata() checkpoints = [cp for name, cp in metadata["checkpoints"].items() if model_name in name] if not checkpoints: return None best = max(checkpoints, key=lambda x: x["metrics"].get(metric, -np.inf)) return best def save_best_model_info(self, model_name: str, best_info: Dict[str, Any]) -> None: """Save best model information.""" best_models = {} if self.best_models_file.exists(): with open(self.best_models_file) as f: best_models = json.load(f) best_models[model_name] = { **best_info, "timestamp": datetime.now().isoformat(), } with open(self.best_models_file, "w") as f: json.dump(best_models, f, indent=2) # ============================================================ # TRAINING ORCHESTRATOR # ============================================================ @dataclass class TrainingResult: """Result of a single training run.""" config_name: str algorithm: str model_path: str metrics: Dict[str, float] total_timesteps: int training_time: float class RLTrainingOrchestrator: """Orchestrates training with checkpointing and best model tracking.""" def __init__( self, csv_path: str, output_dir: str = "./models/rl_training", results_dir: str = "./results/rl_training", checkpoint_dir: str = "./checkpoints/rl_training", augment: bool = True, augment_copies: int = 1, ): self.csv_path = csv_path self.output_dir = Path(output_dir) self.results_dir = Path(results_dir) self.checkpoint_dir = CheckpointManager(checkpoint_dir) # Setup directories self.output_dir.mkdir(parents=True, exist_ok=True) self.results_dir.mkdir(parents=True, exist_ok=True) # Load sessions once print("[info] Loading and preparing training data...") self.sessions = prepare_training_sessions( csv_path=csv_path, apply_augmentation=augment, copies_per_session=augment_copies, ) print(f"[ok] Loaded {len(self.sessions)} sessions") def train_single( self, config_name: str, algorithm: str, env_config: IntradayEnvConfig, train_config: PPOTrainingConfig | DQNTrainingConfig, resume_from: Optional[str] = None, ) -> TrainingResult: """Train a single model with automatic checkpoint resumption.""" print(f"\n{'-'*70}") print(f" Training: {config_name}") print(f" Algorithm: {algorithm.upper()}") print(f" Timesteps: {train_config.total_timesteps:,}") print(f"{'-'*70}") model_path = self.output_dir / f"{config_name}_{algorithm}" model_path.parent.mkdir(parents=True, exist_ok=True) # Setup checkpoint directory for this model model_checkpoint_dir = self.checkpoint_dir.checkpoint_dir / f"{config_name}_{algorithm}" start_time = datetime.now() try: if algorithm == "ppo": model = train_ppo( sessions=self.sessions, env_config=env_config, train_config=train_config, model_output_path=str(model_path), checkpoint_dir=str(model_checkpoint_dir), checkpoint_save_freq=50_000, ) elif algorithm == "dqn": model = train_dqn_baseline( sessions=self.sessions, env_config=env_config, train_config=train_config, model_output_path=str(model_path), checkpoint_dir=str(model_checkpoint_dir), checkpoint_save_freq=50_000, ) else: raise ValueError(f"Unknown algorithm: {algorithm}") elapsed = (datetime.now() - start_time).total_seconds() print(f"[ok] Training complete in {elapsed:.1f}s") print(f" Model saved -> {model_path}") # Get setup description for metrics setup = describe_training_setup(env_config, train_config) result = TrainingResult( config_name=config_name, algorithm=algorithm, model_path=str(model_path), metrics=setup, total_timesteps=train_config.total_timesteps, training_time=elapsed, ) # Save checkpoint self.checkpoint_dir.save_checkpoint( model_name=config_name, model_path=model_path, metrics=setup, config=vars(train_config), epoch=1, ) return result except Exception as e: print(f"[error] Training failed: {e}") import traceback traceback.print_exc() raise def backtest_model( self, model_path: str | Path, algorithm: str, env_config: IntradayEnvConfig, timeframe: str = "5min", ) -> Dict[str, Any]: """Backtest a trained model.""" resolved_model_path = resolve_model_path(model_path) print(f"\n Backtesting: {resolved_model_path.name}") # Load and resample data raw = load_ohlcv_csv(self.csv_path) featured = build_intraday_features(raw) transformed = resample_ohlcv(featured, timeframe=timeframe) sessions = split_sessions(transformed) # Load model if algorithm == "ppo": model = PPO.load(str(resolved_model_path)) elif algorithm == "dqn": model = DQN.load(str(resolved_model_path)) else: raise ValueError(f"Unknown algorithm: {algorithm}") # Run backtest env_config_copy = IntradayEnvConfig( lookback=env_config.lookback, initial_cash=env_config.initial_cash, lot_size=env_config.lot_size, transaction_cost_bps=env_config.transaction_cost_bps, trailing_stop_pct=env_config.trailing_stop_pct, trend_bonus=env_config.trend_bonus, volume_bonus=env_config.volume_bonus, invalid_action_penalty=env_config.invalid_action_penalty, volume_confirmation_multiplier=env_config.volume_confirmation_multiplier, morning_minutes=env_config.morning_minutes, random_reset=False, reward_scale=env_config.reward_scale, ) report = backtest_model(model=model, sessions=sessions, env_config=env_config_copy) return report # ============================================================ # METRICS CALCULATOR # ============================================================ def calculate_backtest_metrics(report: Dict[str, Any]) -> Dict[str, float]: """Extract and calculate performance metrics from backtest report.""" summary = report.get("summary", {}) session_results = report.get("session_results", []) if not session_results: return { "avg_return": 0.0, "total_return": 0.0, "win_rate": 0.0, "sharpe": 0.0, "max_drawdown": 0.0, "profit_factor": 0.0, } returns = [s.get("return_pct", 0.0) for s in session_results] wins = sum(1 for r in returns if r > 0) losses = sum(1 for r in returns if r < 0) # Calculate profit factor gross_profit = sum(r for r in returns if r > 0) gross_loss = abs(sum(r for r in returns if r < 0)) profit_factor = gross_profit / max(gross_loss, 1.0) # Calculate Sharpe (simple approximation) returns_array = np.array(returns) sharpe = (returns_array.mean() / (returns_array.std() + 1e-8)) * np.sqrt(252) return { "avg_return": float(np.mean(returns)), "total_return": float(np.sum(returns)), "win_rate": float(wins / len(returns) * 100 if returns else 0), "sharpe": float(sharpe), "max_drawdown": float(summary.get("max_drawdown_pct", 0.0)), "profit_factor": float(profit_factor), "num_trades": int(sum(s.get("trade_count", 0) for s in session_results)), } # ============================================================ # PARAMETER SWEEP # ============================================================ def run_parameter_sweep( csv_path: str, output_dir: str = "./models/rl_sweep", results_dir: str = "./results/rl_sweep", checkpoint_dir: str = "./checkpoints/rl_sweep", force_fresh: bool = False, ) -> List[Dict[str, Any]]: """ Sweep hyperparameters and track best models. Progress is checkpointed - can resume mid-sweep. """ print("\n" + "=" * 80) print("RL HYPERPARAMETER SWEEP") print("=" * 80) # Define sweep ranges learning_rates = [1e-4, 3e-4, 1e-3] net_architectures = [(128, 128), (256, 256), (512, 512)] algorithms = ["ppo"] # Can add "dqn" for comparison total_combos = len(learning_rates) * len(net_architectures) * len(algorithms) print(f"\nTotal combinations: {total_combos}") checkpoint_mgr = CheckpointManager(checkpoint_dir) results_file = Path(results_dir) / f"sweep_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" results_file.parent.mkdir(parents=True, exist_ok=True) # Load checkpoint if resuming checkpoint_path = Path(checkpoint_dir) / "sweep_checkpoint.json" completed_combos = [] results_list = [] if checkpoint_path.exists() and not force_fresh: print("[info] Resuming from checkpoint...") with open(checkpoint_path) as f: ckpt = json.load(f) completed_combos = ckpt.get("completed", []) results_list = ckpt.get("results", []) print(f" [ok] {len(completed_combos)} combinations already done") elif force_fresh and checkpoint_path.exists(): os.remove(checkpoint_path) print("[ok] Checkpoint cleared") orchestrator = RLTrainingOrchestrator( csv_path=csv_path, output_dir=output_dir, results_dir=results_dir, checkpoint_dir=checkpoint_dir, augment=True, augment_copies=1, ) print(f"\n{'-'*80}") try: combo_idx = 0 for algo in algorithms: for lr in learning_rates: for arch in net_architectures: combo_idx += 1 combo_key = f"{algo}_lr{lr}_arch{arch}" if combo_key in completed_combos: print(f" [{combo_idx}/{total_combos}] {combo_key} - SKIPPED (already done)") continue print(f" [{combo_idx}/{total_combos}] {combo_key}...") try: # Create configs env_config = IntradayEnvConfig(lookback=30, morning_minutes=60) train_config = PPOTrainingConfig( total_timesteps=100_000, learning_rate=lr, net_arch=arch, ) # Train train_result = orchestrator.train_single( config_name=combo_key, algorithm=algo, env_config=env_config, train_config=train_config, ) # Backtest backtest_report = orchestrator.backtest_model( model_path=train_result.model_path, algorithm=algo, env_config=env_config, timeframe="5min", ) metrics = calculate_backtest_metrics(backtest_report) result_entry = { "combo": combo_key, "algorithm": algo, "learning_rate": lr, "net_arch": arch, **metrics, "model_path": train_result.model_path, "training_time": train_result.training_time, } results_list.append(result_entry) completed_combos.append(combo_key) print(f" [ok] Return: {metrics['total_return']:+.2f}% | " f"Sharpe: {metrics['sharpe']:.2f} | Win: {metrics['win_rate']:.1f}%") # Save checkpoint with open(checkpoint_path, "w") as f: json.dump( {"completed": completed_combos, "results": results_list}, f, indent=2, ) except Exception as e: print(f" [error] Failed: {str(e)[:80]}") continue except KeyboardInterrupt: print(f"\n\nPaused. {len(completed_combos)}/{total_combos} done. Re-run to resume.\n") # Sort results by total return results_list.sort(key=lambda x: x["total_return"], reverse=True) # Print summary table print("\n" + "=" * 100) print("SWEEP RESULTS - RANKED BY TOTAL RETURN") print("=" * 100) print( f"{'Rank':<5} {'Algorithm':<8} {'Learning Rate':<15} {'Net Arch':<15} " f"{'Return %':<12} {'Win %':<10} {'Sharpe':<10} {'Max DD %':<10}" ) print("-" * 100) for i, r in enumerate(results_list[:10], 1): print( f"{i:<5} {r['algorithm']:<8} {r['learning_rate']:<15.0e} {str(r['net_arch']):<15} " f"{r['total_return']:<12.2f} {r['win_rate']:<10.1f} {r['sharpe']:<10.2f} " f"{abs(r['max_drawdown']):<10.2f}" ) print("-" * 100) # Save results with open(results_file, "w") as f: json.dump(results_list, f, indent=2) print(f"\n[ok] Results saved -> {results_file}") # Cleanup checkpoint if checkpoint_path.exists(): os.remove(checkpoint_path) print("[ok] Checkpoint cleaned up") return results_list # ============================================================ # CLI & MAIN # ============================================================ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Comprehensive RL Training + Backtesting + Checkpoint System" ) subparsers = parser.add_subparsers(dest="command", help="Command to run") # Train command train_parser = subparsers.add_parser("train", help="Train a single RL model") train_parser.add_argument("--csv", type=str, required=True, help="Path to OHLCV CSV") train_parser.add_argument("--name", type=str, default="ppo_intraday", help="Model name") train_parser.add_argument("--algo", type=str, default="ppo", choices=["ppo", "dqn"]) train_parser.add_argument("--timesteps", type=int, default=250_000) train_parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate") train_parser.add_argument("--net-arch", type=str, default="256,256", help="Network arch (comma-separated)") train_parser.add_argument( "--output-dir", type=str, default="./models/rl_training", help="Output directory" ) train_parser.add_argument( "--results-dir", type=str, default="./results/rl_training", help="Results directory" ) train_parser.add_argument( "--checkpoint-dir", type=str, default="./checkpoints/rl_training", help="Checkpoint directory", ) train_parser.add_argument("--no-augment", action="store_true", help="Disable data augmentation") train_parser.add_argument("--augment-copies", type=int, default=1) # Backtest command backtest_parser = subparsers.add_parser("backtest", help="Backtest trained model") backtest_parser.add_argument("--csv", type=str, required=True, help="Path to OHLCV CSV") backtest_parser.add_argument("--model", type=str, required=True, help="Path to model") backtest_parser.add_argument("--algo", type=str, default="ppo", choices=["ppo", "dqn"]) backtest_parser.add_argument("--timeframe", type=str, default="5min", choices=["1min", "5min", "15min"]) backtest_parser.add_argument( "--results-dir", type=str, default="./results/rl_backtest", help="Results directory" ) # Sweep command sweep_parser = subparsers.add_parser("sweep", help="Hyperparameter sweep") sweep_parser.add_argument("--csv", type=str, required=True, help="Path to OHLCV CSV") sweep_parser.add_argument( "--output-dir", type=str, default="./models/rl_sweep", help="Output directory" ) sweep_parser.add_argument( "--results-dir", type=str, default="./results/rl_sweep", help="Results directory" ) sweep_parser.add_argument( "--checkpoint-dir", type=str, default="./checkpoints/rl_sweep", help="Checkpoint directory", ) sweep_parser.add_argument("--force-fresh", action="store_true", help="Clear and restart sweep") return parser.parse_args() def main() -> None: args = parse_args() if not args.command: print("Usage: train_backtest_rl_system.py {train,backtest,sweep} [options]") return if args.command == "train": net_arch = tuple(map(int, args.net_arch.split(","))) orchestrator = RLTrainingOrchestrator( csv_path=args.csv, output_dir=args.output_dir, results_dir=args.results_dir, checkpoint_dir=args.checkpoint_dir, augment=not args.no_augment, augment_copies=args.augment_copies, ) env_config = IntradayEnvConfig(lookback=30, morning_minutes=60) train_config = PPOTrainingConfig( total_timesteps=args.timesteps, learning_rate=args.lr, net_arch=net_arch, ) if args.algo == "ppo" else DQNTrainingConfig( total_timesteps=args.timesteps, learning_rate=args.lr, net_arch=net_arch, ) result = orchestrator.train_single( config_name=args.name, algorithm=args.algo, env_config=env_config, train_config=train_config, ) print("\n[ok] Training complete") print(f" Model: {result.model_path}") print(f" Time: {result.training_time:.1f}s") elif args.command == "backtest": orchestrator = RLTrainingOrchestrator( csv_path=args.csv, output_dir="./models", results_dir=args.results_dir, ) env_config = IntradayEnvConfig( lookback=30, morning_minutes=60, trend_bonus=0.002, volume_bonus=0.002, invalid_action_penalty=0.0001, ) # Use 1min by default (same as training) unless specified timeframe = args.timeframe if args.timeframe != "5min" else "1min" report = orchestrator.backtest_model( model_path=args.model, algorithm=args.algo, env_config=env_config, timeframe=timeframe, ) metrics = calculate_backtest_metrics(report) print("\nBacktest Summary:") print(json.dumps(metrics, indent=2)) elif args.command == "sweep": results = run_parameter_sweep( csv_path=args.csv, output_dir=args.output_dir, results_dir=args.results_dir, checkpoint_dir=args.checkpoint_dir, force_fresh=args.force_fresh, ) print("\n[ok] Sweep complete. Best model:") print(f" {results[0]['combo']}") print(f" Return: {results[0]['total_return']:.2f}%") if __name__ == "__main__": main()