Differentiable rocket science

View on GitHub Updated 2026-07

Abstract

In 2026 I found the bugs in my 2021 attempt to train a reinforcement learning (RL) agent to fly a spacecraft from Earth to the Moon, and rebuilt the project. The policy, a neural network, now trains on the exact gradient of fuel cost with respect to its weights through a differentiable physics simulation.

On a 3-D orbit-circularization task with quaternion attitude dynamics, one checkpoint succeeds on 92.3% of 4,096 fresh random orbits. A second checkpoint, re-flown in float64, uses a median 1.03x the fuel of the apoapsis burn, a single prograde burn at the orbit's highest point. That burn is the analytic optimum for this task, so matching it without beating it is the correct result.

Where the textbook maneuvers are not optimal, policies trained the same way rediscover J2 dive-drift phasing, an operational technique, from the raw objective and find combined burns on their own. One flies a verified ballistic lunar capture with no capture burn. On a differentiable model of the real solar system, the same method finds how to change orbital plane by using gravity assists to crank inclination up to the kinematic ceiling arcsin(v∞/v_P) and to pump the flyby excess speed, which raises that ceiling. Flown against the real ephemeris, the tour it finds, four pumping flybys and a crank walk, is ballistic at every encounter. Excess speed rises from 5.95 to 16.3 km/s. Inclination reaches 97% of the ceiling (92% on the tightly closed legs alone).

Every circularization result passed these checks, and the multi-body and low-thrust results passed theirs:

• Circularization: a fresh evaluation set, a float64 re-flight, an integrator energy audit, and a check that no episode gamed the 5% success band.
• Multi-body and low-thrust: timestep-sensitivity and conservation checks, and comparison with analytic results.

The code and the round-by-round experiment records in the repo docs are public. The 2021 code is included, so the Section 3 audit can be rerun.

1. The 2021 course project

For a graduate astrodynamics course in 2021 I tried to have an RL agent find an optimal Earth-to-Mars trajectory under real N-body gravity and compare it with the Hohmann transfer. One simulated Mars flight, or episode, would have taken days, so I switched to the Moon. My final report admitted I had taken on more than I could finish.

The agent's training score rose while its distance to the Moon grew. Two of its four discrete actions changed orientation, and it spent most of each episode switching between them.

In the report I blamed the reward function and the training time, as well as my unfamiliarity with Python and the size of the observation space. Some of that was right. The report proposed a reward for tracking a precomputed transfer ellipse, which I ran out of time to try. It also asked what use an RL agent is if the path can be defined in advance, since the point was to learn the optimum path without that supervision. Section 4 makes that a design rule.

2. The circularization task

A spacecraft on an elliptical Earth orbit has to reach a circular orbit on as little fuel as possible. Fuel is measured as delta-v (Δv), the total velocity change the engine produces.

For the apoapsis burn the spacecraft coasts to apoapsis, where it moves slowest, and burns prograde (along the direction of travel) up to circular speed. A burn at any other point or in any other direction costs more.

Red is burning and blue is coasting. The right panel tracks the semi-major-axis ratio and the eccentricity into the 5% band. The semi-major axis is half the ellipse's long diameter.

Across the sampled orbits the apoapsis burn costs 0.44 to 1.18 km/s (median about 0.93). When it makes that burn, the spacecraft is moving 3.6 to 6.3 km/s. The budget is about twice the median. Circularization fuel is quoted as a multiple of the apoapsis burn, so 1.00x is the analytic single-impulse cost and a result under 1.00x would need an explanation. Later results that use a different baseline name it.

I chose circularization over Earth-to-Moon because it is the smallest problem that still needs burn timing, burn direction, attitude control, and fuel accounting, and it has a closed-form optimum. The interplanetary goal returns in Sections 11 and 12. This time I added difficulty in steps instead of starting with the hardest version.

The repo also implements the Hohmann transfer and the three-burn bi-elliptic transfer as analytic baselines. Bi-elliptic can beat Hohmann once the radius ratio exceeds 11.94 and always beats it above 15.58, so the textbook optimum depends on the case.

3. Audit of the 2021 code

Run unchanged, the 2021 code reproduced the failure. Over one episode the distance to the Moon rose monotonically from 379,900 to 393,800 km, and episodes never terminated. A line-by-line audit found these bugs:

The first three bugs in the table sank the project. Smaller bugs included degenerate observation bounds and a reward that was computed before the action was applied and leaked into the observation.

