feat(training): curriculum generation 4 — MultiDiscrete action space redesign

Three curriculum generations (2026-07-21 through 2026-08-04) all tried
gating *when* the policy could use vertical thrust/pitch-roll on top of a
continuous Gaussian action space, and all three failed the same way: PPO's
action-distribution std collapsed within ~10% of steps and never recovered,
landing at a 15-32% win rate vs the grounded reference regardless of
mechanism (hard mask, then a gradual ramp). Generation 3's final attempt
just landed at 24% — the worst of the three.

Root cause, verified against this project's own physics: hovering this ship
requires *holding* thrust.y ~= 0.408 continuously (mass 5.0, vertical_thrust
120, gravity 9.8). A collapsed near-zero-mean Gaussian can brush that value
but never sustain it long enough to earn the reward gradient that would
move the mean — no amount of gating *when* the axis acts fixes a problem in
*how* the policy represents a decision on it. This also independently found
and fixes a real bug: godot_rl never marks an episode timeout as a
truncation, so PPO was bootstrapping V(s)=0 on every 30s draw in every
generation to date.

- Game/scripts/ship_action_codec.gd (new): single source of truth for a
  per-axis MultiDiscrete action space (7 heads, nvec [5,5,5,5,5,5,2]) shared
  by training and in-game inference, replacing the continuous Gaussian.
  thrust_y's bins are deliberately asymmetric so a random policy drifts
  through the volume instead of floor-pinning. Legacy continuous decode
  (ai_ship_controller.gd's old logic) preserved verbatim so every
  pre-generation-4 export (e.g. Game/bots/promoted/easy.json) keeps working
  unchanged via an optional "action_space" JSON field.
- ship_observations.gd: append own contact state (SIZE 31 -> 35, append-only)
  so the value function can see what wall_contact_penalty fires on.
- ship_ai_controller.gd: action space/decode via the codec; drop the
  vertical_ramp/pitch_roll_ramp mechanism entirely; tilt_penalty default
  lowered 4x (aerial approaches require pitching); flight telemetry
  (airborne_fraction, mean_altitude, air_touch_fraction, vertical_thrust_mean)
  and truncation-snapshot fields on get_info().
- training_mode.gd: new air_drill_chance state-setter branch (ball spawned
  high, ships low, kept clear of walls) so aerial practice is forced by the
  environment instead of relying on reward-driven exploration alone; snapshot
  terminal observations before a timeout reset for the truncation fix.
- cosmic_env.py: remap ShipAIController's truncated/terminal_obs info into
  SB3's TimeLimit.truncated/terminal_observation keys.
- train.py: --reset-logits (+ --reset-logits-heads) replaces the
  now-meaningless --reset-std; new EntropyFloorCallback (a persistent
  per-rollout ent_coef controller replacing the one-shot std-reset shock)
  and per-head entropy logging; FlightTelemetryCallback; --air-drill-chance/
  --tilt-penalty flags; optional AbortIfCallback kill-criterion.
- export_policy.py: writes the action_space block for MultiDiscrete models;
  index-level parity check (argmax per head) instead of comparing floats.
- curriculum.py: full rewrite — 3 stages (bootstrap/selfplay/gauntlet), no
  grounded stage, full action space live from step 1; deletes generation
  1-3's checkpoint-lineage machinery (nothing to resume from); final report
  evaluates against both promoted/easy.json and the new
  promoted/reference-grounded.json (a copy of curric-s5-aggression, the
  strongest grounded-era artifact, kept as a fixed yardstick).
- run_training.sh/.gitignore: commit only final.zip, not the ~2400
  intermediate checkpoint files a single stage was writing (~500MB ->
  ~0.2MB per run); requirements.txt pinned (behaviour here now depends on
  specific library internals, not just public APIs).
- test_action_space.py (new): offline rung-0 check catching a head-order
  mismatch before it silently corrupts 24h of training.

Validated: GDScript compiles clean (Godot --headless --import + script
validation), free_play.tscn and training.tscn both boot headless without
errors, offline action-space assertions pass. Not yet run: the actual
smoke-training/A-B validation ladder steps in TRAINING.md's "Generation 4"
section, before committing to the full ~32h curriculum.

See TRAINING.md's "Generation 4" section for the full design writeup.

training/cosmic_env.py

@@ -10,6 +10,7 @@ learning policy: self-play by construction.
 import pathlib
 import subprocess
 
+import numpy as np
 from godot_rl.core.godot_env import GodotEnv
 from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv
 
@@ -72,8 +73,13 @@ class CosmicClashEnv(GodotEnv):
 class CosmicClashVecEnv(StableBaselinesGodotEnv):
     """SB3 VecEnv over N parallel CosmicClashEnv instances.
 
-    convert_action_space=True flattens the env's (Box(6), Discrete(2)) action
-    space into a single Box(7): thrust xyz, rotation xyz, turbo (>0 means on).
+    convert_action_space=True: godot_rl's ActionSpaceProcessor reports a
+    gym.spaces.MultiDiscrete when every per-axis action entry is Discrete
+    (see ShipActionCodec/ShipAIController.get_action_space) — nvec
+    [5,5,5,5,5,5,2] for rotation xyz, thrust xyz, turbo, in that
+    gymnasium-sorted key order. No conversion logic here needs to change for
+    that; this class's only functional addition is the truncation-info
+    remap below.
     """
 
     def __init__(self, godot_bin: str, n_parallel: int = 1, seed: int = 0, port: int = GodotEnv.DEFAULT_PORT, **kwargs):
@@ -90,3 +96,27 @@ class CosmicClashVecEnv(StableBaselinesGodotEnv):
         self.n_parallel = n_parallel
         self._check_valid_action_space()
         self.results = None
+
+    def step(self, action):
+        """Remap ShipAIController.get_info()'s "truncated"/"terminal_obs"
+        into the keys SB3's on_policy_algorithm looks for
+        ("TimeLimit.truncated"/"terminal_observation") so PPO bootstraps
+        V(s) through an episode timeout instead of treating every 30s draw
+        as a true terminal state.
+
+        Godot_rl's own godot_env.py never sets either key (it returns the
+        same `done` array for both term and trunc, "# TODO update API to
+        term, trunc") and StableBaselinesGodotEnv.step() only ever returns
+        that single collapsed `dones` array to SB3 — so without this, PPO
+        has no way to distinguish "episode ended because a goal was scored"
+        (a genuine terminal, V(s)=0 is correct) from "episode ended because
+        the 30s clock ran out" (an artificial boundary that should be
+        bootstrapped through), and was silently treating every draw as the
+        former in every curriculum generation to date.
+        """
+        obs, rewards, dones, infos = super().step(action)
+        for info in infos:
+            if info.pop("truncated", False):
+                info["TimeLimit.truncated"] = True
+                info["terminal_observation"] = {"obs": np.array(info.pop("terminal_obs"), dtype=np.float32)}
+        return obs, rewards, dones, infos

training/curriculum.py

@@ -18,32 +18,49 @@ forever for the wrong reason. When a stage does fail MAX_RETRIES times in a
 row, the script stops and asks for a human look rather than retrying
 indefinitely or silently advancing past a bad stage.
 
-This is generation 3 of the curriculum. Generation 1 (6 stages: score,
-defend, no_draws, mechanics, aggression, unmask) ran 2026-07-21 through
-2026-07-26 and is archived in curriculum_state_gen1.json — its final stage
-("unmask", full 3D flight on top of the aggression retune) failed 3 straight
-attempts, monotonically worsening (25% -> 20% -> 15% win rate vs
-curric-s5-aggression) because every retry resumed the same drifting
-checkpoint under identical flags instead of actually changing anything.
-Generation 2 (archived in curriculum_state_gen2.json) started a fresh
-single "unmask" stage seeded directly from curric-s5-aggression's own
-checkpoint (FOUNDATION_EXPERIMENT below) with retuned reward weights — it
-also failed 3 attempts, landing at a stable 32% / 28% / 31% win rate each
-time, ruling out both "retune the reward weights" and "just give it more
-time" as fixes. Generation 3 replaces the single all-or-nothing unmask
-flip with a gradual ramp (4 stages: unmask-ramp25/50/75, then unmask at
-full authority) — see TRAINING.md's "Generation 3" section for the full
-postmortem and design.
+This is generation 4 of the curriculum — a full redesign, not a patch.
+Generations 1-3 (archived in curriculum_state_gen1.json/_gen2.json/_gen3.json)
+all tried teaching full 3D flight by training grounded first and then
+opening up vertical/pitch-roll authority (a hard 0/1 mask in gen 1/2, a
+gradual float ramp in gen 3) on top of a continuous Gaussian action space.
+All three failed: gen 1's hard mask went 25% -> 20% -> 15% win rate across 3
+attempts; gen 2's single-flip retune landed at a stable 32%/28%/31%; gen 3's
+gradual ramp landed at 29%/30%/24% — actually the worst of the three by its
+final attempt. Every attempt showed the same signature regardless of
+mechanism: PPO's Gaussian action-distribution std collapsed from ~0.30 to
+~0.13-0.15 within the first ~10% of steps and never recovered. The root
+cause: hovering this ship (mass 5.0, vertical_thrust 120, default gravity
+9.8 — see ship.gd) requires *holding* thrust.y ~= 0.408 continuously; a
+collapsed near-zero-mean Gaussian can brush that value but never sustain it
+long enough to earn the reward gradient that would move the mean. No amount
+of gating *when* the axis is allowed to act fixes a problem in *how* the
+policy represents a decision on it.
+
+Generation 4 (see TRAINING.md and Game/scripts/ship_action_codec.gd)
+replaces the action space itself with per-axis MultiDiscrete bins instead of
+a continuous Gaussian, trains the full action space from step 1 with no
+grounded stage at all (no successful self-play RL bot in this problem class
+gates control authority — see the RLGym/RLBot research cited in
+TRAINING.md), and adds a state-setter "air drill" episode-start branch
+(training_mode.gd's air_drill_chance) to force aerial practice instead of
+relying on reward-driven exploration alone. All of generation 1-3's
+checkpoint-lineage machinery (FOUNDATION_EXPERIMENT, locomotion-groundedness
+tracking, resume/reference overrides for skipping a regressed branch) is
+gone because there is nothing to resume from: every prior checkpoint is a
+different, incompatible action/observation shape. The two strongest prior
+artifacts are kept as fixed evaluation references instead (see
+PROMOTED_EASY/PROMOTED_REFERENCE_GROUNDED below) — they remain playable
+opponents forever via PolicyNetwork's format-versioned JSON even though
+their own checkpoints and generation are gone.
 
 Every experiment name this script generates is timestamped
 (YYYYMMDD-HHMM-<name>, applied once in run_stage_attempt) so runs stay
 unique across restarts/generations and sort chronologically in TensorBoard
-and checkpoints/ — plain names like "curric-s1-score" from generation 1
-would otherwise collide with generation 2's own stage 1.
+and checkpoints/.
 
 Usage:
     .venv/bin/python curriculum.py                  # run/resume the curriculum
-    .venv/bin/python curriculum.py --seed-checkpoint checkpoints/run11/final.zip
+    .venv/bin/python curriculum.py --seed-checkpoint checkpoints/some/final.zip
     .venv/bin/python curriculum.py --force-retry     # after fixing something, retry the blocked stage
     .venv/bin/python curriculum.py --skip-to-next-stage  # human judgment call: good enough, move on anyway
 
@@ -61,20 +78,14 @@ from datetime import datetime
 TRAINING_DIR = pathlib.Path(__file__).resolve().parent
 STATE_PATH = TRAINING_DIR / "curriculum_state.json"
 EVAL_HISTORY_PATH = TRAINING_DIR / "eval_history.json"
-ROOKIE_REFERENCE = TRAINING_DIR.parent / "Game" / "bots" / "rookie.json"
 
-# Generation 1's last cleanly-passing checkpoint (see curriculum_state_gen1.json)
-# — generations 2 and 3 both build on this directly instead of re-running
-# stages 1-5.
-FOUNDATION_EXPERIMENT = "curric-s5-aggression"
-
-# Groundedness (locomotion-mask state) for experiments that predate this
-# generation's own log, so _grounded_for_experiment can still answer for
-# them — see that function.
-LEGACY_GROUNDED = {
-    "rookie": False,
-    FOUNDATION_EXPERIMENT: True,
-}
+# Fixed evaluation references — never touched by training scripts (see
+# TRAINING.md's "Promoted bots" section) — kept forever as playable
+# opponents via PolicyNetwork's format-versioned JSON even after their own
+# checkpoints/generation are gone. The final report (not a gate) evaluates
+# generation 4's result against both.
+PROMOTED_EASY = TRAINING_DIR.parent / "Game" / "bots" / "promoted" / "easy.json"
+PROMOTED_REFERENCE_GROUNDED = TRAINING_DIR.parent / "Game" / "bots" / "promoted" / "reference-grounded.json"
 
 MAX_RETRIES = 2
 EVAL_EPISODES = 100
@@ -84,118 +95,108 @@ EVAL_EPISODES = 100
 # module docstring) — advance rather than retry.
 REGRESSION_MARGIN = 0.15
 
-# Standing flags applied to every attempt, mirroring next_run.sh: reset-std
-# reopens exploration every attempt (harmless on fresh starts — train.py
-# only applies it on --resume), ent-coef keeps it from re-collapsing.
-STANDING_ARGS = ["--reset-std", "0.3", "--ent-coef", "0.001"]
+# Standing flags applied to every attempt. Generation 3's "--reset-std 0.3"
+# (a one-shot shock, and meaningless anyway under MultiDiscrete — there is
+# no log_std) is gone; EntropyFloorCallback (see train.py) is a continuous
+# controller instead, which every generation's TensorBoard data argues is
+# what was actually needed (a single reset at attempt start reliably decayed
+# away within ~10% of steps, every time). --ent-coef raised an order of
+# magnitude from generation 3's 0.001: that value was tuned for a Gaussian's
+# unbounded differential entropy, not MultiDiscrete's bounded (~10-nat)
+# entropy.
+STANDING_ARGS = ["--ent-coef", "0.01", "--entropy-floor"]
 
-# Ball-chasing/scoring reward flags shared by every unmask-ramp stage
-# (generation 3 — see below): identical across all 4 stages so the ramp
-# itself is the only studied variable. Lifted from generation 2's single
-# "unmask" attempt (raised from stage 5's 0.05/0.006/0.5/0.02/60 defaults).
-_UNMASK_RAMP_SHARED_FLAGS = [
-    "--opponent-mode", "self_play",
+# Reward-shaping flags shared by every stage so the studied variables (state
+# mix, opponent mode) stay isolated — carried forward unchanged from
+# generation 2/3, which the reward-farmability analysis in TRAINING.md
+# confirmed were never the actual problem. draw_penalty and airborne_penalty
+# are deliberately NOT overridden here (both default to 0.0 in
+# training_mode.gd/ship_ai_controller.gd): generation 3's draw_penalty=5 and
+# airborne_penalty ramping up in lockstep with the unmask ramp were both
+# grounded-era, anti-flight pressures that have no place in a curriculum
+# whose entire point is teaching flight.
+_SHARED_REWARD_FLAGS = [
     "--velocity-to-ball-weight", "0.08",
     "--ball-distance-penalty", "0.01",
     "--ball-touch-reward", "0.7",
     "--ball-velocity-to-goal-weight", "0.06",
     "--goal-reward", "80",
-    "--draw-penalty", "5",
 ]
 
+# A stage dict may additionally set "abort_if": {"metric": "rollout/airborne_
+# fraction", "below": 0.05, "at_steps": N} to end that attempt early if a
+# flight-telemetry metric (see train.py's FlightTelemetryCallback) hasn't
+# cleared a bar by N *absolute* PPO timesteps (model.num_timesteps keeps
+# accumulating across --resume, so N must account for whatever this stage
+# inherits from its predecessor, not just this stage's own budget).
+# Deliberately unset on every stage below for now — rung 5 of TRAINING.md's
+# validation ladder (a short controlled A/B) should establish what a
+# sensible threshold actually looks like before any stage bets a real 12h+
+# budget on a guessed one.
 STAGES = [
     {
-        "name": "unmask-ramp25",
-        # Generation 3, step 1/4 of a gradual locomotion-unmask ramp — see
-        # TRAINING.md's "Generation 3" section for the full postmortem.
-        # Generation 2's single all-or-nothing "unmask" stage (flip
-        # vertical_ramp/pitch_roll_ramp 0 -> 1 in one step) failed 3
-        # independent 240M-step attempts in a row, landing at a stable
-        # 32% / 28% / 31% win rate vs curric-s5-aggression each time — not
-        # noise (attempts 2-3 each gave the *same* checkpoint lineage
-        # another full 240M steps with zero improvement) and not fixable by
-        # more time. Every attempt shows train/std collapsing from ~0.30 to
-        # ~0.13-0.15 within the first ~10% of steps and never recovering —
-        # the policy locks the newly-opened axes back down before ever
-        # meaningfully exploring them.
-        #
-        # This stage instead scales vertical_ramp/pitch_roll_ramp to 25%
-        # authority. Ungated (see "gated" below and main()'s loop): this is
-        # a waypoint, not a measured transition — no eval runs, no
-        # regression gate applies, it always advances after training.
-        # airborne_penalty is off (0.0) here: at 25% authority the axis
-        # barely does anything yet, so there's nothing to discourage.
+        "name": "bootstrap",
+        # Stage 1/3: empty-net finishing practice from a random policy — no
+        # live opponent, so the full action space's first behaviour to
+        # emerge is "fly to ball, push it toward the net" without a moving
+        # target complicating credit assignment. Generation 1's own stage 1
+        # (also inert-opponent, also empty-net) was the one stage across all
+        # 3 prior generations that unambiguously passed on its first
+        # attempt — reusing that shape here, just with the full action space
+        # live instead of yaw-only.
         "flags": [
-            *_UNMASK_RAMP_SHARED_FLAGS,
-            "--vertical-ramp", "0.25",
-            "--pitch-roll-ramp", "0.25",
-            "--airborne-penalty", "0.0",
+            "--opponent-mode", "inert",
+            "--attack-goal-bias", "1.0",
+            "--kickoff-chance", "0.10",
+            "--near-goal-chance", "0.50",
+            "--air-drill-chance", "0.20",
+            *_SHARED_REWARD_FLAGS,
         ],
-        "gated": False,
-        "timesteps": 40_000_000,  # ~4h at the standing n-parallel/speedup (20M took ~2h)
-        # Stage 0 MUST set this explicitly — resume_checkpoint()'s stage-0
-        # branch returns None (train from scratch) without it.
-        "resume_from_experiment": FOUNDATION_EXPERIMENT,
+        "gated": False,  # ungated waypoint: trains, checkpoints, always advances — no eval
+        "timesteps": 40_000_000,  # ~4h at the standing n-parallel/speedup
     },
     {
-        "name": "unmask-ramp50",
-        # Step 2/4: 50% authority. airborne_penalty at 1/3 of its final
-        # value — enough to start discouraging unproductive altitude, not
-        # enough to fight the still-partial vertical axis outright.
-        # resume_from_experiment deliberately omitted: chains from
-        # ramp25's pass via _resume_source_experiment's default
-        # (_passing_experiment_for_stage) — do not add an override here.
+        "name": "selfplay",
+        # Stage 2/3: this is where essentially all of the actual learning
+        # happens. Self-play (not frozen) as the main regime — it's what
+        # scales and what Necto/Nexto-class bots actually use; a frozen
+        # target this early would cap skill at "exploits one specific bot"
+        # instead of a moving, improving target. air_drill_chance stays on
+        # at a constant rate throughout (not introduced as a later stage) —
+        # gating *when* a skill is drilled reproduces the exact "gate what
+        # the policy is allowed to do" pattern that failed 3 generations in
+        # a row; only the state mix should vary between stages, never what
+        # the policy can act on.
         "flags": [
-            *_UNMASK_RAMP_SHARED_FLAGS,
-            "--vertical-ramp", "0.5",
-            "--pitch-roll-ramp", "0.5",
-            "--airborne-penalty", "0.001",
+            "--opponent-mode", "self_play",
+            "--kickoff-chance", "0.15",
+            "--near-goal-chance", "0.25",
+            "--air-drill-chance", "0.25",
+            *_SHARED_REWARD_FLAGS,
         ],
-        "gated": False,
-        "timesteps": 40_000_000,
+        "gated": True,
+        "timesteps": 160_000_000,  # ~16h
     },
     {
-        "name": "unmask-ramp75",
-        # Step 3/4: 75% authority, airborne_penalty at 2/3 of its final
-        # value. Also chains automatically — no resume_from_experiment.
+        "name": "gauntlet",
+        # Stage 3/3: a stationary opponent (this stage's own predecessor's
+        # export) gives a low-variance measurement — important when the gate
+        # is a 100-episode sample with a lenient 15-point margin — and
+        # catches a self-play fixed point: a policy that only learned to
+        # beat itself will look fine in stage 2 and stall here.
+        # opponent_model_from_previous_stage resolves --opponent-model at
+        # run time to whatever stage 2's own passing export turns out to be
+        # (see run_stage_attempt) rather than a hardcoded name.
         "flags": [
-            *_UNMASK_RAMP_SHARED_FLAGS,
-            "--vertical-ramp", "0.75",
-            "--pitch-roll-ramp", "0.75",
-            "--airborne-penalty", "0.002",
+            "--opponent-mode", "frozen",
+            "--kickoff-chance", "0.15",
+            "--near-goal-chance", "0.25",
+            "--air-drill-chance", "0.25",
+            *_SHARED_REWARD_FLAGS,
         ],
-        "gated": False,
-        "timesteps": 40_000_000,
-    },
-    {
-        "name": "unmask",
-        # Step 4/4, the measured transition: full ramp (100% authority),
-        # airborne_penalty at its full value — behaviourally and eval-wise
-        # identical to generation 2's "unmask" stage's flags/config, so
-        # this stage's result is a direct, apples-to-apples comparison
-        # against the 3 failed all-or-nothing attempts (same reference,
-        # same opponent mode, same budget). Gated (default True): evaluated
-        # against FOUNDATION_EXPERIMENT exactly like every prior attempt.
-        #
-        # No resume_from_experiment here (deliberately, unlike generation
-        # 2's single-stage version) — this stage chains from ramp75's own
-        # checkpoint via the default resume path, not back to
-        # FOUNDATION_EXPERIMENT; only reference_experiment (the *eval*
-        # opponent) stays FOUNDATION_EXPERIMENT.
-        "flags": [
-            *_UNMASK_RAMP_SHARED_FLAGS,
-            "--vertical-ramp", "1.0",
-            "--pitch-roll-ramp", "1.0",
-            "--airborne-penalty", "0.003",
-        ],
-        "grounded": False,
-        "timesteps": 240_000_000,  # unchanged from the 3 failed attempts — same budget for a controlled comparison
-        "reference_experiment": FOUNDATION_EXPERIMENT,
-        # On retry, reset to the clean ramp75 checkpoint rather than
-        # compounding a failed full-ramp attempt's own drift — mirrors the
-        # generation 1 -> 2 postmortem (blind same-checkpoint retries only
-        # made things worse).
-        "reset_retry_checkpoint": True,
+        "gated": True,
+        "timesteps": 120_000_000,  # ~12h
+        "opponent_model_from_previous_stage": True,
     },
 ]
 
@@ -227,50 +228,6 @@ def _logged_experiment_name(stage_index: int, attempt: int) -> str:
     raise RuntimeError(f"No logged experiment for stage {stage_index} attempt {attempt}")
 
 
-def resume_checkpoint(stage_index: int, attempt: int, seed_checkpoint: str | None) -> str | None:
-    if attempt > 0 and not STAGES[stage_index].get("reset_retry_checkpoint"):
-        # Retry: keep training the same stage's own last attempt.
-        prev = _logged_experiment_name(stage_index, attempt - 1)
-        return str(TRAINING_DIR / "checkpoints" / prev / "final.zip")
-    if stage_index == 0 and seed_checkpoint:
-        return seed_checkpoint
-    if stage_index == 0 and not STAGES[0].get("resume_from_experiment"):
-        # Deliberately fresh by default: the curriculum exists because
-        # resuming self-play across a regime change (run10, run11) didn't
-        # work, so a from-scratch stage 1 starts from a random policy under
-        # its own regime unless --seed-checkpoint or resume_from_experiment
-        # says otherwise.
-        return None
-    # Either a later stage chaining off its predecessor, or
-    # reset_retry_checkpoint: this stage's own retries have been drifting
-    # rather than converging (see the "unmask" stage's comment) — resume
-    # from the stage's normal resume source instead of compounding the last
-    # failed attempt's drift.
-    prev_experiment = _resume_source_experiment(stage_index)
-    return str(TRAINING_DIR / "checkpoints" / prev_experiment / "final.zip")
-
-
-def reference_bot(stage_index: int) -> str:
-    if stage_index == 0 and not STAGES[0].get("reference_experiment"):
-        return str(ROOKIE_REFERENCE)
-    prev_experiment = _reference_source_experiment(stage_index)
-    return str(TRAINING_DIR.parent / "Game" / "bots" / f"{prev_experiment}.json")
-
-
-# A stage normally chains off "whatever passed at the previous index," but a
-# stage can instead name an explicit resume_from_experiment/reference_experiment
-# to skip a since-regressed branch, or (stage 0) to seed from a fixed
-# foundation checkpoint instead of a from-scratch policy.
-def _resume_source_experiment(stage_index: int) -> str:
-    override = STAGES[stage_index].get("resume_from_experiment")
-    return override if override else _passing_experiment_for_stage(stage_index - 1)
-
-
-def _reference_source_experiment(stage_index: int) -> str:
-    override = STAGES[stage_index].get("reference_experiment")
-    return override if override else _passing_experiment_for_stage(stage_index - 1)
-
-
 def _passing_experiment_for_stage(stage_index: int) -> str:
     state = load_state()
     for entry in state["log"]:
@@ -279,14 +236,31 @@ def _passing_experiment_for_stage(stage_index: int) -> str:
     raise RuntimeError(f"No passing attempt recorded for stage {stage_index} ({STAGES[stage_index]['name']})")
 
 
-def _grounded_for_experiment(experiment: str) -> bool:
-    if experiment in LEGACY_GROUNDED:
-        return LEGACY_GROUNDED[experiment]
-    state = load_state()
-    for entry in state["log"]:
-        if entry["experiment"] == experiment:
-            return STAGES[entry["stage_index"]]["grounded"]
-    raise ValueError(f"Unknown experiment for groundedness lookup: {experiment}")
+def resume_checkpoint(stage_index: int, attempt: int, seed_checkpoint: str | None) -> str | None:
+    if attempt > 0:
+        # Retry: keep training the same stage's own last attempt. No
+        # per-stage "reset to a clean upstream checkpoint" override in
+        # generation 4 (unlike generation 3's "unmask" stage) — nothing yet
+        # suggests a generation-4 retry needs that; add one if a stage's
+        # retries turn out to be drifting rather than converging.
+        prev = _logged_experiment_name(stage_index, attempt - 1)
+        return str(TRAINING_DIR / "checkpoints" / prev / "final.zip")
+    if stage_index == 0:
+        # Deliberately fresh unless --seed-checkpoint says otherwise: full
+        # action space live from step 1, nothing to inherit — every prior
+        # generation's checkpoints are a different, incompatible
+        # action/observation shape (see module docstring).
+        return seed_checkpoint
+    prev_experiment = _passing_experiment_for_stage(stage_index - 1)
+    return str(TRAINING_DIR / "checkpoints" / prev_experiment / "final.zip")
+
+
+def reference_bot(stage_index: int) -> str:
+    """Only called for gated stages (stage 0 is ungated) — the previous
+    stage's own passing export, exactly like every prior generation's
+    default chaining."""
+    prev_experiment = _passing_experiment_for_stage(stage_index - 1)
+    return str(TRAINING_DIR.parent / "Game" / "bots" / f"{prev_experiment}.json")
 
 
 def run_stage_attempt(stage_index: int, attempt: int, args) -> str:
@@ -294,9 +268,6 @@ def run_stage_attempt(stage_index: int, attempt: int, args) -> str:
     # chronologically in TensorBoard/checkpoints — see module docstring.
     exp = f"{datetime.now().strftime('%Y%m%d-%H%M')}-{experiment_name(stage_index, attempt)}"
     resume = resume_checkpoint(stage_index, attempt, args.seed_checkpoint)
-    # A stage can override the run's timesteps budget (see "floor-lock",
-    # which deliberately runs much longer than the ~20M/~2h every stage so
-    # far has used); otherwise it falls back to curriculum.py's own --timesteps.
     timesteps = STAGES[stage_index].get("timesteps", args.timesteps)
     cmd = [
         "./run_training.sh", exp,
@@ -308,6 +279,17 @@ def run_stage_attempt(stage_index: int, attempt: int, args) -> str:
     ]
     if resume:
         cmd += ["--resume", resume]
+    if STAGES[stage_index].get("opponent_model_from_previous_stage"):
+        prev_experiment = _passing_experiment_for_stage(stage_index - 1)
+        opponent_model = TRAINING_DIR.parent / "Game" / "bots" / f"{prev_experiment}.json"
+        cmd += ["--opponent-model", str(opponent_model)]
+    abort_if = STAGES[stage_index].get("abort_if")
+    if abort_if:
+        cmd += [
+            "--abort-metric", abort_if["metric"],
+            "--abort-below", str(abort_if["below"]),
+            "--abort-at-steps", str(abort_if["at_steps"]),
+        ]
     print(f"\n=== Stage {stage_index + 1}/{len(STAGES)} ({STAGES[stage_index]['name']}), "
           f"attempt {attempt + 1}/{MAX_RETRIES + 1}: {exp} ===")
     print(" ".join(cmd))
@@ -315,22 +297,9 @@ def run_stage_attempt(stage_index: int, attempt: int, args) -> str:
     return exp
 
 
-def reference_grounded(stage_index: int) -> bool:
-    if stage_index == 0 and not STAGES[0].get("reference_experiment"):
-        # rookie.json predates the locomotion mask entirely — always full 3D.
-        return False
-    return _grounded_for_experiment(_reference_source_experiment(stage_index))
-
-
-def evaluate_attempt(experiment: str, reference: str, episodes: int, stage_index: int) -> dict:
+def evaluate_attempt(experiment: str, reference: str, episodes: int) -> dict:
     candidate = TRAINING_DIR.parent / "Game" / "bots" / f"{experiment}.json"
     cmd = [".venv/bin/python", "evaluate.py", str(candidate), reference, "--episodes", str(episodes)]
-    # Must match how each side was actually trained — see ai_ship_controller.gd's
-    # allow_vertical/allow_pitch_roll and evaluate.py's --grounded-a/-b.
-    if STAGES[stage_index]["grounded"]:
-        cmd.append("--grounded-a")
-    if reference_grounded(stage_index):
-        cmd.append("--grounded-b")
     print(" ".join(cmd))
     subprocess.run(cmd, cwd=TRAINING_DIR, check=True)
     history = json.loads(EVAL_HISTORY_PATH.read_text())
@@ -345,6 +314,33 @@ def decide(record: dict) -> str:
     return "pass"
 
 
+def final_report(experiment: str) -> None:
+    """Not a gate — the two numbers that actually answer "did generation 4
+    work?" (see TRAINING.md). promoted/easy.json is the shipped bot;
+    promoted/reference-grounded.json (a copy of generation 3's
+    curric-s5-aggression, made before the flat Game/bots/ dump was scrapped)
+    is the strongest grounded-era artifact and the yardstick generations 1-3
+    were all measured against. reference-grounded.json was trained with the
+    locomotion mask on, so needs --grounded-b; easy.json was itself promoted
+    from a *failed* unmask stage (curric-s6-unmask) and is full 3D like
+    every generation-4 candidate, so needs no flag."""
+    candidate = TRAINING_DIR.parent / "Game" / "bots" / f"{experiment}.json"
+    print("\n=== Curriculum complete — final report (informational, not a gate) ===")
+    for label, reference, extra_flags in [
+        ("promoted/easy.json (shipped bot)", PROMOTED_EASY, []),
+        ("promoted/reference-grounded.json (strongest grounded-era bot)", PROMOTED_REFERENCE_GROUNDED, ["--grounded-b"]),
+    ]:
+        if not reference.exists():
+            print(f"  vs {label}: skipped, file not found")
+            continue
+        cmd = [".venv/bin/python", "evaluate.py", str(candidate), str(reference), "--episodes", str(EVAL_EPISODES), *extra_flags]
+        print(" ".join(cmd))
+        subprocess.run(cmd, cwd=TRAINING_DIR, check=True)
+        record = json.loads(EVAL_HISTORY_PATH.read_text())[-1]
+        print(f"  vs {label}: {record['wins_a']}-{record['wins_b']} ({record['draws']} draws), "
+              f"win rate {record['win_rate_a']:.0%}")
+
+
 def commit_progress(experiment: str) -> None:
     subprocess.run(["git", "add", "curriculum_state.json", "eval_history.json"], cwd=TRAINING_DIR, check=True)
     result = subprocess.run(["git", "diff", "--cached", "--quiet"], cwd=TRAINING_DIR)
@@ -364,8 +360,7 @@ def main():
     parser.add_argument("--speedup", type=int, default=16)
     parser.add_argument(
         "--seed-checkpoint", default=None,
-        help="Resume stage 1 from this checkpoint instead of its default resume source "
-        "(FOUNDATION_EXPERIMENT's checkpoint)",
+        help="Resume stage 1 from this checkpoint instead of training from scratch",
     )
     parser.add_argument("--force-retry", action="store_true", help="Retry a blocked stage after human review")
     parser.add_argument("--skip-to-next-stage", action="store_true", help="Human judgment call: treat the blocked stage as good enough, advance anyway")
@@ -413,14 +408,13 @@ def main():
         last_experiment = experiment
 
         if not STAGES[stage_index].get("gated", True):
-            # Ungated ramp waypoint (see the unmask-ramp2X stages): trains,
+            # Ungated waypoint (see the bootstrap stage): trains,
             # checkpoints, and always advances — no eval, no regression
-            # gate, nothing to retry against. See TRAINING.md's
-            # "Generation 3" section.
-            print(f"{experiment}: ungated ramp waypoint — skipping eval, advancing unconditionally")
+            # gate, nothing to retry against.
+            print(f"{experiment}: ungated waypoint — skipping eval, advancing unconditionally")
             state["log"].append({
                 "stage_index": stage_index, "experiment": experiment, "attempt": attempt,
-                "decision": "pass", "note": "ungated ramp waypoint (no eval)",
+                "decision": "pass", "note": "ungated waypoint (no eval)",
             })
             state["stage_index"] += 1
             state["attempt"] = 0
@@ -430,7 +424,7 @@ def main():
             continue
 
         reference = reference_bot(stage_index)
-        record = evaluate_attempt(experiment, reference, EVAL_EPISODES, stage_index)
+        record = evaluate_attempt(experiment, reference, EVAL_EPISODES)
         decision = decide(record)
 
         print(f"{experiment}: candidate {record['wins_a']}-{record['wins_b']} reference "
@@ -469,6 +463,7 @@ def main():
         # None only if the loop above never ran at all (e.g. re-invoking
         # after the curriculum was already "done") — nothing new to commit
         # in that case.
+        final_report(last_experiment)
         commit_progress(last_experiment)
 
 

training/eval_history.json

@@ -178,5 +178,15 @@
     "wins_b": 56,
     "draws": 14,
     "win_rate_a": 0.3
+  },
+  {
+    "timestamp": "2026-08-04T21:26:41+00:00",
+    "model_a": "/Users/jcreek/Documents/repos/GitHub/CosmicClash/Game/bots/20260803-1829-curric-s4-unmask-retry2.json",
+    "model_b": "/Users/jcreek/Documents/repos/GitHub/CosmicClash/Game/bots/curric-s5-aggression.json",
+    "episodes": 100,
+    "wins_a": 24,
+    "wins_b": 56,
+    "draws": 20,
+    "win_rate_a": 0.24
   }
 ]

training/export_policy.py

@@ -1,9 +1,15 @@
 """Export a trained SB3 checkpoint to the JSON format PolicyNetwork.gd loads.
 
-The exported file contains the deterministic policy MLP (obs -> action means);
-the game clamps outputs to [-1, 1] and treats the last value as turbo (> 0).
-A parity self-check compares the JSON forward pass against SB3's own
-deterministic prediction before writing.
+The exported file contains the deterministic policy MLP. For a MultiDiscrete
+(curriculum generation 4+) model, the raw output is 32 per-head logits
+decoded via ShipActionCodec.from_logits (argmax per head, mapped through
+ACTION_HEADS' bin values below) and an "action_space" block is written to
+the JSON so the game knows to decode it that way. For an older continuous
+model, output is 7 action means, clamped to [-1, 1] and the last value
+treated as turbo (> 0) — no "action_space" block, matching every export
+before generation 4 (e.g. Game/bots/promoted/easy.json). A parity self-check
+compares the JSON forward pass against SB3's own deterministic prediction
+before writing, in either case.
 
 Example:
     .venv/bin/python export_policy.py checkpoints/smoke/final.zip ../Game/bots/rookie.json
@@ -13,10 +19,28 @@ import argparse
 import json
 import pathlib
 
+import gymnasium as gym
 import numpy as np
 import torch
 from stable_baselines3 import PPO
 
+# MUST exactly match Game/scripts/ship_action_codec.gd's HEADS (name, order,
+# and bin values) — this is what gets written into every generation-4
+# export's "action_space" block, and PolicyNetwork.gd/AIShipController never
+# re-derive it, they just decode against whatever's in the file. Sizes are
+# cross-checked against the live model's action_space.nvec below (a real
+# assertion), but bin *values* have no automated cross-language check —
+# treat any edit to either file as requiring the other.
+ACTION_HEADS = [
+    {"name": "rot_x", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
+    {"name": "rot_y", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
+    {"name": "rot_z", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
+    {"name": "thrust_x", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
+    {"name": "thrust_y", "bins": [-0.5, 0.0, 0.45, 0.75, 1.0]},
+    {"name": "thrust_z", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
+    {"name": "turbo", "bins": [0.0, 1.0]},
+]
+
 
 def linear_to_layer(linear: torch.nn.Linear, activation: str) -> dict:
     return {
@@ -66,20 +90,55 @@ def main():
     policy = model.policy
     layers = extract_layers(policy)
     input_size = model.observation_space["obs"].shape[0]
+    is_multi_discrete = isinstance(model.action_space, gym.spaces.MultiDiscrete)
 
-    # Parity check: JSON forward pass must match SB3's deterministic action
+    output_data = {"input_size": int(input_size), "layers": layers}
     rng = np.random.default_rng(0)
-    for _ in range(16):
-        obs = rng.uniform(-1, 1, input_size).astype(np.float32)
-        expected, _ = model.predict({"obs": obs}, deterministic=True)
-        actual = np.clip(json_forward(layers, obs), -1.0, 1.0)
-        assert np.allclose(actual, expected, atol=1e-5), f"parity check failed: {actual} vs {expected}"
+
+    if is_multi_discrete:
+        head_sizes = [len(head["bins"]) for head in ACTION_HEADS]
+        nvec = [int(n) for n in model.action_space.nvec]
+        assert nvec == head_sizes, (
+            f"model action_space.nvec {nvec} doesn't match ACTION_HEADS sizes {head_sizes} — "
+            "update ACTION_HEADS to match ship_action_codec.gd's HEADS"
+        )
+        output_data["action_space"] = {"type": "multi_discrete", "heads": ACTION_HEADS}
+
+        # Index-level parity check: deterministic=True now returns one
+        # argmax index per head (not a float to clip), so compare argmax of
+        # the JSON forward pass's raw logits, sliced per head, against SB3's
+        # own chosen indices — a head-order mistake here would otherwise
+        # train/export cleanly and only surface as silently wrong in-game
+        # behaviour (e.g. pitch commands driving strafe thrusters).
+        offsets = []
+        running = 0
+        for size in head_sizes:
+            offsets.append((running, running + size))
+            running += size
+        for _ in range(16):
+            obs = rng.uniform(-1, 1, input_size).astype(np.float32)
+            expected, _ = model.predict({"obs": obs}, deterministic=True)
+            logits = json_forward(layers, obs)
+            actual = np.array([int(np.argmax(logits[start:end])) for start, end in offsets])
+            assert np.array_equal(actual, expected), f"parity check failed: {actual} vs {expected}"
+    else:
+        # Legacy continuous parity check, unchanged: JSON forward pass must
+        # match SB3's deterministic action mean.
+        for _ in range(16):
+            obs = rng.uniform(-1, 1, input_size).astype(np.float32)
+            expected, _ = model.predict({"obs": obs}, deterministic=True)
+            actual = np.clip(json_forward(layers, obs), -1.0, 1.0)
+            assert np.allclose(actual, expected, atol=1e-5), f"parity check failed: {actual} vs {expected}"
 
     output = pathlib.Path(args.output)
     output.parent.mkdir(parents=True, exist_ok=True)
     with open(output, "w") as f:
-        json.dump({"input_size": int(input_size), "layers": layers}, f)
-    print(f"Exported {args.checkpoint} -> {output} (input size {input_size}, {len(layers)} layers, parity OK)")
+        json.dump(output_data, f)
+    action_space_label = "multi_discrete" if is_multi_discrete else "continuous"
+    print(
+        f"Exported {args.checkpoint} -> {output} (input size {input_size}, {len(layers)} layers, "
+        f"action_space={action_space_label}, parity OK)"
+    )
 
 
 if __name__ == "__main__":

training/requirements.txt

@@ -1,5 +1,19 @@
-godot-rl
-stable-baselines3
-tensorboard
+# Pinned as of curriculum generation 4: this codebase now depends on
+# specific library internals (godot_rl's ActionSpaceProcessor discrete-only
+# -> MultiDiscrete branch, stable_baselines3's MultiCategoricalDistribution
+# logit layout — see Game/scripts/ship_action_codec.gd and train.py), not
+# just documented public APIs. An unpinned reinstall (e.g. via
+# setup_linux.sh on the remote training box) could silently resolve a newer
+# version that changes that behaviour without any error, which would be a
+# very expensive thing to discover partway through a 12h+ curriculum stage.
+# torch/gymnasium are pinned too even though they're transitive deps of the
+# two above, for the same reason (torch's log_std/action_net tensor
+# shapes, gymnasium's Dict space key-sorting behaviour that
+# ShipActionCodec's HEADS order relies on).
+godot-rl==0.8.2
+stable-baselines3==2.4.0
+torch==2.13.0
+gymnasium==1.0.0
+tensorboard==2.21.0
 # Optional, for --wandb logging:
 # wandb

training/run_training.sh

@@ -48,7 +48,12 @@ fi
 # Export for in-game use (parity-checked); models live in Game/bots/
 .venv/bin/python export_policy.py "checkpoints/$EXP/final.zip" "../Game/bots/$EXP.json"
 
-git add -A checkpoints logs eval_history.json "../Game/bots"
+# Only final.zip, not the intermediate ppo_*_steps.zip checkpoints (.gitignore
+# excludes them) — --resume only ever points at final.zip, so the "training
+# never stranded on one machine" property is fully preserved at ~0.2MB/run
+# instead of ~500MB/run (a single generation-3 experiment dir was 2401 files/
+# 506MB, of which final.zip was 221KB).
+git add "checkpoints/$EXP/final.zip" logs eval_history.json "../Game/bots"
 if git diff --cached --quiet; then
     echo "Nothing new to commit"
 else

training/test_action_space.py

@@ -0,0 +1,105 @@
+"""Rung 0 of TRAINING.md's validation ladder: offline, no Godot, seconds to
+run. Catches the single most likely silent killer in the generation-4
+action-space redesign — a head-order mismatch between the Python trainer and
+Game/scripts/ship_action_codec.gd's HEADS. If they disagree, training still
+runs happily for 24h+ (pitch commands driving strafe thrusters, say) and
+only surfaces as inexplicably-bad behaviour, not an error. This can't
+directly parse the GDScript file, but it locks the two real, checkable
+invariants an order mismatch would actually depend on: that gymnasium's Dict
+key-sorting produces the exact order ship_action_codec.gd's HEADS is written
+in, and that godot_rl's ActionSpaceProcessor converts that into the
+MultiDiscrete nvec the trainer expects. export_policy.py's own index-level
+parity check plus a real in-game round trip (rung 3) are what catch anything
+this can't.
+
+Usage:
+    .venv/bin/python test_action_space.py
+"""
+
+import sys
+
+import gymnasium as gym
+import numpy as np
+from godot_rl.core.utils import ActionSpaceProcessor
+
+from export_policy import ACTION_HEADS
+
+# The order Game/scripts/ship_action_codec.gd's HEADS is written in — kept
+# here as a literal, independent restatement (not derived from ACTION_HEADS)
+# so this test can actually catch export_policy.py's own list being edited
+# out of order too, not just catch nothing because both sides changed
+# together.
+EXPECTED_ORDER = ["rot_x", "rot_y", "rot_z", "thrust_x", "thrust_y", "thrust_z", "turbo"]
+
+
+def check_action_heads_match_expected_order() -> None:
+    names = [head["name"] for head in ACTION_HEADS]
+    assert names == EXPECTED_ORDER, (
+        f"export_policy.ACTION_HEADS order {names} != expected {EXPECTED_ORDER} — "
+        "this must match Game/scripts/ship_action_codec.gd's HEADS exactly"
+    )
+
+
+def check_gymnasium_sorts_to_expected_order() -> None:
+    # Build the Dict deliberately out of order (reversed) to prove it's
+    # gymnasium's sort doing the work here, not insertion order — this is
+    # exactly what godot_env.py does with the dict Godot sends over the wire
+    # (see godot_env.py's from_dict, which builds a spaces.Dict from
+    # ShipAIController.get_action_space()'s Dictionary).
+    sizes = {head["name"]: len(head["bins"]) for head in ACTION_HEADS}
+    reversed_dict = gym.spaces.Dict({name: gym.spaces.Discrete(sizes[name]) for name in reversed(EXPECTED_ORDER)})
+    sorted_names = list(reversed_dict.keys())
+    assert sorted_names == EXPECTED_ORDER, (
+        f"gymnasium.spaces.Dict sorted {sorted_names}, expected {EXPECTED_ORDER} — "
+        "if this changed, every export from this generation onward would be silently "
+        "mis-ordered relative to ship_action_codec.gd"
+    )
+
+
+def check_action_space_processor_produces_expected_multi_discrete() -> None:
+    sizes = [len(head["bins"]) for head in ACTION_HEADS]
+    tuple_space = gym.spaces.Tuple([gym.spaces.Discrete(n) for n in sizes])
+    processor = ActionSpaceProcessor(tuple_space, convert=True)
+    assert isinstance(processor.action_space, gym.spaces.MultiDiscrete), (
+        f"expected MultiDiscrete, got {type(processor.action_space)} — the all-discrete branch "
+        "in godot_rl's ActionSpaceProcessor may have changed (see requirements.txt's pin note)"
+    )
+    assert list(processor.action_space.nvec) == sizes, (
+        f"MultiDiscrete nvec {list(processor.action_space.nvec)} != expected {sizes}"
+    )
+
+
+def check_round_trip_preserves_per_head_values() -> None:
+    # An integer action per env, one column per head in EXPECTED_ORDER —
+    # confirms to_original_dist splits a MultiDiscrete action back into the
+    # same per-head order it was built from (this is what set_action() on
+    # the Godot side receives, keyed by head name).
+    sizes = [len(head["bins"]) for head in ACTION_HEADS]
+    tuple_space = gym.spaces.Tuple([gym.spaces.Discrete(n) for n in sizes])
+    processor = ActionSpaceProcessor(tuple_space, convert=True)
+
+    n_envs = 3
+    rng = np.random.default_rng(0)
+    action = np.stack([rng.integers(0, n, size=n_envs) for n in sizes], axis=1).astype(np.int64)
+    original = processor.to_original_dist(action)
+    assert len(original) == len(sizes)
+    for head_index, expected_column in enumerate(action.T):
+        np.testing.assert_array_equal(np.asarray(original[head_index]), expected_column)
+
+
+def main() -> int:
+    checks = [
+        check_action_heads_match_expected_order,
+        check_gymnasium_sorts_to_expected_order,
+        check_action_space_processor_produces_expected_multi_discrete,
+        check_round_trip_preserves_per_head_values,
+    ]
+    for check in checks:
+        check()
+        print(f"PASS: {check.__name__}")
+    print(f"\nAll {len(checks)} action-space checks passed.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

training/train.py

@@ -15,6 +15,7 @@ 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
@@ -25,6 +26,12 @@ 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
@@ -54,6 +61,154 @@ class GoalRateCallback(BaseCallback):
             self.logger.record("rollout/goal_rate", safe_mean(rates))
 
 
+class FlightTelemetryCallback(BaseCallback):
+    """Logs rollout/{airborne_fraction,mean_altitude,air_touch_fraction,
+    vertical_thrust_mean} — 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."""
+
+    _KEYS = ("airborne_fraction", "mean_altitude", "air_touch_fraction", "vertical_thrust_mean")
+
+    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(
@@ -75,18 +230,57 @@ def parse_args():
     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(
+        "--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-std",
+        "--reset-logits",
         type=float,
         default=None,
-        help="On resume, reset the policy action std to this value (recovers exploration after entropy collapse)",
+        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("--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")
+    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"
@@ -112,12 +306,13 @@ def parse_args():
     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)",
+        "--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(
-        "--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)",
+        "--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,
@@ -158,8 +353,8 @@ def _curriculum_kwargs(args) -> dict:
         "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,
+        "air_drill_chance": args.air_drill_chance,
+        "ai_tilt_penalty": args.tilt_penalty,
         "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,
@@ -192,7 +387,16 @@ def main():
         speedup=args.speedup,
         **_curriculum_kwargs(args),
     )
-    env = VecMonitor(env, info_keywords=("goal_scored",))
+    env = VecMonitor(
+        env,
+        info_keywords=(
+            "goal_scored",
+            "airborne_fraction",
+            "mean_altitude",
+            "air_touch_fraction",
+            "vertical_thrust_mean",
+        ),
+    )
 
     if args.resume:
         model = PPO.load(
@@ -207,14 +411,22 @@ def main():
             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
-
+        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():
-                model.policy.log_std.fill_(math.log(args.reset_std))
-            print(f"Reset policy action std to {args.reset_std}")
+                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",
@@ -232,12 +444,18 @@ def main():
         save_path=str(checkpoint_dir),
         name_prefix="ppo",
     )
-    goal_rate_callback = GoalRateCallback()
+    # 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=[checkpoint_callback, goal_rate_callback],
+            callback=callbacks,
             tb_log_name=args.experiment,
             reset_num_timesteps=not args.resume,
         )