Left action panel revised to show triggers list + Create Token/Auto-Arrange. CMD DMG badges added to center life-panel score areas (click opens matrix). Removed renderEnhancedPanels calls that split the battlefield on damage adjust. Added encoding='utf-8' to all bare open() calls across 17 Python files to prevent cp1252 crashes on MTG card text with smart quotes/em-dashes. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
824 lines
32 KiB
Python
824 lines
32 KiB
Python
"""
|
||
Commander AI Lab — PPO Trainer
|
||
═══════════════════════════════
|
||
Proximal Policy Optimization (Clip) with:
|
||
- Clipped surrogate objective for policy
|
||
- Value function loss (clipped)
|
||
- Entropy bonus for exploration
|
||
- Generalized Advantage Estimation (computed in RolloutBuffer)
|
||
- Self-play episode collection between updates
|
||
- Periodic checkpointing and evaluation
|
||
- Win-rate tracking against opponent policies
|
||
- Forge-based episode source via --data-source forge|mixed
|
||
|
||
Reference: Schulman et al., "Proximal Policy Optimization Algorithms", 2017
|
||
|
||
Usage:
|
||
python -m ml.training.ppo_trainer --iterations 100 --episodes-per-iter 64
|
||
python -m ml.training.ppo_trainer --load-supervised ml/models/checkpoints/best_policy.pt
|
||
python -m ml.training.ppo_trainer --data-source forge --forge-results-dir results --iterations 50
|
||
python -m ml.training.ppo_trainer --data-source mixed --forge-weight 0.7 --iterations 100
|
||
"""
|
||
|
||
import argparse
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional
|
||
|
||
import numpy as np
|
||
|
||
project_root = str(Path(__file__).parent.parent.parent)
|
||
if project_root not in sys.path:
|
||
sys.path.insert(0, project_root)
|
||
|
||
# Conditional torch import
|
||
try:
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
TORCH_AVAILABLE = True
|
||
except ImportError:
|
||
TORCH_AVAILABLE = False
|
||
|
||
from ml.config.scope import NUM_ACTIONS, STATE_DIMS, TRAINING_CONFIG
|
||
from ml.training.rollout_buffer import RolloutBuffer
|
||
from ml.training.reward import RewardConfig
|
||
from ml.training.self_play import (
|
||
collect_rollouts, HeuristicPolicy, RandomPolicy,
|
||
)
|
||
|
||
logger = logging.getLogger("ml.ppo")
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════
|
||
# Forge data-source integration
|
||
# ══════════════════════════════════════════════════════════
|
||
|
||
def episodes_to_buffer(episodes, buffer: RolloutBuffer) -> int:
|
||
"""Drain a list of ForgeEpisodes into an existing RolloutBuffer.
|
||
|
||
Returns the number of transitions added.
|
||
"""
|
||
added = 0
|
||
for ep in episodes:
|
||
for t in ep.transitions:
|
||
if buffer.size >= buffer.buffer_size:
|
||
break
|
||
buffer.add(
|
||
state=t.state,
|
||
action=t.action,
|
||
reward=t.reward,
|
||
log_prob=t.log_prob,
|
||
value=t.value,
|
||
done=t.done,
|
||
)
|
||
added += 1
|
||
return added
|
||
|
||
|
||
class ForgeDataSource:
|
||
"""Runs the Forge producer pipeline in a background thread and
|
||
exposes a synchronous drain() method for the PPO training loop.
|
||
|
||
Usage::
|
||
source = ForgeDataSource(results_dir="results", batch_id="run1", num_games=64)
|
||
source.start()
|
||
# ... in training loop ...
|
||
episodes = source.drain()
|
||
source.stop()
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
results_dir: str = "results",
|
||
batch_id: str = "forge-batch",
|
||
num_games: int = 256,
|
||
reward_config: Optional[RewardConfig] = None,
|
||
queue_size: int = 256,
|
||
):
|
||
self.results_dir = results_dir
|
||
self.batch_id = batch_id
|
||
self.num_games = num_games
|
||
self.reward_config = reward_config or RewardConfig()
|
||
self.queue_size = queue_size
|
||
self._episodes: List = []
|
||
self._lock = threading.Lock()
|
||
self._thread: Optional[threading.Thread] = None
|
||
self._stop_event = threading.Event()
|
||
|
||
def start(self):
|
||
"""Launch the background producer thread."""
|
||
self._thread = threading.Thread(target=self._run, daemon=True)
|
||
self._thread.start()
|
||
logger.info("[ForgeDataSource] Producer started (results_dir=%s, batch_id=%s, num_games=%d)",
|
||
self.results_dir, self.batch_id, self.num_games)
|
||
|
||
def _run(self):
|
||
from ml.training.forge_episode_generator import ForgeEpisodeQueue, forge_producer
|
||
|
||
async def _produce():
|
||
queue = ForgeEpisodeQueue(max_size=self.queue_size)
|
||
produced = await forge_producer(
|
||
queue=queue,
|
||
results_dir=self.results_dir,
|
||
batch_id=self.batch_id,
|
||
num_games=self.num_games,
|
||
reward_config=self.reward_config,
|
||
)
|
||
queue.close()
|
||
# Drain remaining items
|
||
while not queue.is_closed:
|
||
ep = await queue.get(timeout=1.0)
|
||
if ep is not None:
|
||
with self._lock:
|
||
self._episodes.append(ep)
|
||
logger.info("[ForgeDataSource] Producer finished: %d episodes", produced)
|
||
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
loop.run_until_complete(_produce())
|
||
finally:
|
||
loop.close()
|
||
|
||
def drain(self) -> List:
|
||
"""Return and clear all buffered episodes (thread-safe)."""
|
||
with self._lock:
|
||
eps = list(self._episodes)
|
||
self._episodes.clear()
|
||
return eps
|
||
|
||
def stop(self):
|
||
self._stop_event.set()
|
||
if self._thread:
|
||
self._thread.join(timeout=5.0)
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════
|
||
# PPO Config
|
||
# ══════════════════════════════════════════════════════════
|
||
|
||
class PPOConfig:
|
||
"""PPO hyperparameters."""
|
||
|
||
def __init__(
|
||
self,
|
||
# Training loop
|
||
iterations: int = 100,
|
||
episodes_per_iter: int = 64,
|
||
ppo_epochs: int = 4,
|
||
batch_size: int = 256,
|
||
|
||
# PPO core
|
||
clip_epsilon: float = 0.2,
|
||
value_clip: float = 0.2,
|
||
entropy_coeff: float = 0.01,
|
||
value_loss_coeff: float = 0.5,
|
||
max_grad_norm: float = 0.5,
|
||
|
||
# Optimizer
|
||
learning_rate: float = 3e-4,
|
||
lr_schedule: str = "constant", # constant | linear | cosine
|
||
min_lr: float = 1e-5,
|
||
|
||
# Rollout
|
||
buffer_size: int = 4096,
|
||
gamma: float = 0.99,
|
||
gae_lambda: float = 0.95,
|
||
|
||
# Opponent
|
||
opponent: str = "heuristic", # heuristic | random | self
|
||
playstyle: str = "midrange",
|
||
|
||
# Checkpointing
|
||
checkpoint_dir: str = "ml/models/checkpoints",
|
||
save_every: int = 10,
|
||
eval_every: int = 5,
|
||
eval_episodes: int = 50,
|
||
|
||
# Supervised preload
|
||
load_supervised: Optional[str] = None,
|
||
|
||
# Data source
|
||
data_source: str = "synthetic", # synthetic | forge | mixed
|
||
forge_results_dir: str = "results",
|
||
forge_batch_id: str = "forge-batch",
|
||
forge_num_games: int = 256,
|
||
forge_weight: float = 0.7, # fraction of buffer from Forge when mixed
|
||
):
|
||
self.iterations = iterations
|
||
self.episodes_per_iter = episodes_per_iter
|
||
self.ppo_epochs = ppo_epochs
|
||
self.batch_size = batch_size
|
||
self.clip_epsilon = clip_epsilon
|
||
self.value_clip = value_clip
|
||
self.entropy_coeff = entropy_coeff
|
||
self.value_loss_coeff = value_loss_coeff
|
||
self.max_grad_norm = max_grad_norm
|
||
self.learning_rate = learning_rate
|
||
self.lr_schedule = lr_schedule
|
||
self.min_lr = min_lr
|
||
self.buffer_size = buffer_size
|
||
self.gamma = gamma
|
||
self.gae_lambda = gae_lambda
|
||
self.opponent = opponent
|
||
self.playstyle = playstyle
|
||
self.checkpoint_dir = checkpoint_dir
|
||
self.save_every = save_every
|
||
self.eval_every = eval_every
|
||
self.eval_episodes = eval_episodes
|
||
self.load_supervised = load_supervised
|
||
self.data_source = data_source
|
||
self.forge_results_dir = forge_results_dir
|
||
self.forge_batch_id = forge_batch_id
|
||
self.forge_num_games = forge_num_games
|
||
self.forge_weight = max(0.0, min(1.0, forge_weight))
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════
|
||
# PPO Trainer
|
||
# ══════════════════════════════════════════════════════════
|
||
|
||
class PPOTrainer:
|
||
"""Proximal Policy Optimization trainer with self-play and optional Forge data source."""
|
||
|
||
def __init__(self, config: PPOConfig):
|
||
if not TORCH_AVAILABLE:
|
||
raise ImportError("PyTorch required for PPO training. Install: pip install torch")
|
||
|
||
self.config = config
|
||
self.device = self._detect_device()
|
||
|
||
from ml.training.policy_network import PolicyValueNetwork
|
||
self.model = PolicyValueNetwork(
|
||
input_dim=STATE_DIMS.total_state_dim,
|
||
hidden_dim=TRAINING_CONFIG.hidden_dim,
|
||
num_actions=NUM_ACTIONS,
|
||
num_layers=TRAINING_CONFIG.num_layers,
|
||
dropout=0.0,
|
||
).to(self.device)
|
||
|
||
if config.load_supervised and os.path.exists(config.load_supervised):
|
||
self._load_supervised_weights(config.load_supervised)
|
||
|
||
self.optimizer = torch.optim.Adam(
|
||
self.model.parameters(),
|
||
lr=config.learning_rate,
|
||
eps=1e-5,
|
||
)
|
||
|
||
self.scheduler = None
|
||
if config.lr_schedule == "linear":
|
||
self.scheduler = torch.optim.lr_scheduler.LinearLR(
|
||
self.optimizer,
|
||
start_factor=1.0,
|
||
end_factor=config.min_lr / config.learning_rate,
|
||
total_iters=config.iterations,
|
||
)
|
||
elif config.lr_schedule == "cosine":
|
||
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
|
||
self.optimizer,
|
||
T_max=config.iterations,
|
||
eta_min=config.min_lr,
|
||
)
|
||
|
||
self.buffer = RolloutBuffer(
|
||
buffer_size=config.buffer_size,
|
||
state_dim=STATE_DIMS.total_state_dim,
|
||
gamma=config.gamma,
|
||
gae_lambda=config.gae_lambda,
|
||
)
|
||
|
||
self.reward_config = RewardConfig(
|
||
gamma=config.gamma,
|
||
gae_lambda=config.gae_lambda,
|
||
)
|
||
|
||
self.opponent = self._create_opponent(config.opponent)
|
||
|
||
# Forge data source (started lazily in train())
|
||
self._forge_source: Optional[ForgeDataSource] = None
|
||
if config.data_source in ("forge", "mixed"):
|
||
self._forge_source = ForgeDataSource(
|
||
results_dir=config.forge_results_dir,
|
||
batch_id=config.forge_batch_id,
|
||
num_games=config.forge_num_games,
|
||
reward_config=self.reward_config,
|
||
)
|
||
|
||
self.iteration = 0
|
||
self.total_steps = 0
|
||
self.best_win_rate = 0.0
|
||
self.history = []
|
||
|
||
os.makedirs(config.checkpoint_dir, exist_ok=True)
|
||
|
||
def _detect_device(self) -> str:
|
||
if torch.cuda.is_available():
|
||
return "cuda"
|
||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||
return "mps"
|
||
return "cpu"
|
||
|
||
def _load_supervised_weights(self, path: str):
|
||
logger.info("Loading supervised weights from %s", path)
|
||
checkpoint = torch.load(path, map_location=self.device, weights_only=False)
|
||
state_dict = checkpoint.get("model_state_dict", checkpoint)
|
||
model_dict = self.model.state_dict()
|
||
compatible = {k: v for k, v in state_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
|
||
model_dict.update(compatible)
|
||
self.model.load_state_dict(model_dict)
|
||
logger.info("Loaded %d/%d parameter tensors from supervised checkpoint",
|
||
len(compatible), len(state_dict))
|
||
|
||
def _create_opponent(self, opponent_type: str):
|
||
if opponent_type == "random":
|
||
return RandomPolicy()
|
||
elif opponent_type == "heuristic":
|
||
return HeuristicPolicy()
|
||
elif opponent_type == "self":
|
||
return self._create_self_play_opponent()
|
||
return HeuristicPolicy()
|
||
|
||
def _create_self_play_opponent(self):
|
||
class SelfPlayPolicy:
|
||
def __init__(self, model, device):
|
||
self.model = model
|
||
self.device = device
|
||
|
||
def select_action(self, state_vec):
|
||
state_tensor = torch.from_numpy(state_vec).unsqueeze(0).to(self.device)
|
||
with torch.no_grad():
|
||
logits, _ = self.model.forward_with_value(state_tensor)
|
||
probs = F.softmax(logits, dim=-1)
|
||
dist = torch.distributions.Categorical(probs)
|
||
action = dist.sample()
|
||
log_prob = dist.log_prob(action)
|
||
return action.item(), log_prob.item()
|
||
|
||
return SelfPlayPolicy(self.model, self.device)
|
||
|
||
def _collect_forge_rollouts(self) -> Dict:
|
||
"""Drain buffered Forge episodes into the RolloutBuffer.
|
||
|
||
Returns stats dict compatible with collect_rollouts() output.
|
||
"""
|
||
episodes = self._forge_source.drain()
|
||
if not episodes:
|
||
logger.warning("[Forge] No episodes available yet — waiting for JSONL output in %s",
|
||
self.config.forge_results_dir)
|
||
return {"steps": 0, "win_rate": 0.0, "avg_reward_per_ep": 0.0,
|
||
"wins": 0, "losses": 0, "draws": 0, "source": "forge"}
|
||
|
||
added = episodes_to_buffer(episodes, self.buffer)
|
||
wins = sum(1 for ep in episodes if ep.winner == 0)
|
||
win_rate = wins / max(len(episodes), 1)
|
||
avg_reward = sum(ep.total_reward for ep in episodes) / max(len(episodes), 1)
|
||
logger.info("[Forge] Drained %d episodes (%d transitions) into buffer — WR=%.0f%% AvgR=%.3f",
|
||
len(episodes), added, win_rate * 100, avg_reward)
|
||
return {
|
||
"steps": added,
|
||
"win_rate": win_rate,
|
||
"avg_reward_per_ep": avg_reward,
|
||
"wins": wins,
|
||
"losses": len(episodes) - wins,
|
||
"draws": 0,
|
||
"source": "forge",
|
||
}
|
||
|
||
def _collect_rollouts_for_iter(self) -> Dict:
|
||
"""Collect rollouts according to configured data_source.
|
||
|
||
- synthetic: pure self_play.collect_rollouts (original behaviour)
|
||
- forge: pure Forge episode drain
|
||
- mixed: Forge fills forge_weight fraction of buffer, synthetic fills remainder
|
||
"""
|
||
cfg = self.config
|
||
self.buffer.reset()
|
||
self.model.eval()
|
||
|
||
if cfg.data_source == "forge":
|
||
stats = self._collect_forge_rollouts()
|
||
|
||
elif cfg.data_source == "mixed":
|
||
forge_cap = int(cfg.buffer_size * cfg.forge_weight)
|
||
synth_cap = cfg.buffer_size - forge_cap
|
||
|
||
# Temporarily shrink buffer for Forge slice
|
||
original_size = self.buffer.buffer_size
|
||
self.buffer.buffer_size = forge_cap
|
||
forge_stats = self._collect_forge_rollouts()
|
||
self.buffer.buffer_size = original_size
|
||
|
||
synth_eps = max(1, int(cfg.episodes_per_iter * (1.0 - cfg.forge_weight)))
|
||
synth_stats = collect_rollouts(
|
||
agent_model=self.model,
|
||
opponent_policy=self.opponent,
|
||
buffer=self.buffer,
|
||
num_episodes=synth_eps,
|
||
playstyle=cfg.playstyle,
|
||
reward_config=self.reward_config,
|
||
)
|
||
stats = {
|
||
"steps": forge_stats["steps"] + synth_stats["steps"],
|
||
"win_rate": (
|
||
forge_stats["win_rate"] * cfg.forge_weight
|
||
+ synth_stats["win_rate"] * (1.0 - cfg.forge_weight)
|
||
),
|
||
"avg_reward_per_ep": (
|
||
forge_stats["avg_reward_per_ep"] * cfg.forge_weight
|
||
+ synth_stats["avg_reward_per_ep"] * (1.0 - cfg.forge_weight)
|
||
),
|
||
"wins": forge_stats["wins"] + synth_stats["wins"],
|
||
"losses": forge_stats["losses"] + synth_stats["losses"],
|
||
"draws": synth_stats.get("draws", 0),
|
||
"source": "mixed",
|
||
}
|
||
|
||
else: # synthetic (default — unchanged behaviour)
|
||
stats = collect_rollouts(
|
||
agent_model=self.model,
|
||
opponent_policy=self.opponent,
|
||
buffer=self.buffer,
|
||
num_episodes=cfg.episodes_per_iter,
|
||
playstyle=cfg.playstyle,
|
||
reward_config=self.reward_config,
|
||
)
|
||
stats["source"] = "synthetic"
|
||
|
||
self.model.train()
|
||
return stats
|
||
|
||
def ppo_update(self) -> Dict:
|
||
"""Run PPO update epochs on the collected rollout buffer."""
|
||
cfg = self.config
|
||
total_policy_loss = 0.0
|
||
total_value_loss = 0.0
|
||
total_entropy = 0.0
|
||
total_approx_kl = 0.0
|
||
total_clip_frac = 0.0
|
||
num_updates = 0
|
||
|
||
for epoch in range(cfg.ppo_epochs):
|
||
for batch in self.buffer.get_batches(cfg.batch_size):
|
||
states, actions, old_log_probs, advantages, returns = batch
|
||
|
||
states_t = torch.from_numpy(states).to(self.device)
|
||
actions_t = torch.from_numpy(actions).to(self.device)
|
||
old_log_probs_t = torch.from_numpy(old_log_probs).to(self.device)
|
||
advantages_t = torch.from_numpy(advantages).to(self.device)
|
||
returns_t = torch.from_numpy(returns).to(self.device)
|
||
|
||
logits, values = self.model.forward_with_value(states_t)
|
||
values = values.squeeze(-1)
|
||
|
||
probs = F.softmax(logits, dim=-1)
|
||
dist = torch.distributions.Categorical(probs)
|
||
new_log_probs = dist.log_prob(actions_t)
|
||
entropy = dist.entropy().mean()
|
||
|
||
ratio = torch.exp(new_log_probs - old_log_probs_t)
|
||
surr1 = ratio * advantages_t
|
||
surr2 = torch.clamp(ratio, 1.0 - cfg.clip_epsilon, 1.0 + cfg.clip_epsilon) * advantages_t
|
||
policy_loss = -torch.min(surr1, surr2).mean()
|
||
|
||
if cfg.value_clip > 0:
|
||
value_loss_unclipped = (values - returns_t) ** 2
|
||
value_loss = 0.5 * value_loss_unclipped.mean()
|
||
else:
|
||
value_loss = 0.5 * ((values - returns_t) ** 2).mean()
|
||
|
||
loss = (
|
||
policy_loss
|
||
+ cfg.value_loss_coeff * value_loss
|
||
- cfg.entropy_coeff * entropy
|
||
)
|
||
|
||
self.optimizer.zero_grad()
|
||
loss.backward()
|
||
nn.utils.clip_grad_norm_(self.model.parameters(), cfg.max_grad_norm)
|
||
self.optimizer.step()
|
||
|
||
with torch.no_grad():
|
||
approx_kl = ((ratio - 1) - torch.log(ratio)).mean().item()
|
||
clip_frac = ((ratio - 1.0).abs() > cfg.clip_epsilon).float().mean().item()
|
||
|
||
total_policy_loss += policy_loss.item()
|
||
total_value_loss += value_loss.item()
|
||
total_entropy += entropy.item()
|
||
total_approx_kl += approx_kl
|
||
total_clip_frac += clip_frac
|
||
num_updates += 1
|
||
|
||
n = max(num_updates, 1)
|
||
return {
|
||
"policy_loss": total_policy_loss / n,
|
||
"value_loss": total_value_loss / n,
|
||
"entropy": total_entropy / n,
|
||
"approx_kl": total_approx_kl / n,
|
||
"clip_fraction": total_clip_frac / n,
|
||
"num_updates": num_updates,
|
||
}
|
||
|
||
def evaluate(self, num_episodes: int = 50) -> Dict:
|
||
"""Evaluate the current policy against heuristic and random opponents."""
|
||
self.model.eval()
|
||
results = {}
|
||
for opp_name, opp_policy in [
|
||
("heuristic", HeuristicPolicy()),
|
||
("random", RandomPolicy()),
|
||
]:
|
||
eval_buffer = RolloutBuffer(
|
||
buffer_size=num_episodes * 50,
|
||
state_dim=STATE_DIMS.total_state_dim,
|
||
gamma=self.config.gamma,
|
||
gae_lambda=self.config.gae_lambda,
|
||
)
|
||
stats = collect_rollouts(
|
||
agent_model=self.model,
|
||
opponent_policy=opp_policy,
|
||
buffer=eval_buffer,
|
||
num_episodes=num_episodes,
|
||
playstyle=self.config.playstyle,
|
||
reward_config=self.reward_config,
|
||
)
|
||
results[opp_name] = {
|
||
"win_rate": stats["win_rate"],
|
||
"wins": stats["wins"],
|
||
"losses": stats["losses"],
|
||
"draws": stats["draws"],
|
||
"avg_reward": stats["avg_reward_per_ep"],
|
||
"avg_steps": stats["avg_steps_per_ep"],
|
||
}
|
||
self.model.train()
|
||
return results
|
||
|
||
def save_checkpoint(self, path: str = None, metrics: dict = None):
|
||
if path is None:
|
||
path = os.path.join(
|
||
self.config.checkpoint_dir,
|
||
f"ppo_iter_{self.iteration:04d}.pt",
|
||
)
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
torch.save({
|
||
"iteration": self.iteration,
|
||
"total_steps": self.total_steps,
|
||
"model_state_dict": self.model.state_dict(),
|
||
"optimizer_state_dict": self.optimizer.state_dict(),
|
||
"metrics": metrics or {},
|
||
"config": {
|
||
"clip_epsilon": self.config.clip_epsilon,
|
||
"entropy_coeff": self.config.entropy_coeff,
|
||
"learning_rate": self.config.learning_rate,
|
||
"opponent": self.config.opponent,
|
||
"playstyle": self.config.playstyle,
|
||
"data_source": self.config.data_source,
|
||
},
|
||
"best_win_rate": self.best_win_rate,
|
||
"model_config": {
|
||
"input_dim": STATE_DIMS.total_state_dim,
|
||
"hidden_dim": TRAINING_CONFIG.hidden_dim,
|
||
"num_actions": NUM_ACTIONS,
|
||
"num_layers": TRAINING_CONFIG.num_layers,
|
||
"dropout": 0.0,
|
||
},
|
||
}, path)
|
||
logger.info("Checkpoint saved: %s", path)
|
||
|
||
def train(self, progress_callback=None) -> Dict:
|
||
"""Run the full PPO training loop."""
|
||
cfg = self.config
|
||
logger.info("=" * 60)
|
||
logger.info("PPO Training — Commander AI Lab")
|
||
logger.info("=" * 60)
|
||
logger.info(" Iterations: %d", cfg.iterations)
|
||
logger.info(" Episodes/iter: %d", cfg.episodes_per_iter)
|
||
logger.info(" PPO epochs: %d", cfg.ppo_epochs)
|
||
logger.info(" Batch size: %d", cfg.batch_size)
|
||
logger.info(" Clip epsilon: %.2f", cfg.clip_epsilon)
|
||
logger.info(" Entropy coeff: %.4f", cfg.entropy_coeff)
|
||
logger.info(" Learning rate: %.6f", cfg.learning_rate)
|
||
logger.info(" Opponent: %s", cfg.opponent)
|
||
logger.info(" Device: %s", self.device)
|
||
logger.info(" Supervised: %s", cfg.load_supervised or "None")
|
||
logger.info(" Data source: %s", cfg.data_source)
|
||
if cfg.data_source in ("forge", "mixed"):
|
||
logger.info(" Forge results: %s / %s", cfg.forge_results_dir, cfg.forge_batch_id)
|
||
logger.info(" Forge weight: %.0f%%", cfg.forge_weight * 100)
|
||
logger.info("")
|
||
|
||
# Start Forge producer background thread if needed
|
||
if self._forge_source:
|
||
self._forge_source.start()
|
||
|
||
t_start = time.time()
|
||
|
||
try:
|
||
for iteration in range(1, cfg.iterations + 1):
|
||
self.iteration = iteration
|
||
t_iter = time.time()
|
||
|
||
rollout_stats = self._collect_rollouts_for_iter()
|
||
self.total_steps += rollout_stats["steps"]
|
||
|
||
update_metrics = self.ppo_update()
|
||
|
||
current_lr = self.optimizer.param_groups[0]["lr"]
|
||
if self.scheduler:
|
||
self.scheduler.step()
|
||
|
||
eval_results = None
|
||
if iteration % cfg.eval_every == 0 or iteration == 1:
|
||
eval_results = self.evaluate(cfg.eval_episodes)
|
||
heuristic_wr = eval_results.get("heuristic", {}).get("win_rate", 0)
|
||
if heuristic_wr > self.best_win_rate:
|
||
self.best_win_rate = heuristic_wr
|
||
best_path = os.path.join(cfg.checkpoint_dir, "best_ppo.pt")
|
||
self.save_checkpoint(best_path, {"eval": eval_results})
|
||
logger.info(" ★ New best win rate: %.1f%%", heuristic_wr * 100)
|
||
|
||
if iteration % cfg.save_every == 0:
|
||
self.save_checkpoint()
|
||
|
||
iter_time = time.time() - t_iter
|
||
metrics = {
|
||
"iteration": iteration,
|
||
"steps": rollout_stats["steps"],
|
||
"total_steps": self.total_steps,
|
||
"win_rate": rollout_stats["win_rate"],
|
||
"avg_reward": rollout_stats["avg_reward_per_ep"],
|
||
"policy_loss": update_metrics["policy_loss"],
|
||
"value_loss": update_metrics["value_loss"],
|
||
"entropy": update_metrics["entropy"],
|
||
"approx_kl": update_metrics["approx_kl"],
|
||
"clip_fraction": update_metrics["clip_fraction"],
|
||
"lr": current_lr,
|
||
"iter_time_s": round(iter_time, 1),
|
||
"data_source": rollout_stats.get("source", cfg.data_source),
|
||
"eval": eval_results,
|
||
}
|
||
self.history.append(metrics)
|
||
|
||
eval_str = ""
|
||
if eval_results:
|
||
h_wr = eval_results.get("heuristic", {}).get("win_rate", 0)
|
||
r_wr = eval_results.get("random", {}).get("win_rate", 0)
|
||
eval_str = f" | eval: h={h_wr:.0%} r={r_wr:.0%}"
|
||
|
||
src_tag = f"[{rollout_stats.get('source', '?')[:3]}]"
|
||
logger.info(
|
||
" [%3d/%d]%s wr=%.0f%% rew=%.3f ploss=%.4f vloss=%.4f ent=%.3f kl=%.4f%s (%.1fs)",
|
||
iteration, cfg.iterations, src_tag,
|
||
rollout_stats["win_rate"] * 100,
|
||
rollout_stats["avg_reward_per_ep"],
|
||
update_metrics["policy_loss"],
|
||
update_metrics["value_loss"],
|
||
update_metrics["entropy"],
|
||
update_metrics["approx_kl"],
|
||
eval_str,
|
||
iter_time,
|
||
)
|
||
|
||
if progress_callback:
|
||
progress_callback(iteration, metrics)
|
||
|
||
finally:
|
||
if self._forge_source:
|
||
self._forge_source.stop()
|
||
|
||
total_time = time.time() - t_start
|
||
self.save_checkpoint()
|
||
|
||
history_path = os.path.join(cfg.checkpoint_dir, "ppo_history.json")
|
||
with open(history_path, "w", encoding="utf-8") as f:
|
||
json.dump(self.history, f, indent=2, default=str)
|
||
|
||
summary = {
|
||
"iterations": cfg.iterations,
|
||
"total_steps": self.total_steps,
|
||
"best_win_rate": self.best_win_rate,
|
||
"final_win_rate": self.history[-1]["win_rate"] if self.history else 0,
|
||
"total_time_s": round(total_time, 1),
|
||
"device": self.device,
|
||
"data_source": cfg.data_source,
|
||
"checkpoint_dir": cfg.checkpoint_dir,
|
||
"history_path": history_path,
|
||
}
|
||
|
||
logger.info("")
|
||
logger.info("=" * 60)
|
||
logger.info("PPO Training Complete")
|
||
logger.info(" Total time: %.1f s", total_time)
|
||
logger.info(" Total steps: %d", self.total_steps)
|
||
logger.info(" Best win rate: %.1f%%", self.best_win_rate * 100)
|
||
logger.info(" Data source: %s", cfg.data_source)
|
||
logger.info(" History: %s", history_path)
|
||
logger.info("=" * 60)
|
||
|
||
return summary
|
||
|
||
|
||
def main():
|
||
"""CLI entry point for PPO training."""
|
||
parser = argparse.ArgumentParser(description="Commander AI Lab — PPO Training")
|
||
|
||
# Training loop
|
||
parser.add_argument("--iterations", type=int, default=100)
|
||
parser.add_argument("--episodes-per-iter", type=int, default=64)
|
||
parser.add_argument("--ppo-epochs", type=int, default=4)
|
||
parser.add_argument("--batch-size", type=int, default=256)
|
||
|
||
# PPO hyperparams
|
||
parser.add_argument("--clip-epsilon", type=float, default=0.2)
|
||
parser.add_argument("--entropy-coeff", type=float, default=0.01)
|
||
parser.add_argument("--value-loss-coeff", type=float, default=0.5)
|
||
parser.add_argument("--lr", type=float, default=3e-4)
|
||
parser.add_argument("--lr-schedule", default="constant", choices=["constant", "linear", "cosine"])
|
||
|
||
# Rollout
|
||
parser.add_argument("--buffer-size", type=int, default=4096)
|
||
parser.add_argument("--gamma", type=float, default=0.99)
|
||
parser.add_argument("--gae-lambda", type=float, default=0.95)
|
||
|
||
# Opponent
|
||
parser.add_argument("--opponent", default="heuristic", choices=["heuristic", "random", "self"])
|
||
parser.add_argument("--playstyle", default="midrange")
|
||
|
||
# Checkpointing
|
||
parser.add_argument("--checkpoint-dir", default="ml/models/checkpoints")
|
||
parser.add_argument("--save-every", type=int, default=10)
|
||
parser.add_argument("--eval-every", type=int, default=5)
|
||
parser.add_argument("--eval-episodes", type=int, default=50)
|
||
|
||
# Init
|
||
parser.add_argument("--load-supervised", default=None,
|
||
help="Path to supervised checkpoint to initialize from")
|
||
|
||
# Data source
|
||
parser.add_argument(
|
||
"--data-source", default="synthetic",
|
||
choices=["synthetic", "forge", "mixed"],
|
||
help="Episode source: synthetic (self-play), forge (Forge JSONL), mixed (blend)",
|
||
)
|
||
parser.add_argument(
|
||
"--forge-results-dir", default="results",
|
||
help="Directory where Forge writes ml-decisions-forge-*.jsonl files",
|
||
)
|
||
parser.add_argument(
|
||
"--forge-batch-id", default="forge-batch",
|
||
help="Batch ID prefix for Forge JSONL files (matches BatchRunner --batchId)",
|
||
)
|
||
parser.add_argument(
|
||
"--forge-num-games", type=int, default=256,
|
||
help="Total Forge games to consume before stopping the producer",
|
||
)
|
||
parser.add_argument(
|
||
"--forge-weight", type=float, default=0.7,
|
||
help="Fraction of RolloutBuffer filled from Forge when --data-source mixed (0.0–1.0)",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||
|
||
config = PPOConfig(
|
||
iterations=args.iterations,
|
||
episodes_per_iter=args.episodes_per_iter,
|
||
ppo_epochs=args.ppo_epochs,
|
||
batch_size=args.batch_size,
|
||
clip_epsilon=args.clip_epsilon,
|
||
entropy_coeff=args.entropy_coeff,
|
||
value_loss_coeff=args.value_loss_coeff,
|
||
learning_rate=args.lr,
|
||
lr_schedule=args.lr_schedule,
|
||
buffer_size=args.buffer_size,
|
||
gamma=args.gamma,
|
||
gae_lambda=args.gae_lambda,
|
||
opponent=args.opponent,
|
||
playstyle=args.playstyle,
|
||
checkpoint_dir=args.checkpoint_dir,
|
||
save_every=args.save_every,
|
||
eval_every=args.eval_every,
|
||
eval_episodes=args.eval_episodes,
|
||
load_supervised=args.load_supervised,
|
||
data_source=args.data_source,
|
||
forge_results_dir=args.forge_results_dir,
|
||
forge_batch_id=args.forge_batch_id,
|
||
forge_num_games=args.forge_num_games,
|
||
forge_weight=args.forge_weight,
|
||
)
|
||
|
||
trainer = PPOTrainer(config)
|
||
summary = trainer.train()
|
||
print(json.dumps(summary, indent=2))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|