4. Analytic solutions used only for scoring

Because of the question in my 2021 report, the differentiable-sim policy is never given the analytic maneuver as a training target. Behavior cloning of the scripted expert runs as one of the compared methods, but an imitator can at best match the expert it copies and inherits its assumptions. The Hohmann-type analysis gives the most efficient two-impulse transfer between coplanar circular orbits only for two-body gravity with instantaneous burns at moderate radius ratios. Outside those limits better solutions are known, such as bi-elliptic transfers, combined plane-change burns, low-thrust spirals and multi-body ballistic capture routes. Hiten reached the Moon on such a route in 1991, spending no fuel on the ballistic arc.

The project tests whether minimizing fuel in simulation recovers the textbook maneuver, and whether it finds better ones where the textbook is not optimal. The only analytic term in the loss is a smooth Δv-to-go shaping term.

Every 200 seconds the policy outputs a thrust direction and a throttle setting. A fixed pointing controller turns the spacecraft toward it at up to 0.05 rad/s, so burn timing still matters. Unlike the 2021 agent, the policy does not command rotations.

5. Training methods compared

All methods use one 2-D environment and the same held-out random orbits. The scripted expert coasts to apoapsis, latches, and burns prograde. It solves every start orbit at 1.00× the analytic Δv, so the target is reachable.

Model-free PPO with correct code fails like the 2021 agent: the reward stays near zero until most of the maneuver is right, and a mistimed burn makes the orbit worse.

Behavior cloning drifts into states the expert never visited and fails there. DAgger has the expert relabel the states the cloned policy visits, which removed the drift, but success stayed near 45% while the imitation loss fell from 0.0115 to 0.0039. The expert's apoapsis latch depends on history, which a feedforward network with no memory cannot reproduce.

6. Backpropagation through the physics

The simulator's gravity is differentiable, so training can backpropagate the loss through all ~1,200 physics steps of an episode to get the exact gradient for all 18,820 weights and biases in the network. The loss is the final distance from a circular orbit plus the Δv spent. RL estimates this gradient from sampled rewards; computing it exactly removes the exploration noise and the critic. The shaping term stays because the raw objective is too flat to descend.

The policy is a 13 -> 128 -> 128 -> 4 multilayer perceptron, and its thrust direction has prograde, normal and radial components in the orbit frame. Its inputs are normalized position and velocity, errors in semi-major axis, eccentricity and radius, thrust-axis direction in the orbit frame, and remaining fuel. The simulator integrates a 13-state rigid body (position 3, velocity 3, attitude quaternion 4, angular rate 3) with RK4 at dt = 10 s, and the policy acts every 20 substeps. I wrote the quaternion algebra, dynamics, orbital elements, controller and environments myself, without Basilisk or poliastro, and covered them with type checks and tests. Training runs on a separate re-implementation of that code in PyTorch, later JAX.

The trained policy coasts around the orbit and burns prograde at apoapsis (left). Measured as mean throttle sensitivity over flown states, the radius error moves the throttle four times as much as the next input (right). The policy learned this burn and its radius-error trigger from the raw fuel objective alone. The differentiable-sim policy's peak 2-D checkpoints reached 91% within about 100 gradient steps.

7. Porting the training loop to JAX

On an RTX 3060 the PyTorch version took 13.3 seconds per training iteration after kernel fusion. The JAX port compiles the episode through XLA with lax.scan and takes gradients with jit(value_and_grad). On the same RTX 3060 it runs at 0.265 seconds per iteration, fifty times faster, and matches the PyTorch reference: one RK4 step agrees to 6×10⁻⁸ and the orbital elements agree bit for bit. A debugging cycle dropped from about 90 minutes to about 90 seconds, so the investigation in Section 8 took days instead of months.

8. Heavy-tailed gradients

The 3-D training started from an imitation oracle, a policy DAgger-cloned from a scripted 3-D expert, which scored 79.9% on the fresh evaluation set. Training collapsed repeatedly when the differentiable-sim gradient was used to refine it.

