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 BaseCallback, CheckpointCallback
from stable_baselines3.common.utils import safe_mean
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"


class GoalRateCallback(BaseCallback):
    """Logs rollout/goal_rate: the fraction of completed episodes in the
    current ep_info_buffer that ended in an actual goal, vs. timing out as a
    draw. rollout/ep_rew_mean mixes dense reward-shaping (ball chasing/
    touching) with the sparse terminal goal reward, so it can trend up from
    better shaping alone without the policy finishing more episodes by
    actually scoring — this isolates that. Requires VecMonitor(...,
    info_keywords=("goal_scored",)), which copies ShipAIController.get_info()
    into each completed episode's info["episode"] dict (see
    training_mode.gd's _on_goal_scored / timeout branch)."""

    def _on_step(self) -> bool:
        return True

    def _on_rollout_end(self) -> None:
        if len(self.model.ep_info_buffer) == 0:
            return
        # The vendored godot_rl sync bridge (Game/addons/godot_rl_agents/sync.gd,
        # _training_process) snapshots each agent's info dict once per tick and
        # has its own "NEEDS REFACTOR" comment on the reset-timing path, so an
        # episode's terminal info entry can arrive without "goal_scored" at all
        # (observed crashing a run after 2026-07-28). Skip those rather than
        # crash training over a monitoring-only metric.
        rates = [ep_info["goal_scored"] for ep_info in self.model.ep_info_buffer if "goal_scored" in ep_info]
        if rates:
            self.logger.record("rollout/goal_rate", safe_mean(rates))


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(
        "--exported-binary",
        default=None,
        help="Path to a pre-built game executable (see export_linux.sh) instead of running the "
        "project from source — skips per-instance script/resource import for faster parallel "
        "startup. Overrides --godot_bin when set.",
    )
    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(
        "--vertical-ramp", type=float, default=None,
        help="0.0-1.0: fraction of vertical thrust that reaches the ship (locomotion-unmask ramp; default 1.0)",
    )
    curriculum.add_argument(
        "--pitch-roll-ramp", type=float, default=None,
        help="0.0-1.0: fraction of pitch/roll rotation that reaches the ship (locomotion-unmask ramp; default 1.0)",
    )
    curriculum.add_argument(
        "--velocity-to-ball-weight", type=float, default=None,
        help="Overrides ShipAIController.velocity_to_ball_weight (dense reward for closing speed toward the ball)",
    )
    curriculum.add_argument(
        "--ball-distance-penalty", type=float, default=None,
        help="Overrides ShipAIController.ball_distance_penalty (dense per-tick cost scaled by distance to the ball)",
    )
    curriculum.add_argument(
        "--ball-touch-reward", type=float, default=None,
        help="Overrides ShipAIController.ball_touch_reward (event reward on ball contact, cooldown-gated)",
    )
    curriculum.add_argument(
        "--airborne-penalty", type=float, default=None,
        help="Overrides ShipAIController.airborne_penalty (dense per-tick cost scaled by height above the floor)",
    )
    curriculum.add_argument(
        "--ball-velocity-to-goal-weight", type=float, default=None,
        help="Overrides ShipAIController.ball_velocity_to_goal_weight (dense reward for the ball's velocity toward the attack goal)",
    )
    curriculum.add_argument(
        "--goal-reward", type=float, default=None,
        help="Overrides TrainingMode.goal_reward (terminal reward for actually scoring)",
    )

    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_vertical_ramp": args.vertical_ramp,
        "ai_pitch_roll_ramp": args.pitch_roll_ramp,
        "ai_velocity_to_ball_weight": args.velocity_to_ball_weight,
        "ai_ball_distance_penalty": args.ball_distance_penalty,
        "ai_ball_touch_reward": args.ball_touch_reward,
        "ai_airborne_penalty": args.airborne_penalty,
        "ai_ball_velocity_to_goal_weight": args.ball_velocity_to_goal_weight,
        "goal_reward": args.goal_reward,
    }
    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)

    exported_binary = args.exported_binary or None
    env = CosmicClashVecEnv(
        godot_bin=exported_binary or args.godot_bin,
        exported=exported_binary is not None,
        n_parallel=args.n_parallel,
        seed=args.seed,
        port=args.port,
        show_window=args.viz,
        speedup=args.speedup,
        **_curriculum_kwargs(args),
    )
    env = VecMonitor(env, info_keywords=("goal_scored",))

    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",
    )
    goal_rate_callback = GoalRateCallback()

    try:
        model.learn(
            args.timesteps,
            callback=[checkpoint_callback, goal_rate_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()