feat(training): add airborne_penalty and a stage-6 "unmask" curriculum run

Stage 5 (aggression) passed (41-47 vs grounded curric-s2-defend, within
the lenient gate but not yet a clear win). Rather than keep the locomotion
mask on indefinitely, stage 6 reopens full 3D controls on top of the
aggression retune and pairs it with a new dense airborne_penalty (scaled
by height above the floor) so the policy learns to prefer staying grounded
through incentives instead of a hard mask — same regime shift that
regressed stage 3, but this time with a mitigation and ~12x the training
time (~240M timesteps / ~24h vs ~20M / ~2h) to actually re-converge
instead of stalling mid-shift.

airborne_penalty follows the existing SHIP_AI_OVERRIDES pattern: default
0 (off) on ship_ai_controller.gd, exposed via train.py's new
--airborne-penalty flag, added to training_mode.gd's allow-list. Also adds
a per-stage timesteps override in curriculum.py (STAGES[n]["timesteps"])
since this is the first stage to need a different budget than the rest.

Game/scripts/ship_ai_controller.gd

@@ -70,6 +70,14 @@ extends AIController3D
 # makes running the clock out strictly worse than scoring as soon as a
 # chance appears, instead of a free way to keep collecting dense reward.
 @export var time_penalty := 0.001
+# Per-tick penalty scaled by height above the floor (0 on the floor, full
+# value at the arena's ceiling) — distinct from the locomotion mask, which
+# only discards *thrust*-driven vertical/pitch-roll input; a masked ship can
+# still be launched airborne by collisions (ball impacts, ship-vs-ship
+# knockback, the wall/ceiling surface-pull field), and nothing previously
+# penalized time spent up there. Default 0 (off) so ordinary runs are
+# unaffected; the floor-lock curriculum stage turns it on.
+@export var airborne_penalty := 0.0
 
 # Locomotion curriculum: when false, the corresponding action axes are
 # discarded in set_action before reaching the ship, so the ship stays
@@ -196,6 +204,14 @@ func _physics_process(delta):
 		var uprightness: float = ship.global_transform.basis.y.dot(Vector3.UP)
 		reward -= tilt_penalty * (1.0 - uprightness) * 0.5
 
+	# Dense penalty: height above the floor (see airborne_penalty). The
+	# floor sits at world y = 0 (see training_mode.gd's FIELD_MIN_Y/
+	# _escaped bounds); normalized so the worst case is pinned at the
+	# ceiling.
+	if airborne_penalty > 0.0:
+		var height := maxf(ship.global_position.y, 0.0)
+		reward -= airborne_penalty * height / ArenaBoundary.INNER_HEIGHT
+
 
 func _wall_or_ceiling_contact() -> bool:
 	var state := PhysicsServer3D.body_get_direct_state(ship.get_rid())

Game/scripts/training_mode.gd

@@ -188,7 +188,7 @@ const SHIP_AI_OVERRIDES := [
 	"ball_touch_reward", "ball_touch_cooldown_ticks", "ball_touch_direction_floor",
 	"velocity_to_ball_weight", "ball_velocity_to_goal_weight", "ball_distance_penalty",
 	"wall_contact_penalty", "tilt_penalty", "speed_reward_weight", "time_penalty",
-	"allow_vertical", "allow_pitch_roll",
+	"airborne_penalty", "allow_vertical", "allow_pitch_roll",
 ]
 
 
@@ -238,6 +238,7 @@ func _ai_default(name: String) -> Variant:
 		"tilt_penalty": return 0.002
 		"speed_reward_weight": return 0.004
 		"time_penalty": return 0.001
+		"airborne_penalty": return 0.0
 		"allow_vertical", "allow_pitch_roll": return true
 		_: return null
 

TRAINING.md

