Does LeWorldModel Preserve Local Action Geometry?

AI summary

Overview

Research area: Representation geometry and evaluation of action-conditioned world models (Joint-Embedding Predictive Architectures), with a focus on robot manipulation benchmarks.

Technical level: Advanced. The paper is short and clearly written, but it relies on Jacobians, Gram matrices, trace normalization, condition numbers and principal angles.

Scope: A diagnostic study of whether a pretrained LeWorldModel (LeWM) reproduces the local sensitivity of two simulated environments (Reacher and OGBench Cube) to small action perturbations, measured with matched simulator counterfactuals across three perturbation scales and several output spaces.

What This Paper Is About

World models are usually judged by how accurately they predict the next state or how well a planner performs using them. Neither test checks whether the model gets the relative effects of nearby actions right — that is, whether nudging an action in one direction changes the predicted outcome by the correct amount and in the correct direction compared with a nudge in another direction. The authors build a matched comparison between LeWM's predicted embeddings and the real simulator's encoded responses, then ask whether the model's local action geometry agrees with the environment's, how that agreement depends on perturbation size, and whether the resulting metrics actually predict responses along unseen action directions.

Key Contributions

1. A matched finite-scale comparison of action geometry. The same five-action block is perturbed by a shared normalized offset in both the model and the simulator, and both branches use matched finite differences, yielding coordinate-wise Gram matrices for Reacher ($d=2$) and OGBench Cube ($d=5$) at three perturbation scales ($\epsilon = 10^{-3}$, $3\times10^{-3}$, $10^{-2}$).
2. Calibration against deliberately weakened references. The model's trace-normalized discrepancy is compared with isotropy ($I_d/d$), random orthogonal conjugation of the environment metric, and shuffled-state matching, so that each reference removes a specific source of information.
3. Validation on held-out action directions in three output metrics. Eight held-out sphere directions per state, reused across scales, are used to test whether the Gram predicts response magnitudes, reported as weighted absolute percentage error (WAPE) in physical pose, RGB pixels and frozen embeddings.
4. A documented protocol revision and reproducibility audit. An audit of action alignment, output coordinates and simulator-state restoration led the authors to regenerate trajectories in the current simulator from the original dataset sampling anchors; the revised estimates supersede earlier results.

Main Findings

• Sensitive directions are misaligned at the smallest scale. At the primary scale ($\epsilon = 10^{-3}$), the most sensitive action directions differ by a median of $36^\circ$ on Reacher and $42^\circ$ on Cube (precisely $36.03^\circ$ and $41.52^\circ$). Cube also shows a median maximum principal angle of $40.27^\circ$ between its leading two-dimensional subspaces.

• Relative geometric discrepancy is substantial. Median trace-normalized Gram discrepancy $D_G$ is 0.465 on Reacher (95% episode-bootstrap interval [0.423, 0.500]) and 0.508 on Cube ([0.487, 0.529]).

• Similar eccentricity can coexist with different directions. Reacher's environment and model condition numbers are both near 4.5 (4.49 and 4.50) despite the directional mismatch. Cube's median model condition number ($1.07\times10^{4}$) is roughly four times the environment value ($2.84\times10^{3}$). A numerical floor affects 32.4% of Cube environment matrices, so the upper tail requires care.

• The model beats some references but not all. Cube's $D_G$ (0.508) is lower than isotropy (0.609), rotation (0.885) and shuffle (0.560). Reacher (0.465) improves on rotation (0.632) and shuffle (0.530), while isotropy (0.449) is closer to its environment metric. Paired model-minus-isotropic differences are $+0.083$ ([0.064, 0.108]) for Reacher and $-0.108$ ([−0.128, −0.094]) for Cube.

• Matching the state helps. Paired model-minus-shuffle median differences are $-0.059$ for Cube ([−0.074, −0.029]) and $-0.047$ for Reacher ([−0.071, −0.014]).

• Absolute sensitivity is lower in the model. The median ratio $\sqrt{\operatorname{tr}(G^{\mathrm{model}})/\operatorname{tr}(G^{\mathrm{env}})}$ indicates about half the environment's sensitivity on Cube and four-fifths on Reacher.

• Geometry is strongly scale-dependent. A tenfold increase in perturbation size reduces median $D_G$ by roughly 70% on Reacher and 47% on Cube. On Cube, median $D_G$ falls 0.508 → 0.408 → 0.271 and the top-1 angle falls $41.52^\circ$ → $34.82^\circ$ → $20.96^\circ$ across the three scales; Reacher falls 0.465 → 0.308 → 0.141 with angles $36.03^\circ$ → $20.70^\circ$ → $8.89^\circ$.

