train.py

#!/usr/bin/env python
"""PPO training CLI for GeneralsBot.

Usage:
    python train.py                                              # defaults (10x10, vs random)
    python train.py --grid 15x15 --opponent expander
    python train.py --config experiments/my_exp.yaml
    python train.py --config experiments/my_exp.yaml --lr 1e-4 --tag test
"""

from __future__ import annotations

import argparse
import dataclasses
import sys
from datetime import datetime
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(0, str(PROJECT_ROOT / "generals-bots"))


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="PPO training for GeneralsBot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # Config / resume
    p.add_argument("--config", "-c", type=str, default=None,
                    help="Path to YAML config file.")
    p.add_argument("--resume", "-r", type=str, default=None,
                    help="Resume training from a checkpoint directory "
                         "(loads config.yaml + train_state.eqx from that dir).")
    p.add_argument("--init-from", type=str, default=None,
                    help="Path to a .eqx checkpoint for network initialisation "
                         "(e.g. from SL pretraining). Cannot be used with --resume.")
    p.add_argument("--tag", type=str, default=None,
                    help="Optional experiment tag appended to the run directory name.")

    # Supervised learning mode
    p.add_argument("--sl", action="store_true",
                    help="Run supervised learning (behaviour cloning) from human replays.")
    p.add_argument("--replay-dir", type=str,
                    default="/root/autodl-tmp/generals_io_replays",
                    help="Path to the HuggingFace replay dataset directory.")
    p.add_argument("--sl-epochs", type=int, default=50,
                    help="Number of SL training epochs (default: 50).")
    p.add_argument("--sl-lr", type=float, default=3e-4,
                    help="SL learning rate (default: 3e-4).")
    p.add_argument("--sl-batch-size", type=int, default=256,
                    help="SL minibatch size (default: 256).")
    p.add_argument("--sl-max-replays", type=int, default=None,
                    help="Limit on replays for SL (default: all).")
    p.add_argument("--sl-cache-dir", type=str, default="sl_data_cache",
                    help="Directory to cache processed replay data (default: sl_data_cache).")
    p.add_argument("--sl-max-files", type=int, default=None,
                    help="Limit on cached .npz files for training (default: all). "
                         "Useful for quick partial runs without rebuilding cache.")
    p.add_argument("--sl-chunk-size", type=int, default=16,
                    help="Files to decompress per training step (default: 16 = 4096 samples). "
                         "Larger batches increase GPU throughput but need more VRAM.")

    # Environment
    p.add_argument("--grid", type=str, default=None,
                    help="Fixed grid size HxW (e.g. 10x10). Overrides min/max-grid.")
    p.add_argument("--min-grid", type=int, default=None,
                    help="Min grid edge for variable-size training.")
    p.add_argument("--max-grid", type=int, default=None,
                    help="Max grid edge for variable-size training.")
    p.add_argument("--truncation", type=int, default=500,
                    help="Max ticks before draw (default: 500).")

    # Rollout
    p.add_argument("--num-envs", type=int, default=256,
                    help="Parallel environments (default: 256).")
    p.add_argument("--num-steps", type=int, default=256,
                    help="Steps per env per rollout (default: 256).")

    # Optimisation
    p.add_argument("--total-iterations", type=int, default=1000,
                    help="Training iterations (default: 1000).")
    p.add_argument("--lr", type=float, default=3e-4,
                    help="Learning rate (default: 3e-4).")
    p.add_argument("--num-epochs", type=int, default=4,
                    help="PPO epochs per rollout (default: 4).")
    p.add_argument("--minibatch-size", type=int, default=512,
                    help="Minibatch size (default: 512).")
    p.add_argument("--clip-epsilon", type=float, default=0.2,
                    help="PPO clip range (default: 0.2).")
    p.add_argument("--entropy-coef", type=float, default=0.01,
                    help="Entropy bonus coefficient (default: 0.01).")

    # Curriculum
    p.add_argument("--opponent", choices=["random", "expander", "self"],
                    default="random",
                    help="Starting opponent (default: random).")
    p.add_argument("--curriculum", type=str, nargs="+",
                    default=["random", "expander", "self"],
                    help="Curriculum phases (default: random expander self).")
    p.add_argument("--curriculum-threshold", type=float, default=0.75,
                    help="Win-rate to advance phase (default: 0.75).")
    p.add_argument("--curriculum-steps", type=int, default=50,
                    help="Min iterations per phase (default: 50).")
    p.add_argument("--self-play-window", type=int, default=5,
                    help="Past checkpoints in self-play pool (default: 5).")

    # Eval & logging
    p.add_argument("--eval-interval", type=int, default=50,
                    help="Eval every N iterations (default: 50).")
    p.add_argument("--save-dir", type=str, default="checkpoints",
                    help="Checkpoint directory (default: checkpoints).")
    p.add_argument("--log-file", type=str, default="training_log.jsonl",
                    help="Metrics log (default: training_log.jsonl).")
    p.add_argument("--visualize", action=argparse.BooleanOptionalAction, default=True,
                    help="Show game visualization during eval. Toggle during training "
                         "by creating/deleting the .vis_enable file (default: True).")
    p.add_argument("--seed", type=int, default=42,
                    help="Random seed (default: 42).")

    return p


