Spaces:
Running
Running
| import torch | |
| from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead | |
| from transformers import AutoTokenizer | |
| from server.openenv_wrapper import SupplyChainEnv | |
| def train_ppo(): | |
| print("Loading SFT Model for PPO Training...") | |
| # Load the SFT model we just trained | |
| # Note: On a machine without a GPU, this will fail or be extremely slow. | |
| model = AutoModelForCausalLMWithValueHead.from_pretrained("tradestorm-sft-model") | |
| tokenizer = AutoTokenizer.from_pretrained("tradestorm-sft-model") | |
| ppo_config = PPOConfig( | |
| batch_size=16, | |
| mini_batch_size=4, | |
| learning_rate=1e-5 | |
| ) | |
| print("Initializing PPO Trainer...") | |
| ppo_trainer = PPOTrainer(config=ppo_config, model=model, tokenizer=tokenizer) | |
| print("Initializing Supply Chain OpenEnv Wrapper...") | |
| env = SupplyChainEnv(max_turns=24) | |
| # Training Loop (500 episodes for demonstration) | |
| print("Starting PPO Training Loop...") | |
| for epoch in range(500): | |
| obs_json, info = env.reset() | |
| done = False | |
| episode_reward = 0 | |
| while not done: | |
| # 1. Format the Observation | |
| prompt = f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: " | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| # If using GPU, move tensors to CUDA | |
| if torch.cuda.is_available(): | |
| inputs = {k: v.to("cuda") for k, v in inputs.items()} | |
| # 2. LLM Generates Action | |
| action_tensors = ppo_trainer.generate(**inputs, max_new_tokens=100) | |
| # Decode the action | |
| action_str = tokenizer.decode(action_tensors[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True) | |
| # 3. Step the Environment | |
| obs_json, reward, done, truncated, info = env.step(action_str) | |
| episode_reward += reward | |
| # 4. PPO Weight Update | |
| reward_tensor = torch.tensor([reward], dtype=torch.float32) | |
| if torch.cuda.is_available(): | |
| reward_tensor = reward_tensor.to("cuda") | |
| ppo_trainer.step([inputs["input_ids"][0]], [action_tensors[0]], [reward_tensor]) | |
| print(f"Epoch {epoch} finished. Total Reward: {episode_reward} | Final Balance: {info.get('balance', 0)}") | |
| print("Saving Reinforcement Learning Model...") | |
| model.save_pretrained("tradestorm-rl-agent") | |
| print("✅ PPO Training complete. Model saved to tradestorm-rl-agent/") | |
| if __name__ == "__main__": | |
| train_ppo() | |