Anisotropic Representations Improve Planning in JEPA World Models
Mingu Kang, Yoori Oh, Sookyung Kim, Joonseok Lee
cs.RO
2026-09-29
AnisoWM's ΛReg learns an anisotropic Gaussian target and beats LeWM in all four envs (TwoRoom 93% vs 87%, Cube 79% vs 74%).
JEPA world models predict action-conditioned dynamics in representation space. For visual goal planning, a planner rolls candidate action sequences in that space and scores them by Euclidean distance from the predicted terminal representation to the goal. The encoder supplies features for prediction and also defines how terminal errors are weighted.
LeWorldModel (LeWM) trains the encoder and an action-conditioned predictor jointly, and uses SIGReg from LeJEPA (sketched isotropic Gaussian regularization) to stop collapse. SIGReg pushes features toward an isotropic Gaussian. That target comes from probing-risk arguments, not from whether Euclidean latent distance ranks actions the way the task cost does. Accurate prediction and a noncollapsed representation can still leave the planner ranking feasible outcomes in the wrong order.
A linear-Gaussian analysis makes the mechanism concrete. Isotropic feature covariance induces an inverse-state-covariance metric in physical coordinates: low-variance training directions become expensive errors at planning time. When a goal sits outside the reachable set and the planner must trade residual error across directions, the latent ranking and the task ranking can diverge, leaving strictly positive planning regret. A 2D MLP toy repeats the split: shrinking process noise by 30× cuts held-out conditional-mean prediction error, while normalized physical planning regret stays near 0.16 over ten seeds, matching the 0.1607 inverse-covariance reference. Better prediction does not remove the ranking mismatch.
AnisoWM changes only the regularization target. ΛReg replaces the fixed isotropic Gaussian with a learnable diagonal covariance under two constraints:
κ=1 recovers the isotropic target, which is LeWM. The target is diagonal in representation coordinates; the encoder can still rotate those axes relative to the underlying state geometry. Features are standardized by Λ^{-1/2} before the original SIGReg statistic, because zN(0,Λ) implies Λ^{-1/2}z N(0,I). Standardization sits only inside the regularizer. Prediction and planning keep the raw z.
Λ is parameterized with clipped logits so the trace and condition-number constraints stay feasible. The prediction loss, the predictor, and the Cross-Entropy Method (CEM) Euclidean planner are unchanged. Λ can be discarded after training; planning introduces no extra metric.
The learnable target expands the limiting state-space metric from qΣ⁻¹ to qΣ^{-1/2} L Σ^{-1/2}, with L inside the constrained set. The prediction term assigns larger target variance to lower-noise directions, which can cancel the inverse-covariance weighting of isotropic SIGReg. In a 2D closed-form case, setting κ to the condition number of the state covariance recovers an isotropic state metric and drives planning regret to zero; larger κ raises regret again.
Evaluation follows the datasets, architecture, and visual goal-planning protocol of LeWM on TwoRoom, Reacher, PushT, and Cube. Representations are 192-dimensional, with λ=0.09. The main comparison uses a shared κ=2 and reports the mean of three training seeds. The LeWM baseline is trained in the same pipeline, which is κ=1. Success rates for PLDM, DINO-WM, GCBC, GCIQL, GCIVL, and Random are taken from the original LeWM paper.
| Env | AnisoWM | LeWM |
| TwoRoom | 93% | 87% |
| Reacher | 89% | 86% |
| PushT | 97% | 96% |
| Cube | 79% | 74% |
AnisoWM wins all four. Some reference baselines score higher on TwoRoom and Cube; AnisoWM keeps LeWM's latent-space planner.
On the same recorded action sequences, both models are scored on how well latent costs order task outcomes. Jpred is the CEM planning cost; Jenc uses the encoded reached observation and skips the predictor rollout:
| Env | Jpred LeWM | Jpred AnisoWM | Jenc LeWM | Jenc AnisoWM |
| TwoRoom | 0.575 | 0.646 | 0.564 | 0.661 |
| Reacher | 0.676 | 0.723 | 0.790 | 0.893 |
| PushT | 0.594 | 0.621 | 0.647 | 0.641 |
| Cube | 0.537 | 0.553 | 0.555 | 0.576 |
Jpred improves in every environment. Jenc on PushT is down 0.005. Spearman ρ between latent cost and task cost rises from 0.15 to 0.91 on Cube and from 0.65 to 0.92 on PushT.
The shared κ=2 still yields different spectra: coordinates above mean target variance range from 95 on PushT to 124 on Reacher. Success versus κ is non-monotone, and the best observed anisotropy differs by environment. Each κ>1 point is a single run; the paper treats the sweep as a sensitivity check, not a tuned comparison.
The change lives in the training regularizer. The planner and predictor stay as they are. For pipelines already running LeWM-style JEPA latent planning, this is a zero-cost inference knob: extra diagonal variances that get dropped at planning time.
Collapse-prevention geometry is planning geometry. An isotropic Gaussian can be a reasonable probing target and still be the wrong Euclidean planning metric. SCALE, TRM, and DecisionMetric Alignment attack the same mismatch by changing distances or adding reachability metrics. AnisoWM puts it inside the Gaussian regularizer and leaves the Euclidean planner alone.
The gains are incremental. PushT moves 1 point, Reacher 3, and the main table does not tune κ per environment. The mechanism, that the regularizer target is the planning metric, is the part worth carrying forward.
There is no standalone Limitations section. Points the authors put in the results and analysis: excess anisotropy can raise planning regret; raising κ does not monotonically raise success; each κ>1 setting is a single training run; some reference baselines beat AnisoWM on TwoRoom and Cube.
The theory is linear-Gaussian with affine encoders. Lemma 1 needs the affine assumption; a nonlinear encoder puts weaker constraints on the local metric. The visual-control experiments are nonlinear, and the success and ranking numbers line up with the story, but that does not mean the theorem carries over. The main comparison uses three seeds, and PushT Jenc is slightly worse. There is no matched comparison against SCALE or TRM, and no real robot. Λ is forced diagonal, so anisotropy axes are left to the encoder; how stable that bias is in a 192-d visual representation is not isolated in the main experiments.