Week 8: Real Orbits, Real Capture

A new environment built around SPICE ephemerides, a 4D action space, and a Mars capture criterion, plus the two debugging weeks it triggered. Visit mazzola.dev for nicer formatting

01 — Where Last Week Ended

Last week’s mars_v11 model arrived inside a 20 Mkm sphere around Mars in every evaluation episode and could not get any closer. The audit traced the ceiling to the facts that the simulator was 2D, the action was a single scalar, Phase 1 corrections were hardcoded to a cross-track direction that cannot change orbital energy, the episode ended at flyby with no Mars Orbit Insertion, and the planets followed circular co-planar orbits that ignored real ephemeris geometry. None of these are things tuning can fix.

This week’s work is the rewrite that addresses all five at once. The new environment lives at envs/mars_nasa_env.py with ENV_NAME = "mars_nasa". The old environment is untouched, and the registry picks up the new one automatically. The new env builds on real SPICE ephemerides, exposes a 4D continuous action, runs a 3-phase finite-state machine that ends in a Mars-centred capture orbit, and is supported by a hand-rolled Lambert solver, a B-plane reference trajectory, an analytic oracle, a behavior-cloning pretrainer, and a 3-stage curriculum.

Two training campaigns came out of the new environment this week. Both failed in interesting ways, and the diagnoses are the most useful thing I have to write down.

02 — A 4D Action and a 36D Observation

The single-scalar action of the previous environment was the deepest architectural limitation of the old setup. The new action is four-dimensional and lives in the J2000 inertial frame:

a = (u_x, u_y, u_z, m),   u ∈ [-1, 1]³,  m ∈ [0, 1]
# Direction is normalised; magnitude scales the burn duration
û = u / max(‖u‖, ε),   t_burn = m × t_max,phase

The agent picks any direction in 3D, and a separate scalar sets how long to fire. A spherical (α, β, m) form was rejected because it has a coordinate singularity at the poles where the gradients of a Gaussian PPO policy explode. The unit-vector form lets the policy learn ‖u‖ → 1 on its own without any constraint penalty. Pointing noise is layered on top by rotating û around a random axis using the existing perturbation infrastructure from Week 6.

The maximum burn duration is phase-dependent. Phase 0 (pre-TMI) keeps the 0.9 fraction from the old env, Phase 1 (cruise) uses 3.5e-5 (about 10 s/step) for trajectory correction maneuvers, and Phase 2 (MOI) uses 0.5 to allow the multi-hundred-m/s capture burn at Mars periapsis. The same action vector applies in all three phases; only the integrator window changes.

Observation

The observation grew from 18 floats to 36. The additions are heliocentric Keplerian elements (sin/cos of inclination, RAAN, argument of perigee), Mars-relative position and velocity (replaces the single distance scalar from before), B-plane targeting error against a Lambert reference (sections 04), reference-trajectory state deviation, a phase code, and the existing TMI alignment scalar gated to Phase 0 only.

Mars Orbit Insertion

The success condition is finally the right one. Once the spacecraft enters Mars’s sphere of influence (about 577 Mkm) and is approaching, the env transitions to Phase 2. From then on, the Mars-relative orbital state is converted to Keplerian elements at every step. Capture is declared when:

a > 0,   0.10 < e < 0.98
r_p = a(1 - e) - R_Mars ∈ [250 km, 10,000 km]
r_a = a(1 + e) < 5 × R_Mars,SOI
# All four conditions held for 3 consecutive steps

The 3-step latch prevents a glancing pass from being mistaken for capture. A graded capture-quality score q_moi ∈ [0, 1] is computed alongside the boolean criterion as a product of Gaussian kernels on (e, r_p) relative to a target of e = 0.85, r_p = 500 km. The terminal reward includes both: +500 for the boolean success plus +200 × q_moi for the graded part, so the agent has a smooth gradient to chase even before it lands inside the strict box.

03 — Real Orbits via SPICE

Hardcoded circular planet orbits go away. Earth and Mars are now driven by JPL’s DE430 ephemeris through SPICE, which Basilisk happens to ship as a support file at ~/.cache/bsk_support_data/. utils/ephemeris.py is a thin spiceypy wrapper that furnishes the kernels once at process start and exposes get_body_state(body, et) and utc_to_et. The new BSK builder, build_nasa_interplanetary_simulation, swaps the hand-written planet-state messages for gravFactory.createSpiceInterface(...) with epochInMsg=True, which lets BSK push live planet states into the simulation as time advances.