def parse_grid(s: str | None) -> tuple[int, int] | None:
    if s is None:
        return None
    parts = s.split("x")
    if len(parts) != 2:
        raise argparse.ArgumentTypeError(f"Expected HxW, got '{s}'")
    return int(parts[0]), int(parts[1])


# Mapping from CLI dest names to PPOConfig field names.
# Only includes fields where the names differ or need special handling.
_CLI_TO_CFG = {
    "truncation": "truncation",
    "num_envs": "num_envs",
    "num_steps": "num_steps",
    "total_iterations": "total_iterations",
    "lr": "learning_rate",
    "num_epochs": "num_epochs",
    "minibatch_size": "minibatch_size",
    "clip_epsilon": "clip_epsilon",
    "entropy_coef": "entropy_coef",
    "opponent": "opponent",
    "curriculum": "curriculum",
    "curriculum_threshold": "curriculum_threshold",
    "curriculum_steps": "curriculum_steps",
    "self_play_window": "self_play_window",
    "eval_interval": "eval_interval",
    "save_dir": "save_dir",
    "log_file": "log_file",
    "visualize": "visualize",
    "seed": "seed",
}


def _explicit_args(parser: argparse.ArgumentParser, args: argparse.Namespace) -> set[str]:
    """Return set of CLI dest names that the user explicitly provided."""
    defaults = vars(parser.parse_args([]))
    explicit = set()
    for key, val in vars(args).items():
        if key in ("config", "tag"):
            if val is not None:
                explicit.add(key)
            continue
        if key not in defaults:
            continue
        if val != defaults[key]:
            explicit.add(key)
    return explicit


def load_yaml_config(path: str | Path) -> dict:
    """Load a YAML config file. Returns a dict keyed by PPOConfig field names."""
    import yaml

    with open(path) as f:
        data = yaml.safe_load(f) or {}

    # Normalise grid_dims if given as a string "HxW"
    if isinstance(data.get("grid_dims"), str):
        data["grid_dims"] = parse_grid(data["grid_dims"])

    return data


def _apply_yaml(cfg: object, data: dict) -> None:
    """Apply YAML-sourced values to a PPOConfig instance (does not override CLI)."""
    for key, value in data.items():
        if not hasattr(cfg, key):
            print(f"  ⚠ Unknown config key: {key}, ignoring")
            continue
        if key == "curriculum" and isinstance(value, list):
            value = tuple(value)
        setattr(cfg, key, value)


def _save_config_yaml(cfg: object, path: Path) -> None:
    """Serialise a PPOConfig to YAML."""
    import yaml

    data = dataclasses.asdict(cfg)
    with open(path, "w") as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False)


def _print_summary(cfg) -> None:
    grid_repr = (
        f"{cfg.grid_dims[0]}x{cfg.grid_dims[1]}"
        if cfg.grid_dims is not None
        else f"{cfg.min_grid_size}–{cfg.max_grid_size} (pad to {cfg.max_grid_size})"
    )
    print(f"{'='*60}")
    print(f"PPO Training — grid={grid_repr}  "
          f"envs={cfg.num_envs}  steps={cfg.num_steps}  "
          f"iter={cfg.total_iterations}")
    print(f"Curriculum: {' → '.join(cfg.curriculum)}  "
          f"threshold={cfg.curriculum_threshold:.0%}")
    print(f"Visualization: {'ON' if cfg.visualize else 'OFF'}  "
          f"(toggle during training via .vis_enable file)")
    print(f"{'='*60}\n")


# ── CLI entry point ────────────────────────────────────────────────────────


