CosmicClash/training/train.py

"""Train the Cosmic Clash self-play PPO policy.

Example (smoke run):
    .venv/bin/python train.py --experiment smoke --timesteps 100000

Long run on the Linux/CUDA box:
    GODOT_BIN=~/godot/Godot_v4.7.1-stable_linux.x86_64 \
        .venv/bin/python train.py --experiment run01 --timesteps 20000000 \
        --n-parallel 6 --speedup 16

See TRAINING.md at the repo root for the full workflow.
"""

import argparse
import os
import pathlib

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback
from stable_baselines3.common.vec_env.vec_monitor import VecMonitor

from cosmic_env import CosmicClashVecEnv

TRAINING_DIR = pathlib.Path(__file__).resolve().parent
DEFAULT_GODOT_MACOS = "/Applications/Godot.app/Contents/MacOS/Godot"


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--godot_bin",
        default=os.environ.get("GODOT_BIN", DEFAULT_GODOT_MACOS),
        help="Path to the Godot binary (or set GODOT_BIN)",
    )
    parser.add_argument("--experiment", default="default", help="Run name for logs/checkpoints")
    parser.add_argument("--timesteps", type=int, default=200_000)
    parser.add_argument("--n-parallel", type=int, default=2, help="Parallel Godot instances (2 agents each)")
    parser.add_argument("--speedup", type=int, default=8, help="Physics speedup factor inside Godot")
    parser.add_argument("--port", type=int, default=11008, help="Base TCP port (one per instance)")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--resume", default=None, help="Checkpoint .zip to resume from")
    parser.add_argument("--ent-coef", type=float, default=0.0001, help="Entropy bonus coefficient (applied on resume too)")
    parser.add_argument("--n-steps", type=int, default=256, help="Rollout length per env between updates (applied on resume too)")
    parser.add_argument("--batch-size", type=int, default=256, help="PPO minibatch size (applied on resume too)")
    parser.add_argument(
        "--reset-std",
        type=float,
        default=None,
        help="On resume, reset the policy action std to this value (recovers exploration after entropy collapse)",
    )
    parser.add_argument("--checkpoint-every", type=int, default=100_000, help="Timesteps between checkpoints")
    parser.add_argument("--viz", action="store_true", help="Show game windows (debugging; slow)")
    parser.add_argument("--wandb", action="store_true", help="Also log to Weights & Biases")

    curriculum = parser.add_argument_group(
        "curriculum", "Stage the training run — see TRAINING.md's Curriculum training section"
    )
    curriculum.add_argument(
        "--opponent-mode",
        choices=["self_play", "inert", "frozen"],
        default=None,
        help="self_play (default): both ships are live trainees. inert: team 1 is a "
        "do-nothing placeholder (isolated scoring practice). frozen: team 1 runs a "
        "fixed exported policy (--opponent-model)",
    )
    curriculum.add_argument("--opponent-model", default=None, help="Exported policy .json for --opponent-mode=frozen")
    curriculum.add_argument(
        "--draw-penalty", type=float, default=None, help="One-time penalty when an episode times out with no goal"
    )
    curriculum.add_argument(
        "--attack-goal-bias",
        type=float,
        default=None,
        help="0.5 = uniform between both goals (default); 1.0 = near-goal resets always target the goal team 0 attacks",
    )
    curriculum.add_argument("--kickoff-chance", type=float, default=None, help="Overrides kickoff_state_chance")
    curriculum.add_argument("--near-goal-chance", type=float, default=None, help="Overrides ball_near_goal_chance")
    curriculum.add_argument(
        "--allow-vertical", action=argparse.BooleanOptionalAction, default=None, help="Allow vertical thrust (default true)"
    )
    curriculum.add_argument(
        "--allow-pitch-roll",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="Allow pitch/roll rotation (default true)",
    )

    return parser.parse_args()


def _curriculum_kwargs(args) -> dict:
    """Maps train.py's curriculum flags to the --key=value args training_mode.gd's
    _parse_curriculum_args() reads, omitting anything not explicitly passed so
    unset flags leave Godot's own @export defaults in place."""
    mapping = {
        "opponent_mode": args.opponent_mode,
        "opponent_model": args.opponent_model,
        "draw_penalty": args.draw_penalty,
        "attack_goal_bias": args.attack_goal_bias,
        "kickoff_state_chance": args.kickoff_chance,
        "ball_near_goal_chance": args.near_goal_chance,
        "ai_allow_vertical": args.allow_vertical,
        "ai_allow_pitch_roll": args.allow_pitch_roll,
    }
    return {key: value for key, value in mapping.items() if value is not None}


def main():
    args = parse_args()
    log_dir = TRAINING_DIR / "logs"
    checkpoint_dir = TRAINING_DIR / "checkpoints" / args.experiment
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    if args.wandb:
        import wandb

        wandb.init(project="cosmic-clash-rl", name=args.experiment, sync_tensorboard=True)

    env = CosmicClashVecEnv(
        godot_bin=args.godot_bin,
        n_parallel=args.n_parallel,
        seed=args.seed,
        port=args.port,
        show_window=args.viz,
        speedup=args.speedup,
        **_curriculum_kwargs(args),
    )
    env = VecMonitor(env)

    if args.resume:
        model = PPO.load(
            args.resume,
            env=env,
            tensorboard_log=str(log_dir),
            ent_coef=args.ent_coef,
            n_steps=args.n_steps,
            batch_size=args.batch_size,
        )
        print(
            f"Resumed from {args.resume} at {model.num_timesteps} timesteps "
            f"(ent_coef={args.ent_coef}, n_steps={args.n_steps}, batch_size={args.batch_size})"
        )
        if args.reset_std is not None:
            import math

            import torch

            with torch.no_grad():
                model.policy.log_std.fill_(math.log(args.reset_std))
            print(f"Reset policy action std to {args.reset_std}")
    else:
        model = PPO(
            "MultiInputPolicy",
            env,
            verbose=1,
            ent_coef=args.ent_coef,
            n_steps=args.n_steps,
            batch_size=args.batch_size,
            learning_rate=3e-4,
            tensorboard_log=str(log_dir),
        )

    checkpoint_callback = CheckpointCallback(
        save_freq=max(args.checkpoint_every // env.num_envs, 1),
        save_path=str(checkpoint_dir),
        name_prefix="ppo",
    )

    try:
        model.learn(
            args.timesteps,
            callback=checkpoint_callback,
            tb_log_name=args.experiment,
            reset_num_timesteps=not args.resume,
        )
    finally:
        final_path = checkpoint_dir / "final.zip"
        model.save(str(final_path))
        print(f"Saved {final_path}")
        env.close()


if __name__ == "__main__":
    main()