• The environment's dominant direction rotates while the model's does not. Between $\epsilon = 10^{-3}$ and $10^{-2}$ on Cube, the environment's dominant direction changes by a median $38.99^\circ$, versus $0.01^\circ$ for the model. Cross-scale ICC for $D_G$ is 0.220 (Cube) and 0.200 (Reacher), with pairwise Spearman correlations of 0.220–0.432 and 0.205–0.375, even though repeated runs at a fixed scale are identical (ICC 1).

• The choice of output metric dominates the measurement. Reacher's median physical-pose Gram change across scales is $9.32\times10^{-7}$, while its pixel change is 0.1125 and its embedding change is 0.4331. Cube shows the same ordering: 0.0270 (physical pose), 0.0808 (pixels), 0.4853 (embeddings).

• Environment Grams predict held-out response magnitudes; model Grams do not, at this scale. Median embedding WAPE on eight held-out directions is 17.0% (Reacher) and 28.0% (Cube) for environment Grams, versus 26.2% and 38.2% for an isotropic reference with the same trace. Model Gram errors are higher: 33.2% and 56.4%. The 95% intervals for the environment Gram's advantage span roughly six to eight percentage points on Reacher and eight to twelve on Cube.

• Gain error explains much of Cube's failure. Giving the model the environment's trace ("oracle gain") reduces Cube's median WAPE from 56.4% to 32.3% — about 24 percentage points — bringing it below isotropy (38.2%) and within five points of the environment Gram (28.0%). Reacher improves from 33.2% to 26.6%.

• Prediction is easy in physical coordinates and hard in pixels and latents. Physical-pose prediction has median WAPE below 0.04% in both environments at the primary scale; pixel errors are about 5% on Reacher and 15% on Cube. One-sided embedding displacements are harder than centered responses: median errors rise from 17% to 32% on Reacher and from 28% to 43% on Cube.

• The finding survives checks on eigenvalue separation and trace regularization. Requiring both Cube top-1 spectral gaps to be at least 0.1 retains 449 states with median angle $41.09^\circ$; the top-2 gap check retains 302 states with median maximum angle $34.66^\circ$. Reacher's top-1 gap check retains 454 states with median angle $35.80^\circ$. Adding a trace constant up to $10^{-8}$ changes Cube $D_G$ by less than $9\times10^{-10}$ and Reacher's by less than $3\times10^{-11}$.

• Retained outliers qualify the medians. One Reacher state (sample 389) reaches about 117% physical-prediction WAPE at the primary scale while the other 457 states are below 1%. Cube's physical-pose cross-scale Gram change has upper quartile 0.111 and maximum 1.172 despite a median of 0.027, and its median physical-pose WAPE rises to about 1.47% at $\epsilon = 0.01$.

Methodology in Plain English

The authors take a pretrained LeWM checkpoint and, for each of 458 Reacher anchors and 500 Cube anchors, fix a context of three observations and a block of five recorded actions. They then push all five actions in the same direction by a small amount, scaled by each action dimension's standard deviation.

Two things are measured for that push. The model branch feeds the perturbed action block to LeWM and records the predicted next embedding. The environment branch restores the simulator to the exact same state, executes the perturbed actions, renders the endpoint, and passes that image through the same frozen encoder, so both responses live in the same output space.

By applying the push along each action coordinate separately (using a symmetric finite difference), they build a matrix describing how strongly the response changes in every direction — the Gram matrix $J^\top J$. Two Gram matrices can then be compared: one from the model, one from the environment. After dividing each by its own trace (which discards overall gain), the Frobenius distance between them, $D_G$, summarizes how differently the two systems shape local action sensitivity.

Three calibration baselines make the resulting number interpretable: an isotropic matrix (no preferred direction), a randomly rotated version of the environment metric (correct spectrum, scrambled orientation), and a shuffled pairing of model states with environment states from a different anchor (tests whether matching the current state matters at all).

Two additional readouts guard against over-interpreting $D_G$. Eigen-decomposition gives the angle between the most sensitive directions and the eccentricity (condition number) of each metric. Separately, eight random held-out directions per state test whether the Gram actually predicts response sizes it was not fit to, scored as WAPE against the true centered response.

The whole procedure is repeated at three perturbation scales, and in four output spaces (physical pose, pose plus velocity, RGB pixels, embeddings) to see how much the answer depends on where you look. Uncertainty comes from 2,000 episode-cluster bootstrap resamples (seed 42); repeatability is checked by running the first 50 eligible anchors three times.

