mirror of
https://github.com/jcreek/CosmicClash.git
synced 2026-09-10 16:04:04 +00:00
Josh Creek
602fa297d0
air_approach_weight alone didn't move productive_air_touch_fraction after a further 180M steps (360M cumulative across all six Stage-5 attempts): an unredirected air-intercept ball falls short of the goal from gravity and just lands on the floor, so the already-solved ground game collects the same episode reward whether or not anything touched the ball in the air. air_touch_bonus_weight adds a conjunctive event bonus on top of ball_touch_reward for a touch that's both genuinely aerial and goal-directed, targeting the actual measured behaviour instead of only the approach to it.
552 lines
25 KiB
Python
552 lines
25 KiB
Python
"""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 gymnasium import spaces
|
|
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"
|
|
|
|
# Must match Game/scripts/ship_action_codec.gd's HEADS order exactly (both
|
|
# are independently the gymnasium-sorted key order of the same 7 names) —
|
|
# training/test_action_space.py's rung-0 check asserts this. Used only for
|
|
# per-head entropy logging/reset-logits head selection below.
|
|
ACTION_HEAD_NAMES = ["rot_x", "rot_y", "rot_z", "thrust_x", "thrust_y", "thrust_z", "turbo"]
|
|
|
|
|
|
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))
|
|
|
|
|
|
class FlightTelemetryCallback(BaseCallback):
|
|
"""Logs flight and handling telemetry — leading indicators for curriculum generation 4's
|
|
core hypothesis (a discrete action space lets the policy actually hold a
|
|
sustained vertical set-point, e.g. hovering), visible from the very
|
|
first rollout instead of only in a win-rate number measured a full
|
|
24h+ run later, which is what made every past generation's failure mode
|
|
expensive to diagnose. Requires VecMonitor(..., info_keywords=(...,
|
|
"airborne_fraction", "mean_altitude", "air_touch_fraction",
|
|
"vertical_thrust_mean")) — see ShipAIController.get_info.
|
|
|
|
grounded_upright_fraction is diagnostic only and deliberately ungated:
|
|
upright_fraction's denominator is sub-3m ticks, most of which are
|
|
ballistic transit rather than driving, so it understates uprightness
|
|
while the ship is actually on the floor. Adding a key here requires
|
|
adding it to VecMonitor's info_keywords below too, or the bare
|
|
info[key] lookup raises KeyError mid-run."""
|
|
|
|
_KEYS = (
|
|
"airborne_fraction",
|
|
"mean_altitude",
|
|
"air_touch_fraction",
|
|
"vertical_thrust_mean",
|
|
"productive_air_touch_fraction",
|
|
"upright_fraction",
|
|
"forward_motion_fraction",
|
|
"grounded_upright_fraction",
|
|
)
|
|
|
|
def _on_step(self) -> bool:
|
|
return True
|
|
|
|
def _on_rollout_end(self) -> None:
|
|
if len(self.model.ep_info_buffer) == 0:
|
|
return
|
|
for key in self._KEYS:
|
|
values = [ep_info[key] for ep_info in self.model.ep_info_buffer if key in ep_info]
|
|
if values:
|
|
self.logger.record(f"rollout/{key}", safe_mean(values))
|
|
|
|
|
|
class EntropyFloorCallback(BaseCallback):
|
|
"""Replaces the old one-shot `--reset-std` shock (meaningless under
|
|
MultiDiscrete — there is no log_std) with a persistent controller.
|
|
Three curriculum generations' TensorBoard runs all show the same
|
|
signature: exploration (train/std, under the previous continuous
|
|
Gaussian) collapsing within the first ~10% of steps and never
|
|
recovering from a single reset applied at attempt start. A controller
|
|
that responds every rollout instead of once should not have that decay-
|
|
and-stay-collapsed failure mode.
|
|
|
|
Reads mean policy entropy each rollout (recomputed from a fresh
|
|
minibatch via the same RolloutBuffer.get() plumbing PPO's own train()
|
|
uses, since _on_rollout_end fires before that iteration's train() call)
|
|
and nudges model.ent_coef multiplicatively toward a target that decays
|
|
linearly from target_start_frac to target_end_frac of the action
|
|
space's maximum possible entropy (sum of ln(n) over each MultiDiscrete
|
|
head) over the run. PPO reads self.ent_coef fresh inside train() each
|
|
update, so mutating it here from a callback takes effect on the very
|
|
next update with no subclassing needed. No-ops (does nothing) for a
|
|
non-MultiDiscrete action space, e.g. a continuous-action A/B run.
|
|
|
|
Also logs train/entropy_head_<name> per action head — the direct
|
|
replacement for the old aggregate train/std scalar, and strictly more
|
|
useful: it identifies *which* axis is collapsing instead of one number
|
|
for all seven.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
total_timesteps: int,
|
|
target_start_frac: float = 0.55,
|
|
target_end_frac: float = 0.20,
|
|
adjust_rate: float = 1.02,
|
|
ent_coef_bounds: tuple[float, float] = (1e-4, 0.05),
|
|
):
|
|
super().__init__()
|
|
self.total_timesteps = total_timesteps
|
|
self.target_start_frac = target_start_frac
|
|
self.target_end_frac = target_end_frac
|
|
self.adjust_rate = adjust_rate
|
|
self.ent_coef_bounds = ent_coef_bounds
|
|
self._is_multi_discrete = False
|
|
self._h_max = 0.0
|
|
self._start_timesteps = 0
|
|
|
|
def _on_training_start(self) -> None:
|
|
import numpy as np
|
|
|
|
self._is_multi_discrete = isinstance(self.model.action_space, spaces.MultiDiscrete)
|
|
if self._is_multi_discrete:
|
|
self._h_max = float(np.sum(np.log(self.model.action_space.nvec)))
|
|
# this call's own timesteps budget, not the resumed total — model.
|
|
# num_timesteps keeps accumulating across --resume calls, but
|
|
# total_timesteps below is this invocation's --timesteps.
|
|
self._start_timesteps = self.model.num_timesteps
|
|
|
|
def _on_step(self) -> bool:
|
|
return True
|
|
|
|
def _on_rollout_end(self) -> None:
|
|
if not self._is_multi_discrete:
|
|
return
|
|
import torch as th
|
|
|
|
batch = next(self.model.rollout_buffer.get(batch_size=self.model.batch_size))
|
|
with th.no_grad():
|
|
distribution = self.model.policy.get_distribution(batch.observations)
|
|
per_head = getattr(distribution, "distribution", None)
|
|
if per_head is None:
|
|
return
|
|
|
|
entropies = [dist.entropy().mean().item() for dist in per_head]
|
|
for name, entropy in zip(ACTION_HEAD_NAMES, entropies):
|
|
self.logger.record(f"train/entropy_head_{name}", entropy)
|
|
|
|
mean_entropy = sum(entropies)
|
|
progress = min((self.model.num_timesteps - self._start_timesteps) / self.total_timesteps, 1.0)
|
|
target_frac = self.target_start_frac + (self.target_end_frac - self.target_start_frac) * progress
|
|
target = target_frac * self._h_max
|
|
if mean_entropy < target:
|
|
self.model.ent_coef = min(self.model.ent_coef * self.adjust_rate, self.ent_coef_bounds[1])
|
|
else:
|
|
self.model.ent_coef = max(self.model.ent_coef / self.adjust_rate, self.ent_coef_bounds[0])
|
|
self.logger.record("train/ent_coef_adaptive", self.model.ent_coef)
|
|
|
|
|
|
class AbortIfCallback(BaseCallback):
|
|
"""Optional kill criterion (see curriculum.py's per-stage `abort_if`):
|
|
ends model.learn() early once `metric` (a rollout/* key logged by
|
|
FlightTelemetryCallback — must run earlier in the callback list so the
|
|
value exists by the time this checks it) is below `below` at or past
|
|
`at_steps`. Stops via SB3's own "_on_step returning False halts
|
|
training" contract rather than an exception, so the enclosing
|
|
try/finally in main() still runs and saves/exports/commits whatever
|
|
checkpoint exists — an aborted stage still leaves a usable, logged
|
|
artifact instead of either running a doomed stage to completion
|
|
unattended or leaving one stranded and uncommitted.
|
|
"""
|
|
|
|
def __init__(self, metric: str, below: float, at_steps: int):
|
|
super().__init__()
|
|
self.metric = metric
|
|
self.below = below
|
|
self.at_steps = at_steps
|
|
self._checked = False
|
|
self._should_stop = False
|
|
|
|
def _on_step(self) -> bool:
|
|
return not self._should_stop
|
|
|
|
def _on_rollout_end(self) -> None:
|
|
if self._checked or self.model.num_timesteps < self.at_steps:
|
|
return
|
|
self._checked = True
|
|
value = self.logger.name_to_value.get(self.metric)
|
|
if value is not None and value < self.below:
|
|
print(
|
|
f"AbortIfCallback: {self.metric}={value:.4f} < {self.below} "
|
|
f"at {self.model.num_timesteps} steps — stopping early"
|
|
)
|
|
self._should_stop = True
|
|
|
|
|
|
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.01,
|
|
help="Entropy bonus coefficient (applied on resume too). Raised from 0.0001 for curriculum "
|
|
"generation 4: that value was tuned for a continuous Gaussian's differential entropy "
|
|
"(unbounded, can go negative); MultiDiscrete entropy is bounded (~10 nats for this action "
|
|
"space) and needs an order of magnitude more coefficient to matter. See --entropy-floor.",
|
|
)
|
|
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-logits",
|
|
type=float,
|
|
default=None,
|
|
help="On resume, multiply the policy's action_net weights/bias by this scale (e.g. 0.1), "
|
|
"pulling every head's softmax back toward uniform without discarding learned features — "
|
|
"the MultiDiscrete analogue of the old continuous --reset-std. Combine with "
|
|
"--reset-logits-heads to reset only specific heads.",
|
|
)
|
|
parser.add_argument(
|
|
"--reset-logits-heads",
|
|
default=None,
|
|
help=f"Comma-separated subset of {ACTION_HEAD_NAMES} to apply --reset-logits to (default: all heads)",
|
|
)
|
|
parser.add_argument(
|
|
"--entropy-floor",
|
|
action="store_true",
|
|
help="Enable EntropyFloorCallback: a persistent per-rollout controller nudging ent_coef to "
|
|
"hold policy entropy near a decaying target, replacing the one-shot --reset-std/"
|
|
"--reset-logits shock as the primary exploration mechanism (that flag remains for "
|
|
"resume-time recovery after a diagnosed collapse; this runs continuously).",
|
|
)
|
|
parser.add_argument(
|
|
"--checkpoint-every", type=int, default=10_000_000,
|
|
help="Timesteps between checkpoints. Raised from 100_000 for curriculum generation 4: at the "
|
|
"old value a single 240M-step stage wrote ~2400 intermediate checkpoint files (only final.zip "
|
|
"is ever committed, see .gitignore/run_training.sh, but they still accumulate in the working "
|
|
"tree during the run).",
|
|
)
|
|
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")
|
|
parser.add_argument(
|
|
"--abort-metric", default=None,
|
|
help="Optional kill criterion (see curriculum.py's per-stage abort_if): a rollout/* metric name to watch",
|
|
)
|
|
parser.add_argument("--abort-below", type=float, default=None, help="Stop early if --abort-metric drops below this")
|
|
parser.add_argument(
|
|
"--abort-at-steps", type=int, default=None,
|
|
help="Don't check --abort-metric until at least this many timesteps have elapsed",
|
|
)
|
|
|
|
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", "league"],
|
|
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); league: sample a fixed policy per episode "
|
|
"from --opponent-pool",
|
|
)
|
|
curriculum.add_argument("--opponent-model", default=None, help="Exported policy .json for --opponent-mode=frozen")
|
|
curriculum.add_argument(
|
|
"--opponent-pool", default=None,
|
|
help="Comma-separated exported policy paths for --opponent-mode=league; one is sampled per episode",
|
|
)
|
|
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(
|
|
"--air-drill-chance", type=float, default=None,
|
|
help="Overrides air_drill_chance: ball spawned high, both ships spawned low and lateral — "
|
|
"unsolvable without climbing (curriculum generation 4's state-setter aerial curriculum)",
|
|
)
|
|
curriculum.add_argument(
|
|
"--air-intercept-chance", type=float, default=None,
|
|
help="Moving high-ball interception starts aimed at a real goal (generation-5 aerial stage)",
|
|
)
|
|
curriculum.add_argument(
|
|
"--ground-start-chance", type=float, default=None,
|
|
help="Fraction of resets that spawn ships level and resting on the floor with a floor-level "
|
|
"ball — the state the handling stage's ground rewards are written for",
|
|
)
|
|
curriculum.add_argument(
|
|
"--team-size", type=int, choices=range(1, 6), default=None,
|
|
help="Ships per team (1-5); generation-5 automated stages remain 1v1 until 2v2 evaluation exists",
|
|
)
|
|
curriculum.add_argument(
|
|
"--tilt-penalty", type=float, default=None,
|
|
help="Overrides ShipAIController.tilt_penalty (dense per-tick cost scaled by non-upright tilt)",
|
|
)
|
|
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(
|
|
"--forward-velocity-to-ball-weight", type=float, default=None,
|
|
help="Low-altitude dense reward for nose-led planar approach toward the ball",
|
|
)
|
|
curriculum.add_argument(
|
|
"--air-approach-weight", type=float, default=None,
|
|
help="Aerial mirror of forward-velocity-to-ball-weight: high-altitude dense reward for "
|
|
"nose-led 3D closing speed toward the ball",
|
|
)
|
|
curriculum.add_argument(
|
|
"--air-touch-bonus-weight", type=float, default=None,
|
|
help="Event bonus on top of ball-touch-reward for a touch that is both genuinely aerial "
|
|
"(ball above AIR_TOUCH_HEIGHT) and goal-directed, scaled by the same alignment factor "
|
|
"as the base touch reward",
|
|
)
|
|
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(
|
|
"--ground-tilt-penalty", type=float, default=None,
|
|
help="Low-altitude-only tilt cost that fades to zero by the handling-height threshold",
|
|
)
|
|
curriculum.add_argument(
|
|
"--non-forward-penalty", type=float, default=None,
|
|
help="Overrides ShipAIController.non_forward_penalty (low-altitude dense cost on sideways/reverse "
|
|
"planar velocity, independent of the ball)",
|
|
)
|
|
curriculum.add_argument(
|
|
"--grounded-upright-reward", type=float, default=None,
|
|
help="Overrides ShipAIController.grounded_upright_reward (dense bonus for genuine floor contact "
|
|
"while upright, countering an incentive to just avoid the floor)",
|
|
)
|
|
curriculum.add_argument(
|
|
"--speed-reward-weight", type=float, default=None,
|
|
help="Overrides the orientation-agnostic own-speed reward (generation 5 handling sets it to zero)",
|
|
)
|
|
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,
|
|
"opponent_model_pool": args.opponent_pool,
|
|
"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,
|
|
"air_drill_chance": args.air_drill_chance,
|
|
"air_intercept_chance": args.air_intercept_chance,
|
|
"ground_start_chance": args.ground_start_chance,
|
|
"team_size": args.team_size,
|
|
"ai_tilt_penalty": args.tilt_penalty,
|
|
"ai_ground_tilt_penalty": args.ground_tilt_penalty,
|
|
"ai_non_forward_penalty": args.non_forward_penalty,
|
|
"ai_grounded_upright_reward": args.grounded_upright_reward,
|
|
"ai_velocity_to_ball_weight": args.velocity_to_ball_weight,
|
|
"ai_forward_velocity_to_ball_weight": args.forward_velocity_to_ball_weight,
|
|
"ai_air_approach_weight": args.air_approach_weight,
|
|
"ai_air_touch_bonus_weight": args.air_touch_bonus_weight,
|
|
"ai_ball_distance_penalty": args.ball_distance_penalty,
|
|
"ai_ball_touch_reward": args.ball_touch_reward,
|
|
"ai_airborne_penalty": args.airborne_penalty,
|
|
"ai_speed_reward_weight": args.speed_reward_weight,
|
|
"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",
|
|
"airborne_fraction",
|
|
"mean_altitude",
|
|
"air_touch_fraction",
|
|
"vertical_thrust_mean",
|
|
"productive_air_touch_fraction",
|
|
"upright_fraction",
|
|
"forward_motion_fraction",
|
|
"grounded_upright_fraction",
|
|
),
|
|
)
|
|
|
|
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_logits is not None:
|
|
import torch
|
|
|
|
heads = args.reset_logits_heads.split(",") if args.reset_logits_heads else ACTION_HEAD_NAMES
|
|
nvec = list(model.action_space.nvec)
|
|
offset = 0
|
|
offsets = {}
|
|
for name, size in zip(ACTION_HEAD_NAMES, nvec):
|
|
offsets[name] = (offset, offset + size)
|
|
offset += size
|
|
with torch.no_grad():
|
|
for name in heads:
|
|
start, end = offsets[name]
|
|
model.policy.action_net.weight[start:end].mul_(args.reset_logits)
|
|
model.policy.action_net.bias[start:end].mul_(args.reset_logits)
|
|
print(f"Reset action_net logits for heads {heads} by scale {args.reset_logits}")
|
|
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",
|
|
)
|
|
# Order matters for AbortIfCallback (must run after FlightTelemetryCallback
|
|
# so the rollout/* metric it watches has already been logged this round).
|
|
callbacks = [checkpoint_callback, GoalRateCallback(), FlightTelemetryCallback()]
|
|
if args.entropy_floor:
|
|
callbacks.append(EntropyFloorCallback(total_timesteps=args.timesteps))
|
|
if args.abort_metric is not None and args.abort_below is not None and args.abort_at_steps is not None:
|
|
callbacks.append(AbortIfCallback(args.abort_metric, args.abort_below, args.abort_at_steps))
|
|
|
|
try:
|
|
model.learn(
|
|
args.timesteps,
|
|
callback=callbacks,
|
|
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()
|