Differentiable Rocket Science

View on GitHub → Updated 2026-07

A neural network learns spacecraft maneuvers by backpropagating through the physics. 30+ pre-registered experiment rounds, more than half of which refuted their own hypotheses. 92.3% success at 1.03× the analytic-optimal fuel, every headline number re-verified at high fidelity before it was stated.

Abstract

In 2021 I tried to train a reinforcement learning (RL) agent to fly a spacecraft from Earth to the Moon for a graduate astrodynamics course. It trained for hundreds of hours and never learned anything: the distance to the target went up, not down, no matter what I tried. In 2026 I went back, audited the old code, found the real bugs, and rebuilt the project on a different idea — instead of letting the agent learn by trial and error, differentiate through the physics itself and hand it the exact gradient of fuel cost with respect to its own weights. On a full 3-D orbit-circularization task with quaternion attitude dynamics, the result is a policy that succeeds on 92.3% of 4,096 fresh random orbits at a median fuel cost 1.03× the analytic optimum. It matches the textbook maneuver; it does not beat it — on this problem the textbook is provably optimal, so matching is the correct result. In the regimes where the textbook is not optimal, the same method rediscovers a real operational technique (J2 nodal-drift phasing) from the raw objective, discovers combined burn maneuvers on its own, and achieves a verified ballistic lunar capture — arriving in orbit around the Moon without a capture burn. Every claim comes from a fresh evaluation set and survives a float64 re-flight, an integrator energy audit, and a check that no episode gamed the success tolerance. All code, experiment logs, and the original 2021 failure are public.

1. The Course Project That Didn't Work

This project started as my final project for a graduate astrodynamics course in 2021: train an RL agent to find an optimal Earth-to-Mars trajectory under real N-body gravity, and compare it to the Hohmann transfer every textbook derives. I scoped it down to Earth-to-Moon when I realized a single Mars episode would take days to simulate. Then I spent weeks watching it fail.

The final report I turned in was honest about the outcome:

"Due to the complexity of the project, and the time constraint imposed, it appears I have bitten off much more than I could chew."

The symptom was always the same. The agent's score improved as it trained, but the distance to the Moon steadily increased. The spacecraft spent most of its episodes flipping its own orientation back and forth — 6 of its 8 discrete actions were orientation changes, and it used them constantly without going anywhere.

That's a training curve going up while the thing you actually care about goes nowhere. I blamed the reward function, my unfamiliarity with Python, the size of the observation space, and the training time. Some of that was right. But the report also contained one paragraph that turned out to matter more than the rest, about a reward function I ran out of time to try — reward the agent for tracking a precomputed transfer ellipse:

"However, this approach begs the question: if such an orbital path can be defined in the first place, what use is it to utilize a RL agent to follow the path? The entire point is to see if the optimum path can be learned without such direct supervision and guidance."

That question became the design rule for the whole revival.

2. The Problem

Take a spacecraft in a stretched-out elliptical orbit around Earth. You want it in a clean circular orbit, and you want to spend as little fuel as possible. Fuel spent is measured in delta-v (Δv) — the total change in velocity your engine produces.

Every orbital mechanics student learns the answer. Coast to apoapsis — the high point of the ellipse, where you're moving slowest — and burn prograde (in the direction of travel) until your speed reaches circular speed. Burn anywhere else, or point the wrong way, and you pay more.

The picture above is the trained policy flying one episode. Red segments are burns, blue is coasting. The right panel shows the two numbers that define success: semi-major axis ratio and eccentricity, both of which have to land inside a 5% tolerance band and stay there.

The concrete task, in both the 2-D and 3-D versions:

Why circularization and not Earth-to-Moon? Because it's the smallest problem that still contains the hard parts — burn timing, burn direction, attitude control, fuel accounting — and it has a closed-form optimal answer to measure against. The interplanetary goal is Section 11. In 2021 I started with the hardest version of the problem and learned nothing from the failures. This time the plan was to earn each level of difficulty.