Why This Matters

Impact on research. Standard world-model benchmarks reward accurate prediction at sampled actions and good planning scores. This paper shows that a model can be useful on those measures while its local action geometry is meaningfully misaligned, and that the measured misalignment depends heavily on perturbation size and on the output space used. That argues for reporting action sensitivity as a separate, scale-specific diagnostic rather than folding it into aggregate scores.

Real-world applications:

• Robot manipulation: a controller that trusts a world model's gradient with respect to action inherits the model's local geometry; misaligned sensitive directions can bias fine-grained correction near contacts and grasps.
• Model-based planning and MPC: shooting-based or gradient-based planners query the model at many nearby actions, so relative rather than absolute accuracy at small perturbations is what determines whether the ranking of candidate actions is preserved.
• Sim-to-real transfer: if the model's action sensitivity is roughly half the simulator's (as measured on Cube), plans tuned in the model may systematically under-actuate on hardware.
• World-model training and selection: the trace/gain separation gives a concrete diagnostic — a cheap oracle gain correction reduced Cube's held-out error from 56.4% to 32.3% — that could guide loss design or checkpoint selection.

Industry relevance. Teams deploying learned dynamics models in robotics, autonomous systems or simulation-heavy pipelines need evaluation signals that predict downstream control quality. This work supplies a reusable measurement protocol (matched counterfactuals, trace-normalized Gram comparison, held-out direction validation) and a caution: metrics computed in pixel space or latent space can vary more across perturbation scales than metrics computed on physical state.

Future Directions

1. Training-seed variability. Only one released checkpoint per environment was evaluated, evaluated without training. The authors state that additional checkpoints are needed to estimate variation across training seeds, and that their bootstrap intervals cover sampled anchors and episodes, not training runs.

2. Signed, goal-directed responses and closed-loop control. Gram agreement cannot distinguish $J_{\mathrm{model}} = -J_{\mathrm{env}}$ from $J_{\mathrm{model}} = J_{\mathrm{env}}$, nor can it detect an orthogonal remapping of the output space, and trace normalization removes overall gain. Measuring signed goal-directed responses and planning success would establish whether better sensitivity estimates actually help action selection.

3. Attributing the scale effect. Rendering, encoding and nonlinear dynamics may each contribute to the observed scale dependence; their individual effects remain unresolved. Convergence to a smooth derivative is explicitly untested, and contact events were not labeled along the live counterfactuals, so contact-specific claims from earlier results were not retained.

4. Replay reproducibility. Cube's initial strict replay attempts hit intermittent one-level RGB differences, and the v3 production run stopped at anchor 292 with a baseline embedding replay ratio of 1.193% against a fixed 1% threshold, before an isolated unchanged rerun matched exactly. The cause of the intermittent discrepancy remains unknown, and the authors note this limits claims about reproducibility; the anchor drift itself is also reported (median Cube per-state maximum coordinate difference $5.88\times10^{-4}$ in qpos and $4.10\times10^{-4}$ in qvel, with maxima 0.0262 and 2.2505).

Target Audience

Researchers and engineers working on action-conditioned world models and JEPA-style architectures, especially those evaluating them for robot manipulation and model-based control. It is also relevant to anyone designing evaluation methodology for learned dynamics — the calibration baselines and held-out-direction protocol transfer readily to other models. Readers interested in representation geometry will find the principal-angle and condition-number analyses directly useful. A background in linear algebra and familiarity with latent dynamics models is assumed; the paper does not introduce LeWM or the JEPA framework from scratch.

Authors’ abstract

We examine how closely pretrained LeWorldModel (LeWM) reproduces the effects of small action changes. Using matched simulator and model responses, we measure action sensitivity with Gram matrices at three perturbation scales. Across nearly a thousand states from Reacher and OGBench Cube, the most sensitive action directions differ by median angles of about $36^\circ$ and $42^\circ$, respectively, at the smallest scale. Increasing the perturbation tenfold reduces median normalized Gram discrepancy by roughly 70\% on Reacher and 47\% on Cube. Geometry measured in pixels and embeddings varies more across scales than the chosen physical-state geometry. On held-out directions, environment Grams achieve median embedding response errors of 17\% and 28\%, compared with 26\% and 38\% for an isotropic reference with the same trace. LeWM's errors are higher at this scale. These findings support scale-specific evaluation of action sensitivity; its effect on control remains an open question.