Measurement showed that gradients through the ~1,200 physics steps are heavy-tailed. The per-episode gradient norm had a median of a few hundred and a 99th percentile of 10⁸, and a few episodes in every few hundred gave finite norms of 10¹² to 10¹⁹. A single near-zero commanded thrust direction is enough to cause this. The normalization d/‖d‖ then has gradient 1/√ε, about 10⁶ at the default ε. A coasting policy commands such a direction at every step, so the product of the step Jacobians blows up. One Adam step on one of those episodes can ruin the policy. At learning rate 5×10⁻⁵ a typical step moved fresh-set success by ±3 points, and steps in the tail moved it by -12 points.

I ran the investigation as pre-registered rounds. The repo's roadmap writes up the refuted hypotheses alongside the ones that held up. The raw round-by-round logs stayed local and are not published. The recipe that worked measures the gradient-norm distribution, drops the extreme tail with a trimmed mean over episodes, clips the rest at the measured 90th percentile, keeps an exponential moving average of the weights, and polishes at a low learning rate. No part of the recipe worked alone.

All runs in the figure share one lineage and seed. Except for a low-learning-rate polish that continues the blue run, they start from one checkpoint and differ only in how per-episode gradients are combined. Miscalibrated clipping (red) collapsed and the measured trim plus clip (blue) climbed; trimming alone oscillated.

The same start orbit, flown at successive training stages, ends closer to the target circle each time.

9. Verified results

Training on the true objective raised success from the oracle's 79.9% to 92.3%. The oracle's 0.989× does not beat the optimum: it succeeds only on the easier 80% of orbits, and the 5% band accepts finishes slightly cheaper than an exact circle.

Forty random start orbits flown to the end rule out one lucky episode.

Every number here passes a verification harness first. All numbers come from evaluation sets made after training, because one Adam step can move in-run success by ±3 points and the best in-run number carries the winner's curse. Fuel numbers come from a float64 re-flight at dt = 1 s on 1,024 fresh episodes, crashes excluded, so training-sim rounding cannot explain them. The 5% band admits about 0.849× the exact Δv at the median, and no re-flown episode finished below its own bound, so the policy is not exploiting it. An RK4 energy audit puts integration error at 0.005 to 0.02 m/s of Δv-equivalent per episode, which is more than three orders of magnitude below the differences reported here.

10. One network, many missions

Commanded target radii

At a radius other than its own apoapsis the headline policy reaches about 18% success, because its imitation expert only flew fixed targets. Three in-place fixes to that policy failed. DAgger from a target-conditioned scripted expert gave the same 13->128->128->4 network ~99% success for commanded radii within ±15% of apoapsis.

Thrust levels

Retrained at up to 25× lower thrust (chemical-engine toward electric-thruster scale), the policy's success falls from 92% to 77%. The gravity-loss penalty of longer burns rises from 1.20× to 1.38× the apoapsis burn, as the physics predicts. Instead of one short burn at apoapsis, the lowest-thrust policy burns (red) across two revolutions and spirals out.

Each specialist fails badly at other thrust levels. A generalist with thrust as a 14th input reaches 82.1% success, about 97% of the five specialists' aggregate 84.3%. It needed distillation plus DAgger; differentiable-sim fine-tuning alone left the thrust input unused. The commanded-radius study found the same: the gradient refines behavior the policy already has, and a new conditioning input enters the network through imitation.

Classical low-thrust results

In an orbit-averaged model, a policy minimizing raw Δv recovers the Edelbaum spiral, the closed-form optimal low-thrust LEO-to-GEO transfer, at a median 1.006× the analytic value. The 28.5° case gives 6.02 against 5.95 km/s, a ratio of 1.01.

For a pure plane change it found a continuous bi-elliptic plane change on its own from the fuel cost: raise the orbit ~10 to 20%, turn the plane where velocity is lower, and descend. It matches the Edelbaum optimum (0.997×) and costs 4 to 7% less than turning at constant altitude. It did so only after I widened yaw past the ±90° clamp that forbade the return leg.

11. Outside the textbook regime

Combined maneuvers

With no maneuver structure built in, a differentiable-sim policy minimizing raw Δv takes an inclined ellipse to a circular equatorial orbit with one combined burn at apoapsis and 84% success. Its median cost is 0.75× that of circularizing and then changing plane, and 1.05× the analytic combined-burn optimum. This is the project's first learned result to beat a reasonable strategy a mission designer might write down.

A verified ballistic lunar capture

In the circular restricted three-body problem (CR3BP), an idealized model of Earth and Moon both pulling on a spacecraft, some trajectories come in along the stable manifold of an orbit around the L2 Lagrange point and reach the Moon already captured (negative Moon-relative energy) with no capture burn. A grid search over departure burns with gradient refinement found direct arrivals but zero ballistic captures, because the capture set is too thin for that search.