The episode now picks a random launch window from a list of four real Mars opportunities:

An earlier candidate of January 2031 was dropped after the oracle could not retarget through it: the Lambert solution diverged from a 374 m/s correction to 18.2 km/s during Phase 1 because the Earth-Mars geometry that month produced a transfer arc too sensitive to the departure timing. Replacing it with May 2033 gave a clean Lambert solve and an oracle that captures successfully.

Because of real ephemerides, none of the simplifying tricks from the old env survive. Mars’s heliocentric distance varies by ~20 Mkm across its year because of its 0.0934 eccentricity, the orbit plane is inclined 1.85° relative to Earth’s, and the launch geometry shifts substantially between 2024 and 2033. The agent has to learn to handle that variability rather than memorize one fixed scenario.

04 — Lambert, B-plane, and the Oracle

To give the agent a stable target, the observation includes the deviation from a Lambert reference trajectory, which is the unique two-body solution that connects the Earth’s position at departure to the Mars’s position at arrival in exactly the chosen flight time.

The Lambert solver lives in envs/lambert_solver.py, a hand-rolled implementation of the Izzo 2015 algorithm. The solver takes the two heliocentric position vectors, the time of flight, and μ_Sun, and returns the velocity at each end of the arc. envs/reference_trajectory.py wraps it in a 200-waypoint sampled trajectory plus a B-plane target, with a disk cache keyed on the rounded epoch and parking altitude.

B-plane Targeting

The B-plane is a 2D plane perpendicular to the asymptotic incoming velocity vector v∞. NASA mission designers use it because hyperbolic flybys map cleanly onto two scalar coordinates (B · T̂ and B · R̂), and the periapsis altitude and inclination of the resulting capture orbit depend almost linearly on those two numbers. envs/bplane.py computes them from a Mars-relative state vector during cruise, and the observation carries the difference between the current and reference B-plane coordinates as two normalised scalars.

Why a Reference at All: Without a reference, the agent has to learn arrival geometry from raw position and velocity, which scrambles every time the launch window changes. With the Lambert reference in the observation, the agent’s job becomes “follow the reference, and do better at the end”, which is well-defined across all four launch windows. The reference weight is annealed to zero by 5M training steps so the policy can outperform Lambert during the final stretch of training.

The Oracle

The oracle (utils/oracles/mars_nasa_oracle.py) is a non-RL expert that captures into Mars orbit using analytic methods only. It exists for two reasons: to validate that the env is solvable in principle, and to provide demonstrations for behavior cloning. Phase 0 fires the Lambert Δv in up to 3 burns at peak TMI alignment. Phase 1 schedules trajectory correction maneuvers at t = 5d, 60d, and 180d using a primer-vector first approximation δv ≈ -r_err / t_remaining. Phase 2 computes the analytic MOI Δv from the hyperbolic state at Mars periapsis and splits it across up to 5 sub-burns.

The oracle gate is the most important test in the project. If the oracle cannot capture reliably, no amount of RL is going to. After two iterations on the Phase 1 cap and the B-plane target offset, the gate test now achieves 78/80 captures across the 4 launch windows × 20 random seeds:

# Oracle gate test — 4 windows × 20 seeds
# 2024-10-15: 20/20
# 2026-11-11: 20/20
# 2028-12-15: 18/20
# 2033-05-05: 20/20
capture rate = 78 / 80 = 97.5%

The Phase 1 dv cap had to grow from 1500 m/s to 5000 m/s before the 2028 window stopped failing, and the B-plane target offset had to drop to zero (the original geometry-dependent offset of 5 × 10⁷ m was actively pushing the spacecraft away from Mars in some configurations).

05 — Bootstrapping with BC + Curriculum

A 36D observation, a 4D action, and four launch windows is enough state space that PPO from random initialisation does not converge in any reasonable time on this problem. The training stack adds two pieces around the standard PPO loop.

Behavior Cloning Warm Start

scripts/generate_bc_dataset.py runs the oracle through 1000 randomised episodes and saves (s_t, a_t) pairs to a single .npz file. scripts/train_bc.py trains a [256, 256] MLP with MSE loss on the actions for 50 epochs at lr = 10⁻⁴. The output checkpoint is loaded into the SB3 PPO actor before the training loop starts:

