ΒΆPaper Feed

Issue 38 Β· Pick 07 Robotics βœ“ read

Robot World Models Are Not Invariant to How the Actions Are Written

Ahmed Karim, Leon Chlon

A robot action can be written two ways: as an absolute joint target ("move joint 3 to 1.2 rad") or as a delta from the current state ("move joint 3 by +0.1 rad"). Given the current state, each determines the other exactly β€” the paper measures mutual reconstructibility at R^2 = 0.996. Yet a latent world model trained on one encoding and handed the identical commanded trajectory in the other doesn't just get noisier: its predictions become nearly orthogonal to what it would have predicted, retrieval degrades up to 13.4Γ—, and goal-conditioned action selection drops from 53% to 15%. This is a failure to be invariant to notation, in a channel β€” the action channel β€” that the invariance literature has never audited. The paper also delivers a clean test for when such a re-parameterization should be free, and a repair that works, with an unusually honest accounting of which part of the repair buys what.

Two ways to write the same command

Every robot learning pipeline commits to an action parameterization early, usually for reasons that have nothing to do with learning: the controller interface, the teleop rig, whatever the last project used. Position control emits absolute joint targets a_{\mathrm{abs}}; delta control emits offsets a_{\mathrm{rel}} = a_{\mathrm{abs}} - s_t, where s_t is the current joint state. Both conventions are common in public datasets and codebases, and both describe the same physical trajectory.

An action-conditioned world model β€” a network that takes a latent state z_t and an action window and predicts the future latent \hat{z}_{t+\Delta} β€” inherits this choice silently. Nothing in its training signal tells it that the action vector is a notation for a trajectory rather than the trajectory itself. If some artifact of the notation (say, the raw magnitude of absolute joint values) correlates with the future on the training distribution, the optimizer will use it. The paper formalizes this with a minimum-description-length argument: ignoring a distinguishing coordinate is preferred only when the architectural saving of an invariant model covers a per-sample information term, and when the labels never exercise that coordinate β€” which is exactly this case β€” the objective prices it in neither direction. The model is simply unconstrained along the axis, and free to break.

One commanded trajectory absolute targets a_abs deltas from state a_rel = a_abs βˆ’ s_t mutually reconstructible given s_t : RΒ² = 0.996 world model g(z_t, a) two beliefs about the same future: cos = 0.067 (PushT) worst = βˆ’0.377
The failure mode. Both encodings carry identical information about the commanded trajectory (given the joint state), yet the trained predictor's latent forecasts under the two notations are nearly orthogonal β€” on the worst input they point in opposite directions.

Why this isn't the usual brittleness story

The obvious dismissal is "that's just covariate shift β€” of course a model breaks off-distribution." The paper's most useful contribution is the test that blocks this dismissal, because covariate-shift results routinely confound different distribution with less information, and the two demand opposite responses. If the shifted input genuinely carries less about the target, degradation is correct behavior; there's nothing to repair.

The test has two parts. Predictive parity is necessary: each view, used alone, must be equally predictive of the target. Theorem 2 (a three-line proof: an orbit-constant predictor is measurable with respect to the quotient of the view space, and conditioning on a coarsening can only raise risk) prices the failure exactly β€” enforcing invariance across views costs at least the parity gap \delta, the risk difference between the best and worst view. Mutual reconstructibility is a sufficient certificate: if each view determines the other, the quotient generates the same \sigma-algebra as either view, and invariance is free.

Two things make the test bite in practice. First, it rejected three of the four axes the authors themselves proposed β€” photometric jitter (an augmentation, not a group action), camera choice (one camera reached retrieval rank 6.6 where another reached 23.3, so forcing agreement would discard the better sensor), and lossy action-window summaries (task R^2 0.986 vs 0.709 across the hierarchy). Averaging over any of those would have paid the parity gap for nothing. Second, admissibility must be evaluated on the joint input: the action channel alone reconstructs at only R^2 = 0.757 β€” a confident false negative β€” while jointly with the state it reaches 0.996, because a_{\mathrm{rel}} = a_{\mathrm{abs}} - s_t is a bijection only given s_t.

The abs/rel axis passes the test cleanly. So when the model collapses, we know the collapse is unjustified: the second presentation provably tells it nothing less.