The analytic yardsticks are implemented in the repo alongside the simulator: the single-impulse apoapsis circularization, the two-burn Hohmann transfer, and the three-burn bi-elliptic transfer (which beats Hohmann when the radius ratio exceeds 11.94 — the first hint that "the textbook answer" is a set of special cases, not a law of nature).

3. The Autopsy

Before rebuilding anything, I got the 2021 code running again, exactly as it was. It reproduced the original failure on the first try: distance to the Moon increasing monotonically, 379.9 to 393.8 Mm over an episode, episodes never terminating. Reproducing your own failure five years later is a strange kind of satisfying.

Then I audited the code line by line. The 2021 report blamed the reward function and the training time. The audit found the actual bugs, and they were worse:

There were more (reward computed before the action was applied, the reward leaking into the observation, degenerate observation bounds), but those four are the story. The 2021 failure wasn't a deep RL mystery. It was a discrete plug in a continuous socket, a reward stub some idiot (me) forgot to fill in, and a learning rate four orders of magnitude too hot.

I'm keeping the 2021 code, plots, and reports in the repo permanently. They're the "before" picture, and the audit is reproducible against them.

4. Don't Imitate the Textbook

The line from the 2021 report — if you can already define the optimal path, what use is it to learn to follow it? — sets the design rule for the revival:

Analytic solutions are yardsticks, never training targets.