@@ -159,21 +159,22 @@ just with different curriculum flags.
 | 2 — defend too | `--opponent-mode self_play --no-allow-vertical --no-allow-pitch-roll` | Reintroduces a live opponent (self-play) and the default episode-start mix — the same near-goal state is now simultaneously a finishing chance for one side and a defensive save for the other. Locomotion stays grounded. |
 | 3 — no draws | `--draw-penalty 5 --reset-std 0.3` | Training episodes are golden-goal (end at the *first* goal), so there's no in-episode goal-margin to penalize — `draw_penalty` is the closest available signal: a one-time penalty when an episode times out with no goal at all, on top of the existing per-tick `time_penalty`. Also lifts the locomotion mask (full 3D controls) by omitting `--allow-vertical`/`--allow-pitch-roll`; pair that with `--reset-std` since the policy never got a reward gradient on those axes before now, so expect a brief re-exploration wobble. |
 | 4 — mechanics/refinement | *(no curriculum flags — plain `next_run.sh`)* | Stock self-play, full controls, default reward/start-state mix. This is what all runs before this feature already did. |
-| 5 — aggression | `--opponent-mode self_play --no-allow-vertical --no-allow-pitch-roll --velocity-to-ball-weight 0.05 --ball-distance-penalty 0.006 --ball-touch-reward 0.5` | **Resumes from stage 2 (`curric-s2-defend`), not stage 4** — see the regression note below. Retunes ball-pursuit reward weights (up from 0.02/0.002/0.4) for much more aggressive, constantly-chasing floor play, deliberately keeping the locomotion mask on so it can't reopen the stage-3 regression. |
+| 5 — aggression | `--opponent-mode self_play --no-allow-vertical --no-allow-pitch-roll --velocity-to-ball-weight 0.05 --ball-distance-penalty 0.006 --ball-touch-reward 0.5` | **Resumes from stage 2 (`curric-s2-defend`), not stage 4** — see the regression note below. Retunes ball-pursuit reward weights (up from 0.02/0.002/0.4) for much more aggressive, constantly-chasing floor play, deliberately keeping the locomotion mask on so it can't reopen the stage-3 regression. Passed 2026-07-22 (41-47 vs grounded stage 2 — close, not yet a clear win). |
+| 6 — unmask | `--opponent-mode self_play --velocity-to-ball-weight 0.05 --ball-distance-penalty 0.006 --ball-touch-reward 0.5 --airborne-penalty 0.003` | Re-opens full 3D controls on top of the aggression retune — this is the same grounded-checkpoint-to-full-3D transition that regressed stage 3, but this time paired with `airborne_penalty` (dense, scaled by height above the floor — see `ship_ai_controller.gd`) so the policy learns to *prefer* staying grounded through incentives instead of a hard mask, and can still pick up genuinely useful aerial/wall plays instead of never touching those axes. Runs much longer (~240M timesteps / ~24h vs every prior stage's ~20M/~2h) to actually re-converge through the regime shift instead of stalling mid-way like stage 3 did in a fifth of the time. |
 
-> **Stages 3-4 regressed and are parked.** The locomotion-mask inference bugfix
-> (`8c15c46`) revealed that stage 3's evals up to that point had been running
-> with an unfairly unmasked grounded reference. Re-evaluated fairly,
-> `curric-s2-defend` (grounded) beats both `curric-s3-no_draws` (26-60) and
-> `curric-s4-mechanics` (24-57) — lifting the locomotion mask to full 3D in
-> stage 3 was a clear regression in floor play that self-play never earned
-> back. Stage 5 sidesteps this by resuming and evaluating against stage 2
-> directly (`curriculum.py`'s `resume_from_experiment`/`reference_experiment`
-> stage-dict overrides) instead of chaining through stages 3-4. Full 3D
-> flight is parked as a separate initiative — see TODO.md — that will need a
-> redesigned unmasking approach (more timesteps and/or reward rebalancing) so
-> it doesn't cost floor fundamentals again. See `curriculum_state.json`'s
-> stage-2/stage-3 log entries for the full eval numbers.
+> **Stages 3-4 regressed; stage 6 deliberately reopens the same transition
+> with a mitigation.** The locomotion-mask inference bugfix (`8c15c46`)
+> revealed that stage 3's evals up to that point had been running with an
+> unfairly unmasked grounded reference. Re-evaluated fairly, `curric-s2-defend`
+> (grounded) beats both `curric-s3-no_draws` (26-60) and `curric-s4-mechanics`
+> (24-57) — lifting the locomotion mask to full 3D in stage 3 was a clear
+> regression in floor play that self-play never earned back in 20M steps.
+> Stage 5 sidesteps this by resuming and evaluating against stage 2 directly
+> (`curriculum.py`'s `resume_from_experiment`/`reference_experiment` stage-dict
+> overrides) instead of chaining through stages 3-4. Stage 6 is where full 3D
+> flight comes back — not masked away this time, but discouraged via
+> `airborne_penalty` and given ~12x the training time to settle. See
+> `curriculum_state.json`'s log for the full eval numbers.
 
 All curriculum flags default to leaving Godot's own `@export` defaults
 alone (`train.py` only forwards a flag when you pass it), so ordinary runs
@@ -181,7 +182,8 @@ are unaffected. Full flag list: `--opponent-mode {self_play,inert,frozen}`,
 `--opponent-model <path>` (for `frozen`), `--draw-penalty`,
 `--attack-goal-bias`, `--kickoff-chance`, `--near-goal-chance`,
 `--allow-vertical`/`--no-allow-vertical`, `--allow-pitch-roll`/`--no-allow-pitch-roll`,
-`--velocity-to-ball-weight`, `--ball-distance-penalty`, `--ball-touch-reward`.
+`--velocity-to-ball-weight`, `--ball-distance-penalty`, `--ball-touch-reward`,
+`--airborne-penalty`.
 
 ### Running it automatically
 

training/curriculum.py

@@ -104,6 +104,32 @@ STAGES = [
         "resume_from_experiment": "curric-s2-defend",
         "reference_experiment": "curric-s2-defend",
     },
+    {
+        "name": "unmask",
+        # Re-opens full 3D controls (no more --no-allow-vertical/
+        # --no-allow-pitch-roll) on top of the aggression retune, instead of
+        # keeping locomotion masked indefinitely. The mask blocked *thrust*-
+        # driven flight outright; the new airborne_penalty (dense, scaled by
+        # height above the floor — see ship_ai_controller.gd) is meant to
+        # teach the policy to prefer staying grounded through incentives
+        # rather than a hard constraint, so it can start learning when the
+        # other axes are actually useful (aerial saves, wall recoveries)
+        # instead of never touching them. This resumes the exact regime
+        # shift (grounded checkpoint -> full 3D) that regressed stage 3 —
+        # the mitigation this time is airborne_penalty plus a much longer
+        # run (24h / ~240M steps vs stage 3's 20M) to actually re-converge
+        # instead of stalling mid-shift like stage 3 did in a fifth of the
+        # time.
+        "flags": [
+            "--opponent-mode", "self_play",
+            "--velocity-to-ball-weight", "0.05",
+            "--ball-distance-penalty", "0.006",
+            "--ball-touch-reward", "0.5",
+            "--airborne-penalty", "0.003",
+        ],
+        "grounded": False,
+        "timesteps": 240_000_000,  # ~24h at the standing n-parallel/speedup (20M took ~2h)
+    },
 ]
 
 
@@ -178,9 +204,13 @@ def _grounded_for_experiment(experiment: str) -> bool:
 def run_stage_attempt(stage_index: int, attempt: int, args) -> str:
     exp = 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,
-        "--timesteps", str(args.timesteps),
+        "--timesteps", str(timesteps),
         "--n-parallel", str(args.n_parallel),
         "--speedup", str(args.speedup),
         *STANDING_ARGS,

training/train.py

@@ -96,6 +96,10 @@ def parse_args():
         "--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)",
+    )
 
     return parser.parse_args()
 
@@ -116,6 +120,7 @@ def _curriculum_kwargs(args) -> dict:
         "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,
     }
     return {key: value for key, value in mapping.items() if value is not None}