How badly it breaks

The experimental setup is deliberately spare, which is its diagnostic virtue. A frozen I-JEPA ViT-H/14 encoder produces latents; a small MLP predicts the \ell_2-normalized latent 4 seconds ahead from (z_t, \text{action window}), trained with InfoNCE. No pixel reconstruction, no encoder fine-tuning β€” so nothing measured can be an augmentation or rendering artifact. The only thing the predictor learns is how actions move latents, and the defect must live in the transition data.

Train the predictor on absolute actions, then hand it the same trajectories written as deltas:

Retrieval rank collapse under the equivalent encoding (lower is better)mean retrieval rank (pool of 384)0501001506.1961.5ALOHA cabinet3.9753.3ALOHA coffee66.2170PushTTrained encodingEquivalent encodingTable 1: 9.9x, 13.4x, and 2.6x collapse. Cross-encoding cosine agreement: 0.370, 0.371, 0.067 (mean); worst case -0.178, -0.096, -0.377.

The two ALOHA datasets (14-DoF bimanual) agree to three decimals in cross-encoding cosine, so the effect isn't task-specific. PushT (a 2-D planar pusher) is worse: mean cosine 0.067 between the two predictions of the same future β€” essentially orthogonal beliefs β€” with a worst case of βˆ’0.377. And it's a control failure, not just a retrieval artifact: scoring 32 candidate action windows under the learned dynamics, the model picks the goal-reaching action 53.1% of the time in its own encoding and 15.4% in the equivalent one.

One more result closes the escape hatch of "scale will fix it." Sweeping training-set size from 150 to 1,200 transitions, the collapse ratio grows monotonically:

More data sharpens the trained encoding, not the equivalent onecollapse ratio (held-out rank / canonical rank)training transitions024681002004006008001,0001,200Collapse ratioSection 5.2 / Figure 3. Cross-encoding cosine stays flat across the sweep (0.357, 0.337, 0.374, 0.370): the disagreement doesn't grow, but its cost does, because the canonical model sharpens while the held-out one doesn't.

The careful reading, which the authors give themselves: the disagreement is flat in n; what grows is its price, since the ratio is a quotient of an improving numerator over a stagnant denominator. The defect is invariant to scale rather than amplified by it. Longer training and more capacity move the cosine the wrong way (0.370 β†’ 0.242 and 0.125 respectively), though the authors flag those runs as partly overfitting and decline to lean on them.

The repair, and where the average goes

Once an axis is certified, the repair is orbit averaging β€” train against both presentations. The interesting question is where the average goes, and the output space decides it.

For probability-valued predictions, averaging the outputs is safe by Jensen: the log loss of a mixture is never worse than the average of its members'. But this world model predicts a direction β€” a unit vector scored by cosine similarity β€” and there the guarantee dies. Proposition 3 gives a stable counterexample:

target y, and u1 (score 1.0) u2 (score βˆ’0.9) u3 (score βˆ’0.9) normalized mean Ε« Ε« scores βˆ’1: strictly worse than every member of the orbit
Why output-averaging fails for direction-valued prediction. With u₁ = (1,0) and uβ‚‚,₃ = (βˆ’0.9, Β±0.436), the sum is (βˆ’0.8, 0), so the normalized mean points at (βˆ’1,0) and scores βˆ’1 β€” below all three members. The inequality is strict, so no perturbation removes the witness.

The mean of disagreeing unit vectors is short and points at their centroid; renormalizing it need not return anything near a member of the orbit. Empirically, a student distilled to the orbit-mean teacher lands at retrieval rank 11.55 β€” roughly twice the error of the alternative.

So the paper averages the objective instead:

\mathcal{L}(\theta) = \frac{1}{K}\sum_v \ell\big(g_\theta(z, \phi_v(a)),\, y\big) + \lambda \cdot \frac{2}{K(K-1)}\sum_{v<w} d\big(g_\theta(z,\phi_v(a)),\, g_\theta(z,\phi_w(a))\big),

where the \phi_v are the equivalent presentations (here K=2: abs and rel), \ell is the task loss, and d is a pairwise cosine-disagreement penalty weighted by \lambda. The predictor is asked to be right under each presentation and to agree with itself β€” never to match a summary of its own outputs.

