Buckets:
| import argparse | |
| import logging | |
| import os | |
| import sys | |
| import json | |
| from datetime import datetime | |
| from pathlib import Path | |
| from typing import Dict, Any, Optional | |
| sys.path.append(str(Path(__file__).parent)) | |
| from config import Config, PRESETS | |
| from environments.math_environment import MathEnvironment | |
| from agents.math_agent import MathReasoningAgent | |
| from algorithms.reward_model import ( | |
| HumanPreferenceSimulator, | |
| RewardModel, | |
| RewardModelTrainer, | |
| ) | |
| from algorithms.ppo_trainer import PPOTrainer | |
| def setup_logging(log_dir: str, level: str = "INFO") -> logging.Logger: | |
| os.makedirs(log_dir, exist_ok=True) | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| log_file = os.path.join(log_dir, f"rl_llm_training_{timestamp}.log") | |
| logging.basicConfig( | |
| level=getattr(logging, level.upper()), | |
| format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", | |
| handlers=[logging.FileHandler(log_file), logging.StreamHandler(sys.stdout)], | |
| ) | |
| logger = logging.getLogger(__name__) | |
| logger.info(f"Logging initialized. Log file: {log_file}") | |
| return logger | |
| def run_baseline_evaluation( | |
| agent: MathReasoningAgent, environment: MathEnvironment, num_episodes: int = 20 | |
| ) -> Dict[str, Any]: | |
| logger = logging.getLogger(__name__) | |
| logger.info(f"Running baseline evaluation with {num_episodes} episodes...") | |
| baseline_results = [] | |
| for episode in range(num_episodes): | |
| problem, correct_answer, metadata = environment.generate_problem("medium") | |
| generation_result = agent.generate_solution(problem) | |
| result = environment.verify_solution(problem, generation_result.text) | |
| baseline_results.append( | |
| { | |
| "episode": episode + 1, | |
| "problem": problem, | |
| "solution": generation_result.text, | |
| "reward": result.reward, | |
| "correct_answer": correct_answer, | |
| "feedback": result.feedback, | |
| } | |
| ) | |
| rewards = [r["reward"] for r in baseline_results] | |
| baseline_metrics = { | |
| "avg_reward": sum(rewards) / len(rewards), | |
| "success_rate": sum(1 for r in rewards if r > 0) / len(rewards), | |
| "min_reward": min(rewards), | |
| "max_reward": max(rewards), | |
| "std_reward": ( | |
| sum((r - sum(rewards) / len(rewards)) ** 2 for r in rewards) / len(rewards) | |
| ) | |
| ** 0.5, | |
| "num_episodes": num_episodes, | |
| "results": baseline_results, | |
| } | |
| logger.info(f"Baseline Results:") | |
| logger.info(f" Average Reward: {baseline_metrics['avg_reward']:.3f}") | |
| logger.info(f" Success Rate: {baseline_metrics['success_rate']:.1%}") | |
| logger.info( | |
| f" Reward Range: [{baseline_metrics['min_reward']:.3f}, {baseline_metrics['max_reward']:.3f}]" | |
| ) | |
| return baseline_metrics | |
| def run_rlvr_training(config: Config) -> Dict[str, Any]: | |
| logger = logging.getLogger(__name__) | |
| logger.info("🚀 Starting RLVR Training Pipeline") | |
| logger.info("Initializing RLVR components...") | |
| environment = MathEnvironment(config=config, use_sympy=True) | |
| agent = MathReasoningAgent(config.model.name, config, device=config.model.device) | |
| ppo_trainer = PPOTrainer(config, agent, environment, device=config.model.device) | |
| logger.info("📊 Running baseline evaluation...") | |
| baseline_metrics = run_baseline_evaluation( | |
| agent, environment, config.experiment.eval_episodes | |
| ) | |
| logger.info("🎯 Starting PPO training...") | |
| training_history = ppo_trainer.train() | |
| logger.info("📈 Running final evaluation...") | |
| final_metrics = ppo_trainer.evaluate(config.experiment.eval_episodes) | |
| results = { | |
| "experiment_type": "RLVR", | |
| "config": { | |
| "model_name": config.model.name, | |
| "num_epochs": config.rl.num_epochs, | |
| "batch_size": config.rl.batch_size, | |
| "ppo_clip": config.rl.ppo_clip, | |
| }, | |
| "baseline_metrics": baseline_metrics, | |
| "final_metrics": final_metrics, | |
| "training_summary": ppo_trainer.get_training_summary(), | |
| "training_history": [epoch.__dict__ for epoch in training_history], | |
| "improvement": { | |
| "reward_improvement": final_metrics["eval_avg_reward"] | |
| - baseline_metrics["avg_reward"], | |
| "success_rate_improvement": final_metrics["eval_success_rate"] | |
| - baseline_metrics["success_rate"], | |
| }, | |
| "timestamp": datetime.now().isoformat(), | |
| } | |
| results_path = os.path.join(config.experiment.output_dir, "rlvr_results.json") | |
| with open(results_path, "w") as f: | |
| json.dump(results, f, indent=2) | |
| plot_path = os.path.join(config.experiment.output_dir, "rlvr_training_curves.png") | |
| ppo_trainer.plot_training_curves(plot_path) | |
| export_path = os.path.join(config.experiment.output_dir, "rlvr_training_data.json") | |
| ppo_trainer.export_training_data(export_path) | |
| logger.info(f"✅ RLVR training completed!") | |
| logger.info(f"Results saved to: {results_path}") | |
| logger.info(f"Training curves saved to: {plot_path}") | |
| return results | |
| def run_rlhf_training(config: Config) -> Dict[str, Any]: | |
| logger = logging.getLogger(__name__) | |
| logger.info("🤝 Starting RLHF Training Pipeline") | |
| logger.info("Initializing RLHF components...") | |
| environment = MathEnvironment(config=config, use_sympy=True) | |
| agent = MathReasoningAgent(config.model.name, config, device=config.model.device) | |
| human_simulator = HumanPreferenceSimulator(config) | |
| reward_model = RewardModel(agent.tokenizer, hidden_dim=256) | |
| reward_trainer = RewardModelTrainer(reward_model, learning_rate=3e-5) | |
| logger.info("Generating preference comparison data...") | |
| preference_data = [] | |
| for i in range(100): | |
| problem, _, _ = environment.generate_problem("medium") | |
| response1 = agent.generate_solution(problem, temperature=0.7).text | |
| response2 = agent.generate_solution(problem, temperature=1.2).text | |
| preference = human_simulator.compare_responses(problem, response1, response2) | |
| preference_data.append(preference) | |
| if i % 20 == 0: | |
| logger.info(f"Generated {i + 1}/100 preference comparisons") | |
| logger.info("Training reward model on preference data...") | |
| reward_training_history = [] | |
| for epoch in range(5): | |
| epoch_metrics = reward_trainer.train_epoch(preference_data, batch_size=8) | |
| reward_training_history.append(epoch_metrics) | |
| logger.info( | |
| f"Reward model epoch {epoch + 1}: " | |
| f"loss={epoch_metrics['loss']:.4f}, " | |
| f"accuracy={epoch_metrics['accuracy']:.4f}" | |
| ) | |
| ppo_trainer = PPOTrainer(config, agent, environment, device=config.model.device) | |
| logger.info("📊 Running baseline evaluation...") | |
| baseline_metrics = run_baseline_evaluation( | |
| agent, environment, config.experiment.eval_episodes | |
| ) | |
| logger.info("🎯 Starting PPO training with learned reward model...") | |
| training_history = ppo_trainer.train() | |
| logger.info("📈 Running final evaluation...") | |
| final_metrics = ppo_trainer.evaluate(config.experiment.eval_episodes) | |
| results = { | |
| "experiment_type": "RLHF", | |
| "config": { | |
| "model_name": config.model.name, | |
| "num_epochs": config.rl.num_epochs, | |
| "batch_size": config.rl.batch_size, | |
| "preference_data_size": len(preference_data), | |
| }, | |
| "baseline_metrics": baseline_metrics, | |
| "final_metrics": final_metrics, | |
| "reward_model_training": reward_training_history, | |
| "training_summary": ppo_trainer.get_training_summary(), | |
| "training_history": [epoch.__dict__ for epoch in training_history], | |
| "improvement": { | |
| "reward_improvement": final_metrics["eval_avg_reward"] | |
| - baseline_metrics["avg_reward"], | |
| "success_rate_improvement": final_metrics["eval_success_rate"] | |
| - baseline_metrics["success_rate"], | |
| }, | |
| "timestamp": datetime.now().isoformat(), | |
| } | |
| results_path = os.path.join(config.experiment.output_dir, "rlhf_results.json") | |
| with open(results_path, "w") as f: | |
| json.dump(results, f, indent=2) | |
| plot_path = os.path.join(config.experiment.output_dir, "rlhf_training_curves.png") | |
| ppo_trainer.plot_training_curves(plot_path) | |
| logger.info(f"✅ RLHF training completed!") | |
| logger.info(f"Results saved to: {results_path}") | |
| logger.info(f"Training curves saved to: {plot_path}") | |
| return results | |
| def run_comparison_experiment(config: Config) -> Dict[str, Any]: | |
| logger = logging.getLogger(__name__) | |
| logger.info("🔬 Starting Comparison Experiment") | |
| logger.info("Running RLVR experiment...") | |
| rlvr_results = run_rlvr_training(config) | |
| logger.info("Running RLHF experiment...") | |
| rlhf_results = run_rlhf_training(config) | |
| comparison = { | |
| "experiment_type": "Comparison", | |
| "rlvr_results": rlvr_results, | |
| "rlhf_results": rlhf_results, | |
| "comparison_metrics": { | |
| "rlvr_final_reward": rlvr_results["final_metrics"]["eval_avg_reward"], | |
| "rlhf_final_reward": rlhf_results["final_metrics"]["eval_avg_reward"], | |
| "rlvr_improvement": rlvr_results["improvement"]["reward_improvement"], | |
| "rlhf_improvement": rlhf_results["improvement"]["reward_improvement"], | |
| "winner": ( | |
| "RLVR" | |
| if rlvr_results["final_metrics"]["eval_avg_reward"] | |
| > rlhf_results["final_metrics"]["eval_avg_reward"] | |
| else "RLHF" | |
| ), | |
| }, | |
| "timestamp": datetime.now().isoformat(), | |
| } | |
| comparison_path = os.path.join( | |
| config.experiment.output_dir, "comparison_results.json" | |
| ) | |
| with open(comparison_path, "w") as f: | |
| json.dump(comparison, f, indent=2) | |
| logger.info(f"🔍 Comparison Results:") | |
| logger.info( | |
| f" RLVR Final Reward: {rlvr_results['final_metrics']['eval_avg_reward']:.3f}" | |
| ) | |
| logger.info( | |
| f" RLHF Final Reward: {rlhf_results['final_metrics']['eval_avg_reward']:.3f}" | |
| ) | |
| logger.info(f" Winner: {comparison['comparison_metrics']['winner']}") | |
| logger.info(f" Results saved to: {comparison_path}") | |
| return comparison | |
| def main(): | |
| parser = argparse.ArgumentParser(description="RL-LLM Training Pipeline") | |
| parser.add_argument( | |
| "--experiment", | |
| type=str, | |
| choices=["rlvr", "rlhf", "comparison"], | |
| default="rlvr", | |
| help="Type of experiment to run", | |
| ) | |
| parser.add_argument("--config", type=str, help="Path to config YAML file") | |
| parser.add_argument( | |
| "--preset", | |
| type=str, | |
| choices=list(PRESETS.keys()), | |
| help="Use preset configuration", | |
| ) | |
| parser.add_argument( | |
| "--output-dir", | |
| type=str, | |
| default="./results", | |
| help="Output directory for results", | |
| ) | |
| parser.add_argument( | |
| "--log-level", | |
| type=str, | |
| default="INFO", | |
| choices=["DEBUG", "INFO", "WARNING", "ERROR"], | |
| help="Logging level", | |
| ) | |
| parser.add_argument("--epochs", type=int, help="Number of training epochs") | |
| parser.add_argument("--batch-size", type=int, help="Batch size for training") | |
| parser.add_argument("--model", type=str, help="Model name to use") | |
| args = parser.parse_args() | |
| logger = setup_logging(args.output_dir, args.log_level) | |
| try: | |
| if args.config: | |
| config = Config.from_yaml(args.config) | |
| logger.info(f"Loaded configuration from {args.config}") | |
| elif args.preset: | |
| config = PRESETS[args.preset] | |
| logger.info(f"Using preset configuration: {args.preset}") | |
| else: | |
| config = Config() | |
| logger.info("Using default configuration") | |
| if args.output_dir: | |
| config.experiment.output_dir = args.output_dir | |
| config.experiment.log_dir = os.path.join(args.output_dir, "logs") | |
| config.experiment.checkpoint_dir = os.path.join(args.output_dir, "models") | |
| if args.epochs: | |
| config.rl.num_epochs = args.epochs | |
| if args.batch_size: | |
| config.rl.batch_size = args.batch_size | |
| if args.model: | |
| config.model.name = args.model | |
| os.makedirs(config.experiment.output_dir, exist_ok=True) | |
| os.makedirs(config.experiment.log_dir, exist_ok=True) | |
| os.makedirs(config.experiment.checkpoint_dir, exist_ok=True) | |
| config_path = os.path.join(config.experiment.output_dir, "config.yaml") | |
| config.to_yaml(config_path) | |
| logger.info(f"Configuration saved to {config_path}") | |
| logger.info(f"🎭 Starting {args.experiment.upper()} experiment") | |
| logger.info(f"Configuration: {config}") | |
| if args.experiment == "rlvr": | |
| results = run_rlvr_training(config) | |
| elif args.experiment == "rlhf": | |
| results = run_rlhf_training(config) | |
| elif args.experiment == "comparison": | |
| results = run_comparison_experiment(config) | |
| logger.info("🎉 Experiment completed successfully!") | |
| if args.experiment in ["rlvr", "rlhf"]: | |
| logger.info(f"Final Results Summary:") | |
| logger.info( | |
| f" Average Reward: {results['final_metrics']['eval_avg_reward']:.3f}" | |
| ) | |
| logger.info( | |
| f" Success Rate: {results['final_metrics']['eval_success_rate']:.1%}" | |
| ) | |
| logger.info( | |
| f" Improvement: {results['improvement']['reward_improvement']:+.3f}" | |
| ) | |
| return 0 | |
| except Exception as e: | |
| logger.error(f"Experiment failed with error: {e}") | |
| logger.exception("Full traceback:") | |
| return 1 | |
| if __name__ == "__main__": | |
| exit(main()) | |
Xet Storage Details
- Size:
- 13.9 kB
- Xet hash:
- a4c7489d23f8b76f52010fe3b5902e2393efb5846bd4b976a0a65d845990a772
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.