# In main_train.py, when --bc-init is given:
model = PPO("MlpPolicy", env, policy_kwargs=dict(net_arch=[256,256]), ...)
bc_state = torch.load("bc_checkpoints/bc_policy_v1.pt")
model.policy.load_state_dict(bc_state, strict=False)
# strict=False keeps the value head + log_std random

The strict=False matters because only the actor MLP weights match the BC checkpoint, so the value function and the action distribution scale start from PPO’s defaults. Loading the value head would corrupt the on-policy advantage estimates from the very first rollout.

Three-Stage Curriculum

utils/curriculum.py manages a 3-stage progression keyed on rolling capture rate via a multiprocessing.Value shared across vectorized environments:

# Stage 1 — fixed 2026-11-11 epoch, fixed parking radius, no perturbations
# Stage 2 — random parking radius (6800–8000 km), still fixed epoch
# Stage 3 — random epoch from all 4 windows, perturbations on

However, the mp.Value object wraps an mmap region under the hood, and the default SubprocVecEnv start method (forkserver on Linux) cannot pickle mmap across the process boundary. Switching to start_method="fork" when the curriculum is enabled fixes it because forked children inherit the memory directly without serialising. main_train.py sets the start method automatically when --curriculum is passed.

06 — nasa_v1: A Reward Landscape Disaster

The first training run was nasa_v1: 10M steps, BC initialisation, curriculum off (to isolate baseline behavior). Eval reward stayed between -605 and -80,000 across the entire run.

1. Unbounded Reference Penalty

The Lambert reference deviation was added as a potential-based shaping term:

r_ref = -w_ref × (Δ‖r_sc - r_ref‖) / (10⁸ m)

Combining w_ref = 20 and any divergent trajectory, the cumulative penalty over a 440-step episode reached -50,000 long before the spacecraft did anything else interesting. The “potential-based” framing makes the term telescope to a constant in the limit, but in practice the spacecraft never returns to the reference, so the telescoping never closes. Setting w_ref = 0 for nasa_v2 disabled the term entirely. Lambert deviation is still in the observation; however, it’s not directly penalized by the agent.

2. B-plane PBRS Init Bug

The B-plane error term carried a similar PBRS structure with a previous-step buffer initialised to zero. The first time the spacecraft entered the B-plane evaluation region (Phase 1 onward), the previous error was zero and the current error was several million metres, so the per-step delta was a one-shot reward of around -1000. With w_bplane = 10⁻⁶ in the original config the magnitude was small enough not to matter, but combined with r_ref it pushed several diagnostic episodes into the truly catastrophic range. Setting w_bplane = 0 removed this entirely.

3. Earth Re-impact False Positives

The episode terminated with a -500 penalty if dist_earth < 6528 km after step 40. The intent was to catch a spacecraft that crashed back into Earth after a botched escape. However, random Phase 0 burns from any starting point regularly nudged the LEO perigee below 6528 km altitude (LEO is at 7000 km radius, parking altitude near 600 km, so any retrograde burn fragment crosses the threshold). Even a “do nothing” policy that just orbited Earth got a -500 penalty around step 41 in a meaningful fraction of episodes.

The fix gates the impact check on a latched _has_left_earth flag, set to true the first time the spacecraft passes 50,000 km from Earth. Re-impact is only flagged for spacecraft that actually escaped first.

Reward Distribution Before vs After: Same env, same launch window, same seed, three policies for comparison:

# Before fix
# Oracle:      +146  (should be high, but was being beaten by penalties)
# Coast:       -598  (should be near zero, but penalties dominated)
# 5 random:    ≈ -50,000 each  (catastrophic)

# After fix
# Oracle:      +621
# Coast:       -6
# 5 random:    -22 to -85

The signal-to-noise gap between success and exploration went from less than 1 reward unit to about 700, which is the gradient PPO needs to climb.

07 — nasa_v2: The log_std Runaway

With the reward fixes in place, nasa_v2 ran for 10M steps with BC initialisation and the 3-stage curriculum enabled. The policy improved cleanly during the first few million steps, peaked at an evaluation reward of about 154, then plateaued at around 75 for the remaining run. The agent didn’t achieve any captures and every evaluation episode terminated by step-budget timeout at the 440-step limit.