def main():
    # ── JAX / XLA thread control (must be set before any JAX import) ──
    import os
    os.environ.setdefault("XLA_FLAGS",
        "--xla_cpu_multi_thread_eigen=false --xla_gpu_force_compilation_parallelism=1")
    os.environ.setdefault("TF_NUM_INTRAOP_THREADS", "1")
    os.environ.setdefault("TF_NUM_INTEROP_THREADS", "1")
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

    parser = build_parser()
    args = parser.parse_args()

    from generals_bot.decision.trainer import PPOConfig, train_ppo, train_sl, load_config_yaml
    from generals_bot.decision.model import PolicyValueNet
    import equinox as eqx
    import jax.random as jrandom

    explicit = _explicit_args(parser, args)

    # ═══════════════════════════════════════════════════════════════════
    # Mode 1: Supervised Learning (behaviour cloning from human replays)
    # ═══════════════════════════════════════════════════════════════════
    if args.sl:
        print("=" * 60)
        print("Supervised Learning mode — behaviour cloning from human replays")
        print("=" * 60)
        network = train_sl(
            replay_dir=args.replay_dir,
            model_save_path="checkpoints/sl_pretrained.eqx",
            num_epochs=args.sl_epochs,
            learning_rate=args.sl_lr,
            batch_size=args.sl_batch_size,
            max_replays=args.sl_max_replays,
            max_files=args.sl_max_files,
            chunk_size=args.sl_chunk_size,
            seed=args.seed,
            cache_dir=args.sl_cache_dir,
        )
        print(f"\nSL training complete. Model: checkpoints/sl_pretrained.eqx")
        return

    # ═══════════════════════════════════════════════════════════════════
    # Mode 2: PPO Training
    # ═══════════════════════════════════════════════════════════════════

    # Resolve init-from early (mutually exclusive with resume)
    init_network = None
    if args.init_from is not None:
        if args.resume is not None:
            print("Error: --init-from and --resume are mutually exclusive.")
            sys.exit(1)
        init_path = Path(args.init_from)
        if not init_path.exists():
            print(f"Error: --init-from path {init_path} not found.")
            sys.exit(1)
        dummy = PolicyValueNet(jrandom.PRNGKey(0))
        init_network = eqx.tree_deserialise_leaves(init_path, dummy)
        print(f"Loaded initial network from: {init_path}")

    resume_dir: Path | None = None

    # ── Resolve config ──────────────────────────────────────────────────
    # Priority: PPOConfig defaults < YAML file < explicit CLI flags
    # Except in resume mode: saved config.yaml < explicit CLI overrides

    if args.resume is not None:
        resume_dir = Path(args.resume)
        config_path = resume_dir / "config.yaml"
        if not config_path.exists():
            print(f"Error: {config_path} not found. Cannot resume.")
            sys.exit(1)
        cfg = load_config_yaml(config_path)
        # Apply explicit CLI overrides on top of saved config
        if "grid" in explicit:
            cfg.grid_dims = parse_grid(args.grid)
            cfg.min_grid_size = None
            cfg.max_grid_size = None
        else:
            if "min_grid" in explicit:
                cfg.grid_dims = None
                cfg.min_grid_size = args.min_grid
            if "max_grid" in explicit:
                cfg.grid_dims = None
                cfg.max_grid_size = args.max_grid
        for cli_key, cfg_key in _CLI_TO_CFG.items():
            if cli_key in explicit:
                val = getattr(args, cli_key)
                if cli_key == "curriculum" and isinstance(val, list):
                    val = tuple(val)
                setattr(cfg, cfg_key, val)
        # Default to the resume dir for output (overridable)
        if "save_dir" not in explicit:
            cfg.save_dir = str(resume_dir)
        if "log_file" not in explicit:
            cfg.log_file = str(resume_dir / "training_log.jsonl")
        print(f"Resuming from: {resume_dir}")
    elif args.config is not None:
        yaml_data = load_yaml_config(args.config)
        cfg = PPOConfig()
        _apply_yaml(cfg, yaml_data)

        # Grid-size special handling
        if "grid" in explicit:
            cfg.grid_dims = parse_grid(args.grid)
            cfg.min_grid_size = None
            cfg.max_grid_size = None
        else:
            if "min_grid" in explicit:
                cfg.grid_dims = None
                cfg.min_grid_size = args.min_grid
            if "max_grid" in explicit:
                cfg.grid_dims = None
                cfg.max_grid_size = args.max_grid

        for cli_key, cfg_key in _CLI_TO_CFG.items():
            if cli_key in explicit:
                val = getattr(args, cli_key)
                if cli_key == "curriculum" and isinstance(val, list):
                    val = tuple(val)
                setattr(cfg, cfg_key, val)

        # Run directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        dir_name = timestamp
        if args.tag:
            dir_name += f"_{args.tag}"
        run_dir = Path("runs") / dir_name
        run_dir.mkdir(parents=True, exist_ok=True)

        _save_config_yaml(cfg, run_dir / "config.yaml")
        cfg.save_dir = str(run_dir / "checkpoints")
        cfg.log_file = str(run_dir / "training_log.jsonl")
        print(f"Run directory: {run_dir}")
    else:
        cfg = PPOConfig()

        # Grid-size special handling
        if "grid" in explicit:
            cfg.grid_dims = parse_grid(args.grid)
            cfg.min_grid_size = None
            cfg.max_grid_size = None
        else:
            if "min_grid" in explicit:
                cfg.grid_dims = None
                cfg.min_grid_size = args.min_grid
            if "max_grid" in explicit:
                cfg.grid_dims = None
                cfg.max_grid_size = args.max_grid

        for cli_key, cfg_key in _CLI_TO_CFG.items():
            if cli_key in explicit:
                val = getattr(args, cli_key)
                if cli_key == "curriculum" and isinstance(val, list):
                    val = tuple(val)
                setattr(cfg, cfg_key, val)

    # ── Print summary & launch training ────────────────────────────────
    _print_summary(cfg)
    network = train_ppo(
        cfg,
        resume_dir=str(resume_dir) if resume_dir else None,
        init_network=init_network,
    )
    print(f"\nDone. Trained model at: {cfg.save_dir}/best.eqx")


if __name__ == "__main__":
    main()