The accounting in Table 2 is refreshingly candid. Plain two-view pooled training (\lambda = 0) β€” the thing any practitioner would do once the axis is known β€” restores task performance by itself: held-out rank 61.5 β†’ 5.98, action selection 15.4% β†’ 40.1%. The penalty adds no average accuracy over that (\lambda=16 gives 41.1%, one point against a seed SD of four), and the authors say so plainly. What the penalty buys is the tail: worst-case cross-encoding agreement goes from 0.779 to 0.995, at a cost of ~3 points of canonical action selection. And the tail is what compounds. Rolling the model forward in latent space (3 hops of 4 s, feeding its own predictions back), the teacher's worst-case agreement erodes (0.515 β†’ 0.242 on ALOHA cabinet), pooled training slows the erosion (0.972 β†’ 0.896), and only \lambda=16 holds it (1.000 β†’ 0.997). On PushT, pooled training does not repair the axis at all: worst-case agreement is still βˆ’0.348 at hop 2, and only the penalty lifts it to positive territory (+0.139). The harder the case, the more of the work the consistency term does.

There's also a nice economics to \lambda: at n=150, a strong penalty (\lambda=256) is worth βˆ’6.59 rank (free regularization); at n=1200 it costs +0.38. The sign flip is exactly what the MDL decomposition predicts β€” below a data threshold, invariance pays for itself; above it, it's a real cost you buy for the guarantee.

Why this particular axis, and not the others

Section 5.3 makes the point that whether a re-parameterization costs anything is decided by the representation, not the transformation. Two control experiments with analytic structure show dispersion at 10^{-15}–10^{-12} β€” not because those models are robust, but because the featurization already discharges the nuisance (an angle encoded as (\sin M, \cos M) literally cannot see a 2\pi wrap). A molecular-dynamics ablation flips the verdict from "make" to "break" purely by swapping an invariant encoder for raw coordinates on the same data and symmetry.

That's why abs-vs-rel is the hard case rather than a curiosity: for a camera angle or a time origin, you can featurize the nuisance away, and should. Here both parameterizations are natural, neither is canonical, and a network receiving an action window as a bag of numbers has no way to know which convention produced it. There is no featurization to reach for β€” which leaves measuring the gap and training it away. Fortunately the orbit has exactly two elements, so the repair is cheap; Appendix D.2 is honest that this whole procedure dies when orbits are combinatorial (S_N) or continuous.

What to make of it

If this holds at scale, the practical upshot is immediate: any pipeline that mixes datasets or codebases with different action conventions β€” increasingly the norm as labs pool teleop data for large VLA and world-model training β€” is silently feeding a model presentations it treats as different worlds. The fix is one line (pool both encodings, optionally add a disagreement penalty for worst-case guarantees), and the audit is a linear regression plus a retrieval evaluation. The certification framework β€” parity as gate, reconstructibility as certificate, joint-input evaluation β€” is a genuinely reusable tool, and the fact that it rejected three of the authors' own candidate axes is good evidence it's not a rubber stamp.

The reasons for skepticism are the ones the authors mostly list themselves. The predictor is a small MLP on frozen latents, trained on 1,200 transitions with 3 seeds β€” this is a diagnostic instrument, not a modern world model, and it's plausible that a large action-tokenizing VLA trained on millions of trajectories with diverse conventions partially learns the invariance from data heterogeneity alone (though the flat-in-n result cautions against assuming so). Evaluation is entirely offline: action selection over 32 candidates is a proxy for control, and no real robot ever moves. The admissibility probe is linear ridge regression, so it certifies linear relabelings trivially and could reject valid nonlinear bijections. And the PushT teacher is weak in absolute terms (rank 66 of 384), so the most dramatic orthogonality number comes from the least competent model β€” the repair there also costs canonical accuracy (66.2 β†’ 79.2 rank), showing repair is only free where the model has slack.

The section most worth your time is Β§3 (the certification test) together with Appendix F's three-line proof β€” that's the transferable idea, applicable well beyond robotics wherever you're tempted to enforce an invariance you haven't verified is one. If you only have five minutes, read Β§3 and Table 2, and note the sentence the paper deserves credit for writing: "the penalty does not improve average task performance over averaging alone, and we do not claim that it does."