nasa_v2 evaluation in the heliocentric (J2000) plane. The policy completes a credible transfer arc and the spacecraft passes within 50 Mkm of Mars (well inside the 577 Mkm SOI), but the burn-command panel shows scattered Phase 1 firing and no concentrated MOI burn at periapsis. No capture in any evaluation episode.

The deterministic mean policy (with action noise stripped out) was actually fine and re-running evaluation in deterministic mode gave +154 reward and brought the spacecraft to within 600 Mkm of Mars at closest approach. The policy knew what to do; the stochastic version it actually trained with was burying every decision under noise.

The Diagnosis

Tensorboard’s train/std trace showed the action distribution width grew monotonically from 0.37 to 2.04 over 10M steps. Inspecting the saved model directly, the per-dimension log_std was [1.16, 2.98, 1.99, 1.13]. The action box is [-1, 1] on each dimension, so a Gaussian with σ = 2 effectively samples uniformly from the box and the deterministic mean carries no information.

The cause is the entropy coefficient. The Hohmann env uses ent_coef=0.01, and that was inherited unchanged into the Mars NASA env’s training defaults. Entropy gradient scales with the dimensionality of the action though, and a 4D Gaussian’s entropy gradient is roughly four times larger per step than a 1D Gaussian’s. The same coefficient that kept Hohmann’s policy healthy was actively inflating Mars’s log_std at every update. By the time training was halfway through, KL ≈ 0 and no further policy learning was happening.

Curriculum Never Promoted Either: A second issue compounded the first. The curriculum’s promotion threshold was 60% capture rate to advance from Stage 1, and zero captures means the agent stayed in Stage 1 for all 10M steps. Stages 2 and 3 (parking radius randomisation, all four launch windows, perturbations) never trained. Lowering the thresholds to 5% and 20% along with a hard 3M-step timestep fallback fixes the curriculum, but the log_std runaway has to be fixed first or even Stage 1 will not get anywhere.

The Fix

# 1. Per-env entropy coefficient default (envs/mars_nasa_env.py)
TRAIN_DEFAULTS["ent_coef"] = 0.002

# 2. New CLI flag pinning the initial action distribution width
python main_train.py --env mars_nasa --log-std-init -2.0 ...

# 3. Reset log_std at the start of a resume run, undoing the runaway
python main_train.py --resume runs/nasa_v2/best_model/best_model
    --reset-log-std -2.0 ...

Resuming nasa_v2’s best deterministic policy with log σ reset to -2.0 (giving σ ≈ 0.135, well inside the action box) and entropy coefficient dropped to 0.002 should let the policy continue improving from where the deterministic mean already is, rather than restarting from scratch.

08 — What Comes Next

The plan for next week is the nasa_v3 result and the tuning that follows it, which will be based on these factors in the eval logs.

Curriculum advancement. With promotion thresholds at 5% capture, Stage 1 should hand off within the first few million steps, and the agent should start seeing parking-radius randomisation. If the agent never gets above 5% even on Stage 1, the entropy fix did not address the root cause and the log_std is still leaking through somewhere else.

Capture rate by launch window. The oracle scores 100% on three of the four windows and 90% on 2028. If the trained agent’s per-window split mirrors that, the policy is generalising; if it collapses to one window only, BC is being undone too quickly.

B-plane error distribution. The whole point of the reference trajectory is that the agent learns to track and improve on it. A trained policy with median B-plane error below 5000 km is the release criterion for this env. Anything above 50,000 km means the BC weights have been lost.

Beyond nasa_v3, the next environmental improvement is enabling Mars’s GGM2B spherical-harmonic gravity model in Phase 2 only. Point-mass Mars is fine for cruise and approach, but the capture orbit elements are sensitive to oblateness during the periapsis pass. Enabling it everywhere doubles the integrator cost, so gating it to Phase 2 is a cheaper compromise. The perturbation pipeline from Week 6 also needs to be plumbed through the new env, which currently ignores its own cfg.perturb in Stage 3 because build_nasa_interplanetary_simulation does not accept a perturbation config.

The interesting question for next week is whether 4D action plus capture criterion plus real ephemerides actually delivers what was missing from mars_v11, or whether the next ceiling is lurking one layer deeper.