Seeding from the stable manifolds of Lyapunov orbits I computed gave a verified temporary ballistic capture. A two-impulse transfer onto the manifold beats a steel-manned Hohmann transfer plus minimal capture burn to the same verified captured state by 0.18 km/s (~5%). The capture lasts about two weeks in a loose high lunar orbit, in a pure CR3BP model with no Sun.

Backpropagating a two-burn control through the CR3BP rollout later produced its own verified capture, which survives a timestep change and stays bound for ~9 to 10 lunar revolutions (~35 days). A one-burn search had stalled: across the departure grid it searched, no single burn produced a captured arrival. Objectives rewarding closeness to the Moon were met by fast crashes. Verification caught both of these reward hacks, and I fixed them.

The headline arc in the figure arrives captured and stays bound for 4.6 lunar revolutions (about 26 days) with no capture burn, verified at a fine timestep.

J2 phasing, and a correction

Earth's equatorial bulge (the J2 term) rotates every orbit's plane faster at lower altitude, and the textbook impulsive plane change ignores this. The same setup as in the combined maneuvers, run under J2, rediscovered the dive-drift maneuver used to fan out satellite constellations: change the semi-major axis so J2 drives the node faster, then return. I score the policy against the repo's analytic dive-drift maneuver.

It reaches the target node for ~0.88× the cost of passive waiting at a 30° node change, and a quarter to a half of it at 60 to 90°. The passive baseline coasts at the starting altitude for the same time budget, then pays one impulsive burn at the best moment to remove the remaining node error. The ratios above divide by that burn's cost.

Optimizing the analytic dive-drift maneuver directly is 8 to 14% cheaper than the learned policy. I first logged this as the project's first beat of a physics-aware strategy, then reclassified it as a rediscovery after the analytic comparison. Both the original claim and the correction are in the repo history.

12. Gravity assists and the inclination ceiling

The rollout here flies under point-mass bodies at their JPL Horizons positions, and the fuel gradient is taken through every RK4 step. It matches the two-body Kepler limit to one part in 10¹² and holds energy to a few parts in 10¹⁵. From a zero-thrust start the optimizer recovers the analytic Lambert transfer to one part in 10¹¹.

A gravity assist rotates v∞, the spacecraft's excess-speed vector past the planet, without lengthening it, so it turns the orbital plane for almost no fuel. Flybys alone tilt the orbit at most arcsin(v∞ / v_P), where v_P is the planet's orbital speed; the planet's mass does not enter. Cranking climbs to that ceiling over repeated flybys of one planet, as Cassini did over its 54 Titan flybys. Pumping raises v∞, and so the ceiling, with a small burn at the far point of a resonant orbit. The leverage is the v∞ change at the next flyby divided by that burn.

A heavier or closer planet needs fewer flybys: one Jupiter flyby sent Ulysses over the Sun's poles, while Venus takes about five. Solar Orbiter reaches about 33° on its sixth of eight Venus flybys. My Venus tour saturates at 32.8° against a 32.9° ceiling.

I gave the optimizer a target inclination and a fuel objective, starting at zero leverage, and told it nothing about cranking or pumping. It rediscovered the textbook plan: it cranked with near-zero burns to targets under the ceiling, and for targets above it raised v∞ first, spending 1.19 km/s against an analytic 1.20 km/s. Reaching an inclination this way costs 2 to 12 times less than a direct engine plane change, and an order of magnitude less only when the leverage is favorable.

Against the real ephemeris the leverage is 15 to 37. The same leverage that raises v∞ also moves the return point, and unlike in an idealized model the real Earth is at one position on its orbit, so the return can miss it. Keeping the return inside Earth's sphere of influence caps the pump near 0.085 km/s of v∞ per lap, and over five pre-registered rounds I found a cap of 0.1 to 0.19 km/s per lap at every resonance. Earth flybys alone stall after raising v∞ from 8 to about 9.7 km/s.

An early round suggested that repeating one resonance removes the leverage; with correctly sized burns that finding did not hold. Bending the flyby cannot fix the return, because the flyby sets the outgoing orbit before the burn. A correctly timed fix breaks the cap but costs fuel one for one with the v∞ it adds. A second planet breaks the cap without a burn: inner planets pump faster with their shorter years, and the handoff approximately conserves the Tisserand parameter.

