| """ |
| Train the Blue policy against a PFSP-sampled pool of Red opponents. |
| |
| This is the practical entry point for the "Blue learner vs Red league" |
| configuration in the report. It reuses the rollout + TRL helpers from |
| `colab_trl_selfplay`, swapping the learner side to `blue` and updating |
| opponent win-rate stats after each round. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from train._bootstrap import ensure_repo_root_on_path |
|
|
| ensure_repo_root_on_path() |
|
|
| import argparse |
| import os |
| from statistics import mean |
| from typing import List |
|
|
| from train.colab_trl_selfplay import ( |
| HAS_TRL_STACK, |
| build_training_dataset, |
| collect_rollouts, |
| maybe_upload_to_hub, |
| run_trl_sft, |
| save_artifacts, |
| ) |
| from train.pfsp import OpponentStats, sample_opponent |
|
|
|
|
| def _seed_pool() -> List[OpponentStats]: |
| return [ |
| OpponentStats(name="R-easy", win_rate_vs_learner=0.20, games=10), |
| OpponentStats(name="R-mid", win_rate_vs_learner=0.50, games=10), |
| OpponentStats(name="R-hard", win_rate_vs_learner=0.80, games=10), |
| ] |
|
|
|
|
| def _update_winrate(opp: OpponentStats, summaries: list[dict]) -> None: |
| """Approximate win-rate-vs-learner update using mean Blue reward sign.""" |
| if not summaries: |
| return |
| avg_r = mean(s.get("avg_learner_reward", 0.0) for s in summaries) |
| blue_won = 1.0 if avg_r > 0 else 0.0 |
| new_games = opp.games + len(summaries) |
| blended = ( |
| (opp.win_rate_vs_learner * opp.games) + ((1.0 - blue_won) * len(summaries)) |
| ) / max(1, new_games) |
| opp.win_rate_vs_learner = float(min(1.0, max(0.0, blended))) |
| opp.games = new_games |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Blue vs Red-pool PFSP trainer.") |
| parser.add_argument("--rounds", type=int, default=int(os.getenv("ROUNDS", "2"))) |
| parser.add_argument("--episodes", type=int, default=int(os.getenv("EPISODES", "10"))) |
| parser.add_argument("--max-steps", type=int, default=int(os.getenv("MAX_STEPS", "60"))) |
| parser.add_argument( |
| "--no-train", |
| action="store_true", |
| help="Collect rollouts + artifacts only, skip TRL fine-tune.", |
| ) |
| args = parser.parse_args() |
|
|
| pool = _seed_pool() |
| last_out_dir = "outputs/cyber-blue-sft" |
|
|
| for r in range(args.rounds): |
| opp = sample_opponent(pool) |
| print( |
| f"[blue-train round {r}] PFSP picked Red={opp.name} " |
| f"(w_vs_learner={opp.win_rate_vs_learner:.2f})" |
| ) |
|
|
| rollouts, summaries = collect_rollouts( |
| num_episodes=args.episodes, |
| max_steps=args.max_steps, |
| side="blue", |
| ) |
| _update_winrate(opp, summaries) |
|
|
| if HAS_TRL_STACK and not args.no_train: |
| ds = build_training_dataset(rollouts, keep_top_frac=0.5) |
| last_out_dir = run_trl_sft(ds, output_dir=f"outputs/cyber-blue-sft-r{r}") |
| save_artifacts(rollouts, summaries, last_out_dir) |
|
|
| print("[blue-train] final pool state:") |
| for o in pool: |
| print( |
| f" {o.name}: win_rate_vs_learner={o.win_rate_vs_learner:.2f} games={o.games}" |
| ) |
|
|
| maybe_upload_to_hub(last_out_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
