Spaces:
Sleeping
Sleeping
File size: 2,630 Bytes
ad77c79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | 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()
|