Planning through a learned world model

A closed-loop episode. The planner only ever sees 64×64 frames through the learned encoder. It has to pump the swing to build energy, then damp it to land the ball in the cup.

Why this project

I wanted to understand JEPA (Joint-Embedding Predictive Architecture) by building one end to end, not by reading about it. Ball-in-cup was a deliberately unimportant proxy task. The point was the whole loop: self-supervised pretraining on pixels, checking what the latent space actually knows, and planning through the learned model. Solving the task at 97% was a bonus.

32latent dims per frame

R² 0.97position recovered from z (MLP probe)

~150real steps of stable rollout

29/30solved from random starts

Architecture

E · encoderOne 64×64 RGB frame → 32-d latent z. Purely spatial.

P · predictorWindow of K=3 latents + action → next latent. All velocity reasoning lives here.

ProbeFit after E and P are frozen. Reads position back out of z for the cost.

CEM / MPCSample action sequences, roll them through P, refit to the elites, execute one step, replan.

The loss is VICReg-style: prediction error plus variance and covariance terms against collapse, with no EMA target and no stop-gradient. Numerical gradient checks found a vanishing-gradient-at-collapse bug in the covariance term, fixed by square-rooting it before the penalty.

Design decision An early version stacked frames into the encoder. A linear probe showed it couldn't recover velocity at all (R² ≈ 0). Ball-in-cup is a second-order ODE, so position alone isn't Markovian. The fix: a single-frame encoder, with a K=3 window in the predictor. K=2 was rejected because a two-point finite difference amplifies position noise about 2500× at dt = 0.02 s.

The drift problem

Planning needs long horizons, since the swing has to be pumped over 100+ steps. But feeding the predictor its own outputs made rollout error explode past 30–50 steps. The model had only ever been trained on real context windows (exposure bias).

Two changes fixed it together:

Unrolled training loss. P is supervised on 5 successive self-fed predictions, not just one.
Predictor stride. One P call advances 5 real control steps under a held action. This needed new action-repeat data, because the original data changed action every step.

After that, ball-position error stayed flat at about 0.08–0.11 m across the full 150-step horizon, with no compounding.

Results: 60% → 97%

Real closed-loop episodes against the simulator, random initial conditions, n=30 unless noted.

Baselinelinear probe, no warm-start, n=10

60%

MLP probe12× more accurate offline

67%

Linear probe, re-runconfirms the null result

63%

+ warm-startstratified elites carried forward

73%

+ 24 s time budgetwas 12 s; search unchanged

97%

Three findings that mattered

A better probe didn't help. The MLP probe cut per-frame position error from 0.083 m to 0.007 m, and closed-loop success didn't move (63% vs 67%, same at n=30). The cost only uses the difference between ball and cup positions, so shared decoding bias largely cancels.
Failures weren't near-misses. Filmstrips of every failure showed the same thing: the ball still mid-swing at the cutoff. The planner re-randomised each step and kept abandoning damping trajectories. Carrying forward a slice of elites from every CEM iteration kept diversity and fixed most of it.
The honest kind of retry. A real robot can't reroll an easier start. It can only keep acting from where it is. Doubling the time budget showed most "failures" were slow, not stuck. Several solved only after 12 s, some as late as 21.7 s.

Principle kept throughout The probe is a readout fitted after E and P are frozen, never a training signal. Training on position would probably have closed the precision gap faster, but then the project would stop testing what self-supervision recovers on its own.

What came next

A Gromov δ-hyperbolicity test on the trained latent space came out low (relative δ ≈ 0.15–0.27), even though it was trained with a plain Euclidean loss. That suggested the latent space already has tree-like structure, and led to the hyperbolic latent-space work.