You could train a network to imitate the apoapsis-burn maneuver directly (that's behavior cloning, and I do it below as a baseline). But an imitator can at best match its teacher, and it inherits every assumption the teacher makes. The Hohmann-type analysis is provably optimal only under narrow conditions: two impulses, coplanar orbits, two-body gravity, instantaneous burns, moderate radius ratios. Relax any of those and better solutions are known to exist — bi-elliptic transfers, combined-plane-change burns, low-thrust spirals, and multi-body ballistic capture routes that have actually been flown (the Hiten mission reached lunar orbit that way in 1991).

The interesting question is not whether a network can learn the textbook maneuver. It is whether the textbook maneuver comes out on its own when you optimize actual fuel spent through actual physics — and whether something better comes out in the regimes where the textbook isn't optimal. To keep that question honest, the agent never sees the analytic answer during training. It only gets measured against it afterward.

One more design decision, made in the same spirit of not repeating 2021: the network does not fly the spacecraft. It makes decisions — every 200 seconds it outputs a desired thrust direction and a throttle setting — and a deterministic controller does the flying, slewing the spacecraft to the commanded attitude at a bounded rate. Burn planning is learned; burn execution is classical control. The 2021 agent spent its episodes flipping around in space because it was allowed to command rotations directly. This one can't.

5. Four Ways to Train It

The 2-D version of the task became the bake-off arena. Same environment, same evaluation (success rate over held-out random orbits, and Δv as a multiple of the analytic optimum), four training methods.

First, a sanity check that the task is solvable: a hand-coded scripted controller — coast to apoapsis, latch, burn prograde — solves 100% of seeds at 1.00× the analytic Δv, so the environment is fair and everything below is a learning method trying to reach what the script proves is reachable.

Each method fails for a specific reason, and the reasons are what motivated the method that won.

Model-free PPO fails the way 2021 failed, just with correct code this time. The reward for circularizing is essentially zero until you've done most of the maneuver correctly — burn at the wrong point and you've spent fuel making your orbit worse. Random exploration almost never stumbles into "coast for forty minutes, then burn precisely prograde." Eight runs, and the best one solved 3% of orbits.

Behavior cloning collapses off the expert's path. It imitates well on the states the expert visits, then drifts somewhere the expert never went, and has no idea what to do there.

DAgger fixes the drift and hits a different wall. DAgger (Dataset Aggregation) asks the expert to relabel the states the student actually visits, which fixed the drift — and success stuck at ~45% anyway, while the imitation loss kept improving (0.0115 down to 0.0039). That gap between "imitating better" and "not succeeding more" is structural. The expert is history-dependent — it latches once it reaches apoapsis and stays committed to the burn. The student is a feedforward network with no memory, so it smooths over that sharp decision boundary and no amount of imitation fixes that.

Both failures point the same direction: exploration can't find the maneuver, and imitation can't represent the teacher. Neither method uses the fact that the dynamics are known — gravity here is not a black box, it's a differentiable function we wrote ourselves.

6. Backprop Through the Physics

The idea: instead of estimating a training signal from sampled rewards the way RL does, run the whole episode inside an automatic-differentiation framework, compute the loss (how far from circular you ended up, plus every m/s of Δv you spent), and take its exact gradient with respect to the network weights — backpropagating through every one of the ~1,200 physics steps in the rollout.

The physics engine is the training signal. If burning 0.4 seconds earlier would have saved fuel, the gradient says so, exactly. No exploration noise, no reward shaping tricks, no critic estimating value from samples. The agent gets told, analytically, how each weight in its network moved the fuel bill.

The pieces, concretely:

• The policy is small. A 13 → 128 → 128 → 4 MLP (multilayer perceptron — a plain feedforward network) — three weight matrices you can look at. The 13 inputs are the normalized position and velocity, two orbit-shape errors, where the thrust axis currently points relative to the orbit frame, and remaining fuel. The 4 outputs are a thrust direction in the orbit frame (prograde/normal/radial components) and a throttle.
• The controller is not learned. A proportional pointing law slews the body toward the commanded direction at a bounded rate (0.05 rad/s max). Slewing takes real time, so when the policy commits to a burn still matters — the attitude physics aren't simplified away, they're just not the network's job.
• The physics is the full 13-state rigid body. Position, velocity, attitude quaternion, angular rate, integrated with RK4 at dt = 10 s. The policy decides every 20 substeps (200 s).
• Everything is hand-rolled. The quaternion algebra, the dynamics, the orbital elements, the controller, the environments — all written for this project, typed, and test-covered. No Basilisk, no poliastro. Owning the whole stack from RK4 up is what makes "differentiate through it" possible.

Those three matrices are the whole policy.

In 2-D, this method hit 81% — nearly double DAgger — and its peak checkpoints reached 91% within about a hundred gradient steps. Not a hundred thousand. A hundred. When the gradient is exact, you don't need many of them.

7. Fifty Times Faster

Backprop through 1,200 physics steps is expensive, and the first 3-D implementation (PyTorch) took 13.3 seconds per training iteration — about ninety minutes per debugging cycle.

Porting the hot path to JAX — the rollout compiled through XLA (Accelerated Linear Algebra, JAX's compiler) with lax.scan, gradients via jit(value_and_grad) — brought that to 0.265 seconds per iteration on the same consumer GPU (an RTX 3060). Fifty times faster, numerically exact against the torch reference: a single RK4 step agrees to 6×10⁻⁸, and the derived orbital elements agree bit-for-bit.

The research loop went from ninety minutes to ninety seconds, and the 30-round forensic campaign in the next section happened in days instead of months.

8. The Hard Part: Gradient Monsters

The 3-D campaign started from imitation — DAgger a scripted 3-D expert into the network to get a competent starting policy (79.9% on the fresh evaluation set) — and then refined it with the differentiable-sim gradient on the true objective. The plan was for the exact gradient to polish the imitator past its ceiling.

Instead, training collapsed repeatedly.

The cause took a measurement campaign to pin down: gradients that flow backward through a thousand steps of orbital mechanics are heavy-tailed. The median per-episode gradient norm was a few hundred. The 99th percentile reached 10⁸. And a few episodes per few hundred produced finite-but-monstrous norms of 10¹²–10¹⁹ — one bad geometry (a trajectory grazing the atmosphere-boundary clamp, a near-degenerate orbit frame) and the product of a thousand Jacobians blows up. One Adam step on one monster and the policy forgets how to fly: measured, a single step at learning rate 5×10⁻⁵ typically moves fresh-set success by ±3 points, and in the tail, −12 points.

I ran the whole campaign as pre-registered rounds — write the hypothesis down, define the measurement, run it, log the verdict. Thirty-plus rounds, and more than half refuted their own hypothesis. The refuted rounds are logged in the repo alongside the survivors, and there are more of them.

What survived is a boring, effective playbook: measure the norm distribution, delete the monster tail (trimmed mean over episodes), clip the survivors at the measured 90th percentile, bank progress with an exponential moving average (EMA) of the weights, and polish at low learning rate. No single trick worked alone. Here are four training runs from the same checkpoint and the same seed, differing only in how the per-episode gradients are aggregated:

The red run's aggregation was miscalibrated and it collapsed. The blue run's was measured from the actual norm distribution and it climbed. The only difference between the collapsing run and the climbing run is how the per-episode gradients were averaged.

And here's what that refinement buys on a single episode — the same start orbit flown at four stages of training, the path tightening onto the target ring:

9. Verified Results

On a fresh 4,096-episode evaluation set (never touched during training):

(The oracle's 0.989× is not "beating the optimum" — it succeeds on an easier 80% of orbits and fails the rest, and the tolerance box admits finishes slightly cheaper than exact circularization.)

The refinement climbs past the oracle it started from — 79.9% to 92.3% — which is the point of optimizing the true objective instead of imitating: the student outgrew the teacher. And it isn't one lucky episode; forty random start orbits flown to termination, colored by outcome:

In 2021 I trusted a rising score and got nothing. This time every headline number passes through a verification harness before it gets stated:

• Fresh data only. All numbers come from evaluation sets generated after training. In-run telemetry is banned from claims — one Adam step can move success ±3 points, so quoting your best in-run number is quoting the winner's curse.
• Float64 re-flight. Every episode behind a claim is re-flown at float64 with dt = 1 s. Fuel numbers come from the re-flight, not the training sim.
• The sim isn't being gamed. The 5% tolerance box admits solutions as cheap as 0.849× the exact circularization Δv. Zero of 4,096 episodes finished below that bound — the policy isn't exploiting the tolerance.
• The integrator isn't lying. An RK4 energy audit bounds integration error at 0.005–0.02 m/s of Δv-equivalent per episode — four orders of magnitude below the differences being claimed.

The honest summary: the agent matches the analytic optimum and does not beat it. On this problem, that's the correct outcome — single-body coplanar circularization is exactly the regime where the textbook analysis is provably optimal, so it functions as a correctness check for the whole method. The regimes where beating is possible are in Section 11.

10. One Network, Many Missions

Everything above is one policy trained for one maneuver at its own apoapsis. Several follow-up studies pushed on the "one" part, and two of them taught me the same lesson from different directions.

Commanded target radii. Ask the headline policy to circularize at a radius other than its own apoapsis and it collapses — about 18% success once the target leaves the radius it trained at. Three attempted in-place fixes failed. The fix that worked was embarrassing in hindsight: the specialist had been cloned from an expert that only ever flew fixed targets, so it had never seen tracking behavior. A target-conditioned scripted expert (drive both apses to the commanded radius) DAgger'd into the same 13→128→128→4 network reaches ~99% success across commanded radii spanning ±15% of apoapsis. The blocker was the imitation source, not network capacity.

Thrust levels. Retrain the policy at lower and lower thrust (a 25× drop, from chemical-engine scale toward electric-thruster scale) and success degrades gracefully, 92% to 77%, while the gravity-loss penalty of long finite burns grows just as the physics predicts (1.20× to 1.38× the impulsive bound). Here is what that regime looks like — the lowest-thrust specialist flying one episode:

At chemical thrust this maneuver is one short burn at apoapsis. At 25× less thrust there is no crisp burn to make: the policy spends 55% of the episode burning (red), smeared across two full revolutions, and takes ten hours to spiral out to the target circle. But a policy trained at one thrust level craters at another. A single thrust-conditioned generalist — one 14-input network, thrust added to the observation — recovers ~97% of the five specialists' aggregate success (82.1% vs 84.3%). Getting there required distillation + DAgger; differentiable-sim fine-tuning alone left the new thrust input dead. Same lesson as the target study: the gradient refines behavior that exists, but imitation is how new conditioning enters the network.

Recovering the classical results. In an orbit-averaged low-thrust model, a policy minimizing raw Δv recovers the Edelbaum spiral — the closed-form optimal low-thrust transfer from low Earth orbit to geostationary orbit (LEO to GEO) — at 1.006× the analytic value (6.02 vs 5.95 km/s including the inclination change). And asked for a pure plane change, it rediscovers the raise-turn-return trick on its own: lift the orbit ~10–20%, rotate the plane up high where velocity is small, come back down. That matches the Edelbaum optimum (0.997×) and beats the naive constant-altitude strategy by 4–7%. It found the continuous version of the bi-elliptic inclination trick from nothing but the fuel bill — though only after the yaw parameterization was widened; the first version clamped yaw to ±90°, which silently forbids the return leg. A control parameterization can hide an entire strategy class from the optimizer.

11. Where the Textbook Isn't Optimal

This is the frontier, and it's where the 2021 question — is the textbook maneuver actually optimal, or just optimal under its assumptions? — finally gets an answer. Three results, in increasing order of how much they mattered to me.

The agent discovers combined maneuvers. Give it an inclined ellipse and ask for a circular and equatorial orbit. The naive approach circularizes, then fixes the plane in a separate burn. The known-better approach folds both into a single combined burn at apoapsis. A diff-sim (differentiable-simulation) policy minimizing raw Δv, with no maneuver structure baked in, reaches 84% success at a median 0.75× the naive two-maneuver cost — right at the analytic combined-burn optimum (1.05× that bound). It's the project's first learned result that beats a reasonable strategy a mission designer might write down, and it found the combined burn without being shown one.

A verified ballistic lunar capture. In the circular restricted three-body problem (CR3BP — Earth and Moon both pulling on the spacecraft, the idealized version of real Earth–Moon dynamics), there exist trajectories that arrive at the Moon already captured — negative Moon-relative energy, no capture burn — threading in along the stable manifold of an orbit around the L2 Lagrange point. Finding them is hard: a straightforward search over departure burns found direct arrivals but zero ballistic captures (that null is logged too — the thin capture set eludes grid-and-gradient search). Building the dynamical-systems machinery — Lyapunov orbits via differential correction, their monodromy matrices, the stable manifolds — and seeding from the manifold produced a verified temporary ballistic capture, and a two-impulse transfer patched onto that manifold beats a steel-manned Hohmann-plus-minimal-capture to the same verified captured state by ~5% (0.18 km/s). Modest, temporary (a couple of weeks bound, in a loose high lunar orbit), and real — in the pure CR3BP, no Sun. Then the learned policy caught up: backpropagating a two-burn control through the CR3BP rollout achieves its own dt-robust verified capture, staying bound for ~9–10 lunar revolutions (~35 days) — where the one-burn search had provably stalled (a single burn cannot produce a captured arrival there). En route, two reward hacks were caught and fixed — naive "get close to the Moon" objectives get exploited into fast suicidal plunges that technically minimize distance. It's the same class of reward hacking that sank the 2021 project; this time the verification caught it.

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 fine timestep. This is the regime where the two-body textbook analysis doesn't apply at all.

The J2 story, including the correction. Earth isn't a sphere; its equatorial bulge (the J2 term in the gravity model) slowly rotates every orbit's plane, faster at lower altitude. The textbook impulsive plane-change is blind to this. A diff-sim policy minimizing true Δv over J2-on dynamics, given no maneuver structure at all, rediscovered the operational trick: dive to a lower altitude, let the plane drift faster for free, come back. It reaches the target plane orientation (the "node") for ~0.88× the cost of passively waiting at a 30° node change, improving to a quarter-to-half of passive at 60–90°.

I initially logged this as the project's first genuine beat of a physics-aware strategy. A follow-up round built the analytic version of the same dive-drift maneuver and optimized it directly — and the analytic dive is 10–15% cheaper than the learned one. So the honest classification is rediscovery, not a beat: the agent independently found a real operational technique from the raw objective, and a human with the same insight still executes it better. Both the original claim and the correction are in the repo history.

12. What Doesn't Work Yet

• From-scratch discovery in 3-D is blocked. The exact gradient can refine a competent policy but can't bootstrap one: a policy that never burns gets no signal about when burning would help (the "coast basin"), and there's no exploration mechanism to escape it. Every 3-D result above starts from imitation. Stochastic exploration on top of the analytic gradient is the obvious next lever, and it hasn't been built.
• The headline numbers are single-evaluation. 4,096 fresh episodes is a real sample, but it's one sample. The claims discipline (no in-run telemetry, fresh sets only) exists precisely because the run-to-run spread is ±3 points.
• No ephemeris N-body result yet. The CR3BP capture is the idealized three-body model. A differentiable engine over real JPL ephemerides is built and verified; no maneuver has been trained on it. The big low-energy wins (permanent capture, Sun-assisted perigee lowering) live there.
• Eclipse avoidance was refuted, twice. I expected a low-thrust policy to learn to avoid burning in Earth's shadow. It doesn't, even when incentivized — cancelled-in-shadow thrust costs no fuel in the current model, so there's nothing to avoid, and decision-level control is too coarse for shadow arcs anyway. There's nothing to learn here until a binding cost (a deadline, battery mass) exists.
• Real-mission benchmarks are a plan, not a result. A registry of nine flown missions (MAVEN, Perseverance, Voyager 2, ...) exists for a future comparison via JPL Horizons data. The benchmark doc opens with a warning to my future self: real missions optimize constrained multi-objective problems, and "we beat NASA on Δv" is almost certainly a sign you're measuring the wrong thing.

13. Conclusion

The 2021 report ended with an apology. This one ends with numbers: 92.3% success at 1.03× the analytic-optimal fuel on the full 3-D task, a policy that outgrew the oracle that taught it, one rediscovered operational maneuver, one modest verified beat in the three-body regime, and a verification harness I trust more than my own enthusiasm.

The method's honest scorecard: differentiable simulation, given a competent starting point, refines it past its teacher and matches the provably-optimal solutions wherever they exist — the Edelbaum spiral at 1.006×, the apoapsis burn at 1.03×. Those matches are what make the other results credible: when the same optimizer, pointed at dynamics with no closed-form answer, produces a combined burn, a dive-drift phasing, or a manifold capture, the calibration on the solved cases is the reason to believe it.

The 2021 version of this project wanted to beat Hohmann and couldn't load its own model weights. This version knows exactly what it has: a small network, an exact gradient, a verified simulator, and a to-do list of the regimes where beating the textbook is actually possible.

Technical Details

• Stack: Python ≥ 3.12 (uv-managed), JAX/XLA for the training hot path, PyTorch reference implementation, NumPy/SciPy, Gymnasium environments, hand-rolled quaternion/orbital/dynamics/attitude libraries (pyright strict, test-covered), astropy/astroquery for JPL Horizons ephemerides.
• Hardware: one RTX 3060 (12 GB). Total policy size: 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 evaluation sets, float64 dt=1s re-flight, tolerance-box lower-bound check (0/4096 gamed), RK4 energy audit, pre-registered experiment rounds logged verbatim.
• Everything is public: code, experiment logs including the refuted rounds, the checkpoints behind the figures (GitHub release), and the unmodified 2021 code, plots, and reports.

github.com/robotrocketscience/trajectory-bot →