"""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( "--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( "--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)", ) 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_allow_vertical": args.allow_vertical, "ai_allow_pitch_roll": args.allow_pitch_roll, "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) 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()