Teaching JEPA world models to fly a drone

The project started as a port. LeWM is a JEPA world model from the open stable-worldmodel project: it encodes an observation into a latent state, predicts how that state changes under candidate actions, and scores the predictions against a goal. We rebuilt its runtime in Rust on Candle CUDA tensors so that decoding, encoding, rollout, scoring and action selection all stay on the GPU. CEM, MPPI and iCEM planners run there too.

The port was checked against the official Python implementation on the same RTX 4090 before any speed claim. Over a fixed PushT batch, the largest difference in predicted latents was 7.3 × 10⁻⁴, and the cost ranking picked the same best action.

From June 12 the same runtime learned a drone. A vector-observation LeWM trained on recorded racing-drone flights, a Bevy simulator replayed them, and a gate-loop test asked the planner to fly the course. Fitting the simulated plant to the recordings let the recorded expert actions complete all four gates, up from two. Plan traces then showed where the remaining work was: the model's objective often preferred its own action sequence to the expert's while the drone drifted off the recorded path. The next model would need to be built for control from the start.

The SkyJEPA paper describes a JEPA world model for quadrotors. Its authors' model, training code and data were not available, so the implementation was reconstructed from the paper, keeping the paper's interfaces and sizes separate from the choices made here. It runs natively in the Rust/Candle runtime with a dedicated 200 Hz rotor-physics simulator and a 20 Hz controller.

Two problems were solved the same day. The first dataset was rejected because it did not excite the rotors differently enough to learn from, and was replaced by 2,000 ten-second flights across 100 randomized vehicles. Sampling raw rotor forces also proved a poor way to start each plan, so a trim-aware geometric controller now proposes the first sequence and MPPI improves it. That hybrid design stayed for the rest of the project.

The first trained controller passed all 63 test flights, against 62 for the geometric controller alone. Its worst trajectory error was 0.41 m, against 0.89 m, at a planning p95 of 8.69 ms.

Before claiming more, the pipeline was audited and hardened. Test vehicles now come from physical configurations never seen in training. Every configuration contains both hover and moving flights. Checkpoints are bound to their normalization, timestep and training history, and interrupted training resumes exactly. The evaluation protocol, with three training seeds, the test populations and the controller settings, was fixed before any final test ran.

Each model trained in about 34 minutes on a shared RTX 4090, from 1,600 ten-second flights. Every controller then flew the same 63 hover, circle and figure-eight cases on unseen vehicles.

This is the work behind the short film Self-Improving Agent Trains a UAV with JEPA: an agent read the paper, wrote the code and trained a drone controller that flies more accurately than the hand-written one.

On September 17 the JEPA-Anything paper introduced Orthogonal Predictive Factorization (OPF). OPF splits the latent state into complementary groups of coordinates, predicts each group, and recombines them into the next state, with the aim of longer, more stable predictions. The release turned out to be larger than its GitHub page suggested: a Hugging Face repository holds 36 checkpoint files. Its core library and three released locomotion models ran on our RTX 4090.

The goal was to compare it with SkyJEPA, not replace it, and to film the comparison. OPF went into SkyJEPA's own structure: the same encoders, the same 24-value recurrent state and four factor groups of six coordinates recombined through a learned orthonormal basis. Its latent model trained in PyTorch in 24 seconds, then a new importer brought it into the Rust runtime, matching the Python rollout to within 9 × 10⁻⁸ over 20 steps. The prober and all flying ran natively.

All three passed 126 of 126 tracking cases with no ground contact and no missed deadlines. SkyJEPA held its lead over this first OPF adaptation in both conditions. Open-loop prediction over 26,400 held-out windows tells the same story, and shows what a better architecture would need to fix:

To make comparisons like this routine, the work moved to MuJoCo 3.13. The brief was a gym for robots: abstract, with training aids and overlays, but good to look at. It became the JEPA Gym: a charcoal floor grid, landing pads, reference hoops, lane markings and live telemetry. MuJoCo simulates rigid-body motion and ground contact at 400 Hz. Each motor has first-order lag, thrust gain and reaction torque, the airframe has linear drag, and control runs at 20 Hz. The training aids have no collision response.

Looking at the drone shaped the gym as much as the tests did. Watching it fly showed that it never turned to face its path, so heading control was added: face the direction of travel, or turn in place independently. A closer look showed landing skids that did not follow the nose and rotors in a plus layout. The skids were realigned, and the drone became an X frame: diagonal arms, motors 0.24 m from the centre at 45° to the body axes, the nose between the front motors, and a motor mixer that matches those positions.

The gym briefly supported several airframes, which made it more complex than intended. It was cut back to one X-frame drone and one motor mixer. The simplified code reproduced the recorded flights exactly, with zero state difference. On a Mac, an interactive viewer opens the same gym, with live telemetry and replay of recorded learned-model flights.

The gym compares two matched models: a plain LE-WM JEPA baseline and the same architecture with OPF. Both are trained on the same data, predict 0.75 s ahead, and fly through the same MPPI controller and geometric action prior. Each representation is then frozen and given two readouts. The physics readout learns corrections inside supplied equations of motion. The learned readout predicts future states directly, without the physics integrator.

The first study, with three training seeds on the original plus-frame drone, gave OPF a promising signal:

The protocol was fixed before training: seeds 7, 17, 29, 43 and 59; 864 ten-second episodes across 72 physical configurations, balanced across fixed, path-facing and turning-in-place heading; 48 configurations for training and eight each for validation, ordinary testing and a held-out region of heavier drones with slower motors. Both models have 40,928 trainable parameters and read ten past states and nine motor commands. Checkpoints were chosen on validation data only, and no test result tuned anything. The full run of data, training and evaluation took 50 minutes 32 seconds on the RTX 4090.

Action regret measures how close a model's top-ranked action comes to the best one actually available, scored by MuJoCo itself; lower is better. OPF picked slightly better actions with both readouts. The learned readout flew every flight to completion without the supplied physics equations, more than four times faster to score than the physics readout.

Everything here is simulated. The result is a working path from a new world-model idea to measured flights, and a set of baselines that any future JEPA has to beat.