My first differentiable version solved each leg as a Lambert boundary-value problem with a flyby-mismatch penalty. Resonant-return legs sit near its singularities, where the gradient exceeds 10¹⁸, and despite gradient clipping and penalty tuning the optimizer moved into them and reported v∞ values that broke the flyby physics. The rewrite propagates each leg forward with a differentiable Kepler propagator and makes each flyby a bounded rotation that conserves v∞ exactly. A Gauss-Newton solve in the loop targets the real moving planet. Leading legs close to within meters, and trailing resonant returns close well inside the sphere of influence.

The model is patched-conic (two-body legs, instantaneous flybys). The result shows the mechanism and claims no fuel record. From a 5.95 km/s launch, four ballistic flybys of Venus and Earth pump v∞ to 16.3 km/s with no maneuvers, and the Earth-to-Venus handoff alone adds 5.0 km/s. Five Venus flybys then crank the inclination from 1° to 27.1°, which is 97% of the ceiling at that v∞. The last of these are loose sphere-of-influence grazes, and the four tightly closed legs alone reach 92%.

Re-hitting the real moving planet uses about half of every turn, so the crank takes five flybys instead of the ideal one or two. Of eight launch epochs over two synodic cycles, five reach 15 to 18 km/s and three fail: two have no viable cheap launch, and one launches but never climbs. The median final v∞ is 15.3 km/s.

13. Open problems

From-scratch discovery in 3-D is blocked, so every 3-D result here starts from imitation. A policy that never burns gets no gradient signal about when a burn would help, and nothing in the method makes it explore. The next step is stochastic exploration on top of the exact gradient. The version I tried, reparameterized noise with an entropy bonus, never learned to burn, and the version that did work needed a rescaled final layer instead.

Each headline number is one sample: a single evaluation on 4,096 fresh episodes.

The Section 12 tour is not yet a trained policy. Tree search over ballistically closed continuations picks each flyby, and gradients tune the continuous parameters. In this formulation, gradient restarts asked to choose flybys ran into boundary-value singularities, and none of five found a valid tour. No transfer policy has been trained end to end on the real ephemeris, and the large low-energy gains (permanent capture, Sun-assisted perigee lowering) remain open.

My expectation that a low-thrust policy would learn to avoid thrusting in Earth's shadow was refuted: even with an incentive it did not. Thrust cancelled in shadow costs no fuel in the current model, and decision-level control is too coarse for shadow arcs, so there is nothing to learn until a binding cost such as a deadline or battery mass exists.

A registry of flown missions (MAVEN, Perseverance, Voyager 2 and others) exists for a future comparison against JPL Horizons data; there are no benchmark results yet. My notes warn that real missions optimize several constrained objectives, so "we beat NASA on Δv" almost always means I am measuring the wrong thing.

14. Conclusion

Given a competent starting point, differentiable simulation refines a policy past its teacher and matches the provably optimal solutions where they exist: the Edelbaum spiral at 1.006x and the apoapsis burn at 1.03x. Those matches are why I trust what the same optimizer finds on dynamics with no closed-form answer: the combined burn, the dive-drift phasing, and the manifold capture. Of those, the three-body transfer is a modest verified improvement on its baseline, and the J2 phasing rediscovers an operational maneuver.

Technical details

• Stack: Python >= 3.12 (uv), JAX/XLA, PyTorch reference, NumPy/SciPy, Gymnasium, astropy/astroquery for Horizons ephemerides, and hand-written quaternion, orbital, dynamics and attitude libraries (pyright strict, tested).
• Hardware: one RTX 3060 (12 GB). Policy: 13->128->128->4, three weight matrices.
• Method: differentiable-simulation policy gradient, jit(value_and_grad) through full rollouts (~1,200 RK4 substeps, 60 decisions); trimmed-mean + measured-p90-clip gradient aggregation; EMA banking; imitation bootstrap via scripted experts + DAgger.
• Verification: fresh 4,096-episode sets, float64 dt=1s re-flight of 1,024, 5%-band lower-bound check (0/1024 gamed), RK4 energy audit, pre-registered rounds logged verbatim.
• Public: code, experiment records including refuted rounds, figure checkpoints (GitHub release), and the unmodified 2021 code, plots, and reports.

github.com/robotrocketscience/trajectory-bot ->