1 Introduction↩︎

World models aim at predicting the consequences of actions in a compact latent space so that agents can plan, imagine, and generalize. Joint-Embedding Predictive Architectures (JEPAs) [7]–[10] pursue this in feature space, forgoing pixel-level reconstruction. The recent LeWM method [1] establishes an end-to-end pixel JEPA with a two-term loss, next-embedding MSE plus a Sketch Isotropic Gaussian Regularizer (SIGReg) [11] on the marginal latent distribution, and demonstrates competitive planning with \(\sim\)​15M parameters trainable on a single GPU.

A recurring observation across LeWM and adjacent work is that these latents capture slow features and exhibit emergent temporal straightening [12], [13], but they do not carry an interpretable notion of progression: there is no axis along the latent whose value corresponds to “how far along” the agent is in an episode, nor any geometric prior that different episodes of the same task should share a common progression manifold.

A natural way to add a designated progression axis is to allocate a low-dimensional subspace of the latent to it, and shape that subspace with a temporal-ordering objective. Cosine-margin triplet losses are a well-established mechanism for this purpose in temporal contrastive learning [14], [15]: they pull embeddings of temporally adjacent frames together and push apart frames that are far in time, thereby imposing a smooth ordering on the resulting representation. Independently, recent work on temporal straightening for latent planning [2] demonstrates that geometric regularisers on latent trajectories can substantially improve gradient-based planners. What has not been studied is whether such structured-progression objectives can be combined with the marginal-distribution anti-collapse mechanisms used by end-to-end JEPAs (specifically, SIGReg in LeWM) without the two regularisers interfering.

This paper asks: can we add a structured progression coordinate to a LeWM-style JEPA latent without breaking its SIGReg anti-collapse guarantee, and does the resulting architecture improve planning?

1.0.0.1 Contributions.

1. We introduce SD-JEPA, an extension of LeWM that splits the JEPA latent into a low-dimensional progression subspace \(\boldsymbol{z}^{\mathrm{prog}}\) and a high-dimensional content subspace \(\boldsymbol{z}^{\mathrm{cont}}\), shaped respectively by a cosine-margin triplet loss and SIGReg.

2. We prove that the two anti-collapse forces (SIGReg on \(\boldsymbol{z}^{\mathrm{cont}}\), triplet on \(\boldsymbol{z}^{\mathrm{prog}}\)) act on disjoint coordinates with disjoint gradient supports, so they compose additively rather than competing on the same dimensions; we connect the triplet loss and SIGReg via the sample/dimension-contrastive duality of [16], and validate the decomposition empirically with a subspace-ablation falsifier.

3. We propose a planning cost that decomposes into a content term (\(\boldsymbol{z}^{\mathrm{cont}}\)-MSE) and an optional angular term on \(\boldsymbol{z}^{\mathrm{prog}}\), and show that SD-JEPA matches or improves on LeWM on the majority of its control benchmarks at matched compute.

4. We characterise the resulting 1-D angular progression coordinate \(\theta_t = \mathrm{atan2}(\boldsymbol{z}^{\mathrm{prog}}_2, \boldsymbol{z}^{\mathrm{prog}}_1)\) as a scene-aware compass on the latent: it tracks task progress, regresses under genuine task regression, and under controlled perturbations decouples the moment of surprise from its post-perturbation re-localisation.

2 Background↩︎

2.1 Notation and setting↩︎

We follow the LeWM setting. An offline dataset of trajectories \(\{(o_{1:T_u}, a_{1:T_u})\}_{u=1}^{U}\) is given, with pixel observations \(o_t\) and action labels \(a_t\). An encoder \(f_\theta\) maps observations to latents \(\boldsymbol{z}_t = f_\theta(o_t) \in \mathbb{R}^D\), and a predictor \(p_\phi\) models dynamics autoregressively in latent space: \(\hat{\boldsymbol{z}}_{t+1} = p_\phi(\boldsymbol{z}_{t-H+1:t}, a_t)\), where \(H\) is the history length. Planning is performed by CEM over action sequences to minimize a latent-space cost to a goal embedding \(\boldsymbol{z}_g = f_\theta(o_g)\).

2.2 LeWM: next-embedding prediction with SIGReg↩︎

LeWM trains the encoder and predictor jointly with \[\mathcal{L}_{\text{LeWM}} = \underbrace{\|\hat{\boldsymbol{z}}_{t+1} - \boldsymbol{z}_{t+1}\|_2^2}_{\mathcal{L}_{\text{pred}}} \;+\; \lambda \,\mathrm{SIGReg}(\boldsymbol{z}), \label{eq:lewm}\tag{1}\] where \(\mathrm{SIGReg}\) enforces that the marginal distribution of \(\boldsymbol{z}\) is isotropic Gaussian. Using the Cramér–Wold theorem, SIGReg aggregates the univariate Epps–Pulley normality statistic over \(M\) random unit directions \(u^{(m)} \in \mathbb{S}^{D-1}\): \[\mathrm{SIGReg}(Z) = \frac{1}{M} \sum_{m=1}^{M} T\!\left(Zu^{(m)}\right), \label{eq:sigreg}\tag{2}\] with \(T(\cdot)\) the Epps–Pulley test statistic. The key property of 2 is that it constrains only the marginal distribution \(p(\boldsymbol{z})\); the conditional \(p(\boldsymbol{z}_{t+1}\mid\boldsymbol{z}_t)\) is unaffected.

2.3 Duality of contrastive and covariance-based regularizers↩︎

[16] show that the sample-contrastive criterion \(\mathcal{L}_c = \|K^\top K - \mathrm{diag}(K^\top K)\|_F^2\) and the dimension-contrastive criterion \(\mathcal{L}_{nc} = \|KK^\top - \mathrm{diag}(KK^\top)\|_F^2\) (with \(K\) the embedding matrix) are equivalent up to the row- and column-norms of \(K\): \[\mathcal{L}_{nc} + \sum_{j=1}^{M} \|K_{j,\cdot}\|_2^4 = \mathcal{L}_c + \sum_{i=1}^{N} \|K_{\cdot,i}\|_2^4. \label{eq:duality}\tag{3}\] Under double-\(L_2\)-normalization of rows and columns, \(\mathcal{L}_{nc} = \mathcal{L}_c + N - M\). SIGReg and the cosine-triplet both belong to this broader variance-repulsion family, a fact we exploit in §4.

3 Method: SD-JEPA↩︎

3.1 Subspace decomposition of the latent↩︎

We split the JEPA latent into a progression subspace and a content subspace: \[\boldsymbol{z}_t = \underbrace{P\, \boldsymbol{z}^{\mathrm{prog}}_t}_{\in \mathbb{R}^D} \;+\; \underbrace{Q\, \boldsymbol{z}^{\mathrm{cont}}_t}_{\in \mathbb{R}^D}, \qquad \boldsymbol{z}^{\mathrm{prog}}_t \in \mathbb{R}^k,\;\boldsymbol{z}^{\mathrm{cont}}_t \in \mathbb{R}^{D-k},\] where \(P \in \mathbb{R}^{D\times k}\) and \(Q \in \mathbb{R}^{D\times(D-k)}\) are fixed orthogonal injections with \(P^\top P = I_k\), \(Q^\top Q = I_{D-k}\), \(P^\top Q = 0\). In practice we take the canonical split \(P = [I_k;\, 0]\), \(Q = [0;\, I_{D-k}]\), so that \(\boldsymbol{z}^{\mathrm{prog}}_t\) and \(\boldsymbol{z}^{\mathrm{cont}}_t\) are simply the first \(k\) and last \(D-k\) coordinates of \(\boldsymbol{z}_t\). We use \(k = 2\) by default.

3.2 Four-term training objective↩︎

SD-JEPA is trained with \[\mathcal{L}_{\text{SD}} \;=\; \mathcal{L}_{\text{pred}}(\boldsymbol{z}) \;+\; \lambda_S\,\mathrm{SIGReg}(\boldsymbol{z}^{\mathrm{cont}}) \;+\; \lambda_T\,\mathcal{L}_{\text{trip}}(\boldsymbol{z}^{\mathrm{prog}}) \;+\; \lambda_{\mathrm{str}}\,\mathcal{L}_{\text{straight}}(\boldsymbol{z}) \label{eq:objective}\tag{4}\] with:

• \(\mathcal{L}_{\text{pred}}(\boldsymbol{z}) = \|\hat{\boldsymbol{z}}_{t+1} - \boldsymbol{z}_{t+1}\|_2^2\) on the full latent (unchanged from LeWM).

• \(\mathrm{SIGReg}(\boldsymbol{z}^{\mathrm{cont}})\): SIGReg (2 ) applied only to the content subspace, using random unit directions sampled in \(\mathbb{R}^{D-k}\).

• \(\mathcal{L}_{\text{trip}}(\boldsymbol{z}^{\mathrm{prog}})\): a cosine-margin triplet loss on \(\boldsymbol{z}^{\mathrm{prog}}\) that orders embeddings by their position within a trajectory: positives are sampled within a temporal window \(\vartheta_t\) of an anchor, negatives outside the window or from a different trajectory.

• \(\mathcal{L}_{\text{straight}}(\boldsymbol{z}) = -\mathbb{E}_t\!\left[\cos(\boldsymbol{v}_t, \boldsymbol{v}_{t+1})\right]\) with \(\boldsymbol{v}_t = \boldsymbol{z}_{t+1} - \boldsymbol{z}_t\): explicit temporal straightening (see §9.3).

The LeWM base is recovered as \(\lambda_T = \lambda_P = \lambda_{\mathrm{str}} = 0\) and applying \(\mathrm{SIGReg}\) to the full \(\boldsymbol{z}\) instead of \(\boldsymbol{z}^{\mathrm{cont}}\).

3.3 Latent-space indicators as predictor conditioning↩︎

From \(\boldsymbol{z}^{\mathrm{prog}}_t\) we compute the angular and radial readouts \[\theta_t = \mathrm{atan2}(\boldsymbol{z}^{\mathrm{prog}}_{t,2},\, \boldsymbol{z}^{\mathrm{prog}}_{t,1}), \qquad r_t = \|\boldsymbol{z}^{\mathrm{prog}}_t\|_2.\] These are scalar, scale- and rotation-aware quantities that summarize where on the progression manifold the current state sits. We expose them to the autoregressive predictor as additional conditioning alongside the action embedding: \[\hat{\boldsymbol{z}}_{t+1} = p_\phi\!\left(\boldsymbol{z}_{t-H+1:t},\; a_t,\; (\theta_t, r_t)\right).\] Concretely, \((\theta_t, r_t)\) (with \(\theta_t\) encoded as \((\sin\theta_t, \cos\theta_t)\) to avoid wrap-around discontinuities) is embedded to the predictor’s hidden dimension and summed with the action embedding before AdaLN-zero conditioning. This gives the predictor an explicit progression compass, analogous in spirit to sinusoidal positional encodings in attention: no new information, but better-structured information.

3.4 Planning with an angular goal-matching cost↩︎

In LeWM, planning minimizes \(\|\hat{\boldsymbol{z}}_H - \boldsymbol{z}_g\|_2^2\) at the end of a rollout. We replace this with a decomposed cost: \[C(\hat{\boldsymbol{z}}_H) = \underbrace{\|\hat{\boldsymbol{z}^{\mathrm{cont}}}_H - \boldsymbol{z}^{\mathrm{cont}}_g\|_2^2}_{\text{content match}} \;+\; \gamma\,\underbrace{\bigl(1 - \cos(\hat{\theta}_H, \theta_g)\bigr)}_{\text{angular progression match}} \;+\; \delta\,\underbrace{(\hat{r}_H - r_g)^2}_{\text{radial match}}. \label{eq:plancost}\tag{5}\] The angular and radial terms make the progression match scale-invariant and independent of the ambient content geometry. The receding-horizon Cross-Entropy Method (CEM) solver is otherwise identical to LeWM: at each replanning step, \(300\) candidate action sequences are sampled, the top-\(30\) retained as elites to update the sampling Gaussian, and we iterate \(30\) times on Push-T and \(10\) times on the other environments.

4 Theoretical analysis↩︎

We state the two propositions that establish why the subspace split composes additively rather than competing on the same dimensions. Full proofs and three further results (consistency with emergent straightening, identifiability of the progression axis, when the split fails) are deferred to App. 9.

Proof sketch. \(\mathrm{SIGReg}(Q^\top\boldsymbol{z})\) depends on \(\boldsymbol{z}\) only through \(Q^\top\boldsymbol{z}\), so by the chain rule \(\nabla_{\boldsymbol{z}}\, \mathrm{SIGReg}(Q^\top\boldsymbol{z}) = Q \,\nabla_{\boldsymbol{z}^{\mathrm{cont}}}\mathrm{SIGReg}(\boldsymbol{z}^{\mathrm{cont}})\), which lies in \(\mathrm{col}(Q)\). Identically for \(\mathcal{L}\) and \(P\). Orthogonality follows from \(P^\top Q = 0\). \(\square\)

These two results together formalise the architectural intuition: triplet on \(\boldsymbol{z}^{\mathrm{prog}}\) and SIGReg on \(\boldsymbol{z}^{\mathrm{cont}}\) are complementary rather than redundant despite belonging to the same variance-repulsion family in the sense of [16]. The empirical falsifier (A2_full \(=\) A0, Tab. 2) confirms the proposition matters in practice: when the split is removed and triplet acts on the full latent, the gain disappears. Disjointness is at the latent level; the encoder’s parameter gradients are still the sum of both terms via the chain rule, so encoder-level conflict is not formally precluded.

5 Experiments↩︎

5.1 Setup↩︎

We evaluate SD-JEPA on the four control benchmarks of [1]: Two-Room (2-D navigation), Reacher (DM-Control), Push-T (2-D manipulation), and OGBench-Cube (3-D manipulation). All evaluations use the paper’s protocol, \(50\)-step horizon, \(25\)-step goal offset, CEM solver with \(300\) candidates / \(30\) iterations on Push-T and \(10\) iterations elsewhere, planning horizon \(5\) at frame-skip \(5\). We train each model for 10 epochs at batch size \(128\), matching the LeWM training budget; encoder and predictor architectures are inherited from LeWM (ViT-tiny encoder, \(6\)-layer transformer predictor, \({\sim}18\)M parameters total). Multi-seed runs use seeds \(\{0, 42, 3072\}\). Implementation details, dataset paths, hyperparameters, and a glossary of the ablation rung names (A0, A2, A2_full, A4, A5, A6) used throughout this section are in App. 11 and Tab. 6.

5.2 Main results↩︎

Figure 1 shows planning success rate per environment for SD-JEPA against the published LeWM numbers and the cross-method baselines reported by [1]. SD-JEPA improves over LeWM on three of the four environments at matched 10-epoch compute:

• Reacher: \(+2\) over LeWM (\(86 \to 88\) at \(k_\text{prog}{=}4\), \(3\)-seed mean across both planning costs); \(+12\) over our A0 rerun (\(76\)).

• Two-Room: \(+3\) over LeWM (\(87 \to 90\) at \(k_\text{prog}{=}2\), \(3\)-seed mean).

• Push-T: \(+1.3\) over LeWM (\(96 \to 97.3\) at \(k_\text{prog}{=}8\), \(3\)-seed mean),\(+5\) over the strongest non-LeWM baseline (DINO-WM with proprioception, \(92\)).

• OGBench-Cube: \(-2\) vs LeWM at \(k_\text{prog}{=}2\) (\(72\) vs \(74\), \(3\)-seed mean); the \(k_\text{prog}{=}\{4, 8\}\) sweep does not close the gap (cube prefers \(k{=}2\) in our 10-epoch regime).

The optimal \(k_\text{prog}\) is task-dependent (§5.4); the architectural framework is robust across \(k \in \{2, 4, 8\}\). Table 1 reports cross-method numbers.

5.3 The subspace split is load-bearing↩︎

To isolate which component of SD-JEPA drives the gain, two falsifiers on Push-T (Tab. 2, single seed). A2_full disables the split: \(96 =\) A0. A2_split_full keeps the split but mis-targets the triplet on the full latent: \(92\), four points below A0. The split itself, not the triplet in isolation, is the load-bearing ingredient: the empirical analogue of Prop. 2 on the actual training dynamics.

5.4 \(k_{\text{prog}}\) scales with task complexity↩︎

The optimal \(k_\text{prog}\) varies per env (Tab. 3). On Push-T, multi-seed evaluation gives a clean monotone \(94 \to 96 \to 97.3\) at \(k \in \{2, 4, 8\}\), with the \(k{=}8\) checkpoint exhibiting a near-perfect \(S^1\) progression manifold (\(\sigma_r \in [0.04, 0.07]\), angular span \(6.1\) rad on seven held-out episodes; App. 13). Two-Room is non-monotone (\(k{=}2\) and \(k{=}8\) both \(90\), \(k{=}4\) at \(88\)), consistent with the 1-D navigation topology saturating at two progression dimensions; Reacher prefers a moderate \(k{=}4\) (\(92\) vs.\(84\)/\(82\)). The cross-env pattern is not “higher \(k\) always wins” but “the right \(k\) matches the task’s intrinsic progression dimensionality.” A negative control with \(\boldsymbol{z}^{\mathrm{prog}}\)-only MSE planning collapses to \(28\%\) on Push-T: the progression subspace is too low-dimensional to carry goal information alone; the gain at higher \(k\) comes from the encoder devoting more budget to a coherent progression manifold while the planning cost still operates primarily on \(\boldsymbol{z}^{\mathrm{cont}}\).

5.5 The latent acts as a scene-aware compass↩︎

The canonical SD-JEPA architecture exposes a 1-D angular readout \(\theta_t = \mathrm{atan2}(\boldsymbol{z}^{\mathrm{prog}}_{t,2}, \boldsymbol{z}^{\mathrm{prog}}_{t,1})\) on the trained progression subspace. We probe what this coordinate captures along three axes: the cross-episode geometry it produces in \(\boldsymbol{z}^{\mathrm{prog}}\) versus \(\boldsymbol{z}^{\mathrm{cont}}\), the physical task quantities it correlates with across environments, and how it behaves under controlled observation perturbations. The picture is environment-dependent and partly clock-like, partly task-coupled; the full per-environment breakdown is in App. 13.

5.5.0.1 Cross-episode geometry matches the disjoint-supports prediction.

Prop. 1 predicts that \(\boldsymbol{z}^{\mathrm{prog}}\) should carry shared cross-episode progression structure while \(\boldsymbol{z}^{\mathrm{cont}}\) should be episode-specific. The Push-T A2 t-SNE (Fig. 4, App. 13) makes the asymmetry visible: \(\boldsymbol{z}^{\mathrm{cont}}\) separates into per-episode clusters while \(\boldsymbol{z}^{\mathrm{prog}}\)’s arcs partially mix across episodes. The second moments confirm it: \(\overline{\cos}(\boldsymbol{z}^{\mathrm{prog}}[t]) \approx 0.5\) vs.\(\overline{\cos}(\boldsymbol{z}^{\mathrm{cont}}[t]) \approx 0\), with the full \(\boldsymbol{z}\) inheriting the \(\boldsymbol{z}^{\mathrm{cont}}\)-side near-orthogonality.

5.5.0.2 \(\theta_t\) is partly a clock and partly a task-phase coordinate, with the balance varying by env and \(k_{\mathrm{prog}}\).

We score the per-episode Spearman \(\rho\) between \(\theta_t\) and candidate proxies: the elapsed-time clock (step_idx) plus env-specific physical quantities (block_target_dist on Cube, ee_target_dist on Reacher, block_angle_err on Push-T). Mean \(|\rho|\) tables, heatmaps, and per-episode signed values are in App. 13. Cube is the only env where a task-physical signal beats the clock (block_target_dist, \(|\rho|{=}0.59\) vs \(0.47\)); elsewhere the two are within \(0.1\) or the clock leads. On Push-T, \(|\rho|\) with the clock falls from \(0.91\) at \(k{=}2\) to \(0.56\) at \(k{=}8\) as the physical signals catch up, consistent with \(\theta\) becoming more cyclic-phase-like at higher \(k\). The cube state-traj overlay (Fig. 4, App. 13) makes one episode concrete: \(\theta\) traces a smooth phase gradient along the planar trajectory.

5.5.0.3 \(\theta_t\) as a richer surprise signal: separating the moment of surprise from its meaning.

Standard latent VoE, \(\|\hat{\boldsymbol{z}}_{t+1}-\boldsymbol{z}_{t+1}\|\), says when the model was wrong but not about what; identical scalar magnitudes may correspond to different counterfactual interpretations. We test the angular signal on Push-T with a teleport-and-continue perturbation: at step \(50\) of held-out episode \(0\), the next observation is replaced by a frame from episode \(1000\) (whose T-block sits on the opposite side of the goal) and the planner resumes. Two distinct signals appear in \(\theta_t\) (Fig. 3). (i) A sharp single-frame spike of \({\sim}1.5\) rad in \(|\Delta\theta_t|\) (\(\approx 75\times\) baseline drift) marks the moment, like a scalar VoE spike. (ii) After the spike, \(\theta_t\) does not return: it relocalises from sector \(\theta\!\approx\!-2.5\) to sector \(\theta\!\approx\!-4.0\), corresponding to the new solving mode (T-block on the other side), marking the meaning, the model’s updated belief about which task phase is active. Standard MSE-VoE reports only the spike, not the re-localisation. Three further perturbation modes (splice, 1-frame teleport, reverse) reproduce this pattern (App. 14).

5.5.0.4 \(|\Delta\theta_t|\) outperforms scalar prediction-error at localising semantic phase events.

The qualitative argument above motivates a quantitative test: does \(|\Delta\theta_t|\) actually pinpoint task-meaningful events more reliably than the latent prediction-error magnitude \(\|\hat{\boldsymbol{z}}_{t+1} - \boldsymbol{z}_{t+1}\|\) (z-MSE), the standard surprise metric in latent world models? On OGBench-Cube the gripper-contact transitions provide an objective ground-truth event signal: each pick-and-place episode has roughly four such transitions (close, open, close, open) marking semantically meaningful phase boundaries. We compute the per-step AUROC of each surprise metric against labels marking “step is within \(\pm \mathrm{tol}\) of any contact transition” on \(40\) held-out episodes (\(160\) ground-truth events, \(1480\) steps; SD-JEPA \(k_{\mathrm{prog}}{=}8\) checkpoint). Pooled AUROC and the per-episode head-to-head appear in Tab. 4; the qualitative phase overlay on episode \(500\) is shown in Fig. 2 and the four-panel summary across all episodes in Fig. 14 (App. 15).

5.5.0.5 \(\boldsymbol{z}^{\mathrm{prog}}\) packs within-episode progress information densely.

A complementary test asks how much task-progress information is packed into the progression subspace per dimension. We fit linear-regression probes from various latent features to a per-episode-normalised target-distance signal, in two regimes: a per-episode probe (within each held-out episode, leave-one-step-out CV; mean R\(^2\) over \(40\) episodes) and a pooled probe (leave \(25\%\) of episodes out, \(30\) bootstraps). On cube, the per-episode probe gives \(\boldsymbol{z}^{\mathrm{prog}}\) (\(8\) dims, \(4.2\%\) of the latent) mean R\(^2 = \mathbf{0.905}\) (\(\mathbf{100\%}\) positive over \(40\) episodes); \((\sin\theta, \cos\theta)\) packs \(\mathbf{55.5\%}\) of the variance into just \(2\) dimensions, \(2\times\) a random-\(2\)-d projection of \(\boldsymbol{z}\) (\(0.263\)) and the elapsed-time clock (\(0.291\)). Cross-env replication on the same \(40\)-episode-per-env protocol confirms the claim is not cube-specific (Tab. 5): \(\boldsymbol{z}^{\mathrm{prog}}\) R\(^2 \in \{0.91, 0.91, 0.95, 0.72\}\) on cube/Push-T/Reacher/Two-Room, \(100\%\) positive on the first three and \(95\%\) on Two-Room, never below \(0.72\) across any env using only \(4.2\%\) of the latent. The largest probe-vs-clock gap, \(+0.66\) on Reacher, lines up with Reacher being the only env with a robust \(\boldsymbol{z}^{\mathrm{prog}}\)-aware planning lift (§5.6, App. 12). The pooled probe is the honest caveat: across episodes with different absolute initial distances, only step_idx stays reliably positive (App. 15, Tab. 9, Fig. 16, 17); the compass is a per-trajectory phase coordinate, not a globally calibrated cross-episode distance estimate.

5.5.0.6 Three operationalisations of the same intuition.

The phase-event AUROC (Tab. 4, \(|\Delta\theta_t|\) wins on \(39/40\) cube episodes) asks where the semantic events are; the regime-change CPD (App. 15) asks where the regime boundaries are; the linear probe (Tab. 9, App. 15) asks how much progress information is packed per dimension. All three confirm that the trained \(\boldsymbol{z}^{\mathrm{prog}}\) subspace, and \(\theta\) in particular, carries task-phase structure that the standard scalar prediction-error metric does not surface.

5.6 Cross-environment summary and limitations↩︎

SD-JEPA improves on three of four LeWM envs (Tab. 1, Tab. 3, \(3\)-seed multi-seed); the framework is robust across \(k_\text{prog} \in \{2, 4, 8\}\) but the per-env optimum varies. Whether full-z planning helps is also task-dependent (Reacher \(+3.3\) pp at \(k{=}8\), other envs within \(\pm 1\) pp; App. 12). The latent-geometry diagnostics establish that the architecture produces the predicted progression-content geometry independently of any planning-success delta, providing a second axis of evidence on top of the empirical one.

5.6.0.1 Known limitations.

(i) Per-environment \(k_\text{prog}\) tuning means we report best-\(k\) values rather than a single universal \(k\). (ii) The triplet and angular readouts both use cosine geometry; other distance metrics or non-angular readouts (e.g.a learned scalar progression head) may surface complementary structure. (iii) The fixed orthogonal \([P, Q]\) is a modeling choice; a learned orthogonal decomposition is the natural extension. Extended ablations on auxiliary terms (straightening, polar conditioning, angular planning cost) are in App. 12.

6 Related work↩︎

6.0.0.1 Latent world models.

We complete the related work mentioned in the introduction with pixel-space architectures such as DreamerV3 [17] and feature-space JEPAs include I-JEPA, V-JEPA, PLDM, DINO-WM, and most directly LeWM [1], [7]–[9], [13]. Latent-action world models such as LAWM, Genie, Gaia-2, Cosmos and explore alternative conditioning regimes [18]–[23]. SD-JEPA inherits LeWM’s encoder-predictor stack and the SIGReg anti-collapse regularizer, adding only a fixed orthogonal subspace split and a triplet term on the progression side; the parameter count is identical to LeWM (\(18.04\)M, no new learnable weights).

6.0.0.2 Progression geometry and temporal straightening.

The cosine-margin triplet on \(\boldsymbol{z}^{\mathrm{prog}}\) is a temporal-ordering loss in the family of [14], [15]. Our straightening regulariser \(\mathcal{L}_{\mathrm{straight}} = -\mathbb{E}[\cos(\boldsymbol{v}_t,\boldsymbol{v}_{t+1})]\) is identical to the curvature loss of [2] (stop-gradient anti-collapse, gradient-descent planner). On Push-T at the matched protocol, their best published configuration reaches \(91.3 \pm 2.5\); SD-JEPA at \(k_\text{prog}{=}8\) reaches \(97.3 \pm 1.2\), \(6\) points above. Adding the curvature regulariser on top of SIGReg in our setup (rung A4, App. 12) costs \(14\) points on Push-T, consistent with the principle that anti-collapse forces should not compete on the same coordinates. Latent straightening has applications in OOD detection [24]. The SSL duality framework of [16] grounds our no-double-counting result.

7 Discussion↩︎

7.0.0.1 The latent as a navigable map.

The diagnostics in §5.5 support reading \(\theta_t\) as less of a chronological counter and more of a state-based navigation signal: it advances with task progress, regresses in lockstep when the agent backtracks, and under perturbation jumps to a semantically appropriate new sector. The qualitative parallel with mammalian head-direction cells [25], [26] is suggestive; we make no mechanistic claim. App. 14 reproduces the moment-and-meaning decoupling under three further perturbation modes.

7.0.0.2 When does \(\theta\)-coupling translate into a planning lift?

The full-z planning lift is non-uniform: Reacher gains \(+3.3\) pp at \(k_{\mathrm{prog}}{=}8\), Cube \(+2.0\), Two-Room and Push-T nothing. Cross-referencing per-environment \(\theta\)-task coupling (App. 13), \(\theta\)-coupling is necessary but not sufficient: Reacher (coupling \(|\rho|=0.49\)) sees the cleanest lift; Two-Room (\(0.34\)) sees none. The two anomalies have separate bottlenecks: Cube has the strongest coupling (\(0.59\)) but the predictor’s contact-rich rollout error caps the cost-function gain, while Push-T at \(k_{\mathrm{prog}}{=}8\) is already near the ceiling (\(97.3\%\)). Reacher sits in the sweet spot: phase-coupled enough for the angular cost to add information, simple-enough dynamics that the predictor is not the bottleneck. The cross-environment probe (Tab. 5) gives a complementary mechanistic story for the same pattern: Reacher has the largest within-episode probe-vs-clock gap (\(+0.66\) R\(^2\), the largest in our suite), so the encoder concentrates the most progress information into \(\boldsymbol{z}^{\mathrm{prog}}\) exactly where the planning cost can exploit it.

7.0.0.3 What \(\theta_t\) adds beyond scalar VoE.

The angular readout differs from \(\|\hat{\boldsymbol{z}}_{t+1}-\boldsymbol{z}_{t+1}\|\) in three ways: it reports a phase sector, it is signed, and it is cross-environment comparable on a unit circle. The three tests of §5.5 confirm complementarity rather than competition with VoE; the broader design space of geometric readouts of structured-latent world models is an open agenda.

8 Conclusion↩︎

SD-JEPA splits the JEPA latent into a low-dimensional progression subspace shaped by a cosine-margin triplet loss and a high-dimensional content subspace regularised by SIGReg. We prove the two anti-collapse forces act on disjoint coordinates and validate this empirically: removing the split collapses the architectural gain back to the LeWM baseline. At matched 10-epoch compute (3-seed multi-seed) SD-JEPA improves on three of four LeWM control benchmarks (Two-Room \(+3\), Reacher \(+2\), Push-T \(+1.3\); cube \(-2\)). Beyond planning, the 1-D angular coordinate \(\theta_t\) behaves as a scene-aware compass that tracks progress, regresses on backtrack, and under perturbations separates the moment of surprise from its meaning. Three independent quantitative tests back this up: \(|\Delta\theta_t|\) beats z-MSE at semantic-event localisation on \(40\) cube episodes (\(+0.18\) pooled AUROC at \(\pm 1\)-step tolerance, \(39/40\) per-episode wins), change-point detection on the angular-velocity trace beats prediction-error metrics at loose tolerance, and a within-episode linear probe replicates across all four envs (\(40\) episodes per env) with \(\boldsymbol{z}^{\mathrm{prog}}\) R\(^2 \in [0.72, 0.95]\) using only \(4.2\%\) of the latent. All three back the same conclusion: the angular readout extends the standard interpretability toolkit with information the scalar prediction-error metric does not supply.

9 Extended theoretical analysis↩︎

This appendix reproduces the full theoretical content compressed into §4, including the corollary on the marginal product structure, the Garrido-style duality discussion, and three further results (consistency with emergent straightening, identifiability of the progression axis, and conditions under which the split fails).

9.1 Non-conflict under the subspace split↩︎

Proposition 1 formalises the informal claim that an arbitrary loss on \(\boldsymbol{z}^{\mathrm{prog}}\) alone does not interfere with SIGReg on \(\boldsymbol{z}^{\mathrm{cont}}\). The triplet loss on \(\boldsymbol{z}^{\mathrm{prog}}\) and the prediction loss on the full \(\boldsymbol{z}\) do have components in \(\mathrm{col}(Q)\) (through the content subspace of the predictor output), so the non-conflict is specifically between the marginal-distributional term on \(\boldsymbol{z}^{\mathrm{cont}}\) and any progression-geometric term on \(\boldsymbol{z}^{\mathrm{prog}}\).

9.2 Relation to the sample/dimension-contrastive duality↩︎

The cosine-margin triplet loss on \(\boldsymbol{z}^{\mathrm{prog}}\) is a sample-contrastive criterion in the sense of [16]: it repels \(\boldsymbol{z}^{\mathrm{prog}}_a\) from \(\boldsymbol{z}^{\mathrm{prog}}_n\) while attracting it to \(\boldsymbol{z}^{\mathrm{prog}}_p\), in a space of \(L_2\)-normalized embeddings (cosine similarity is the inner product of unit vectors). SIGReg on \(\boldsymbol{z}^{\mathrm{cont}}\) is not exactly a dimension-contrastive criterion, but it belongs to the same variance-repulsion family: it penalizes any departure of the marginal from isotropic Gaussian, and in particular penalizes any direction along which the variance is abnormally low (collapse) or the covariance is nonzero (feature correlation). Under the double-\(L_2\)-normalization assumption of [16], the anti-collapse force of the triplet loss on \(\boldsymbol{z}^{\mathrm{prog}}\) equals (up to additive constants) that of a VICReg-style dimension-contrastive criterion. SIGReg subsumes that variance/covariance criterion as a special case (matching the second moments of an isotropic Gaussian).

Proposition 2 is the practical consequence of the duality: even though triplet and SIGReg are of the same family (in the Garrido sense), on the subspace split they are complementary rather than redundant. This gives a clean theoretical justification for using both terms simultaneously rather than treating one as subsuming the other.

9.3 Consistency with emergent straightening↩︎

LeWM reports (Sec. 5 and Appendix H of [1]) that as training proceeds, the cosine similarity between consecutive latent velocity vectors \(\boldsymbol{v}_t = \boldsymbol{z}_{t+1} - \boldsymbol{z}_t\) increases emergently, without any explicit regularization. This perceptual straightening property, inspired by [12], is an invariant of good representations of natural dynamics: straight paths in latent space are linearly extrapolable, and so compound less error over multi-step rollouts.

Our straightening term \(\mathcal{L}_{\text{straight}}(\boldsymbol{z}) = -\mathbb{E}[\cos(\boldsymbol{v}_t, \boldsymbol{v}_{t+1})]\) makes this trend explicit. It is:

• scale-invariant in \(\boldsymbol{z}\): cosine depends only on directions of velocity, not magnitudes; this means \(\mathcal{L}_{\text{straight}}\) does not compete with the variance-related component of \(\mathrm{SIGReg}\);

• subspace-consistent: applying \(\mathcal{L}_{\text{straight}}\) on the full \(\boldsymbol{z}\) does not prefer one of \(\boldsymbol{z}^{\mathrm{prog}},\boldsymbol{z}^{\mathrm{cont}}\) over the other; in practice we may apply it only on \(\boldsymbol{z}^{\mathrm{cont}}\) if \(\boldsymbol{z}^{\mathrm{prog}}\) is already sufficiently constrained by the triplet ordering;

• a soft regularizer: it does not force \(\cos=1\) (perfectly straight dynamics), only moves the population mean upward.

Empirically this is the weakest of our claims, but it rests on the reported emergence from previous works. It motivates treating straightening as a warm-start regularizer rather than a hard constraint.

9.4 Identifiability of the progression axis↩︎

Under the canonical \((P, Q)\) split with \(k = 2\), the angular coordinate \(\theta_t = \mathrm{atan2}(\boldsymbol{z}^{\mathrm{prog}}_2, \boldsymbol{z}^{\mathrm{prog}}_1)\) is identifiable up to a global rotation (the triplet loss enforces a cyclic ordering on the unit-circle but does not pin a unique reference direction). This rotation freedom is harmless when the planning cost itself is rotation-invariant (the canonical \(\boldsymbol{z}^{\mathrm{cont}}\)-MSE cost is, by construction).

The identifiability is local to episodes: episodes with different durations or solving modes can each present a coherent angular progression in \(\boldsymbol{z}^{\mathrm{prog}}\) without any global reference direction shared across them.

9.5 When does the subspace split fail?↩︎

The construction above assumes the progression subspace \(\boldsymbol{z}^{\mathrm{prog}}\) is a reasonable \(k\)-dimensional summary of the episode phase. This assumption fails when:

1. The environment lacks coherent episodic temporal structure (e.g.purely exploratory rollouts with no consistent task progression). In that case the triplet loss has no useful ordering signal and the split degenerates to the LeWM baseline.

2. The task has multiple branching progressions (e.g.different goals per episode with no shared notion of completion). The triplet loss can still impose local orderings within each branch, but a single \(k\)-dimensional progression manifold may not accommodate the branching cleanly; goal-conditioned variants are a natural extension.

3. \(k\) is too large: the progression subspace absorbs content information, polluting the \((\theta, r)\) reading. We sweep \(k \in \{2, 4, 8\}\) as an ablation.

10 Implementation notes↩︎

We build on the public reference implementation of LeWM [1] for the encoder, predictor, action-conditioning, and SIGReg modules. SD-JEPA additions are localised to (a) computing the canonical injections \(P, Q\) on the projected latent, (b) instantiating the cosine-margin triplet on \(\boldsymbol{z}^{\mathrm{prog}}\), and (c) the optional polar-conditioning / angular-planning-cost paths described in App. 12. The triplet uses within-batch sampling: for each anchor, a positive is drawn from the same trajectory within a temporal window \(\vartheta_t\) and a negative from outside that window or from a different trajectory in the same batch, so triplet computation adds negligible overhead to the standard JEPA training step.

11 Hyperparameter summary↩︎

Per-environment configs, dataset paths, hyperparameter sweeps, and training-time stop conditions are documented in the released code. The base hyperparameters: AdamW, lr \(5\!\times\!10^{-5}\), weight decay \(10^{-3}\), cosine-annealed over the 10 training epochs, batch size \(128\), image resolution \(224 \times 224\), ViT-tiny encoder (patch size \(14\)), \(6\)-layer transformer predictor with AdaLN-zero action conditioning, embedding dimension \(192\). The frame-skip is \(5\) for all environments. Default loss weights: \(\lambda_S = 0.09\) (SIGReg), \(\lambda_T = 0.10\) (triplet). Optional auxiliary weights (\(\lambda_{\text{str}}\) for straightening, \(\gamma_\theta, \delta_r\) for the angular planning cost) default to \(0\); see App. 12 for the A4 / A5 / A6 sweep.

12 Additional auxiliary-term explorations↩︎

While exploring the design space of SD-JEPA we tested several optional auxiliary terms that, in our setup, do not improve the canonical A2 configuration. We document them here for completeness and as a reference for follow-up work. (i) A4 adds the temporal-straightening regulariser \(\mathcal{L}_{\text{straight}} = -\mathbb{E}[\cos(\boldsymbol{v}_t, \boldsymbol{v}_{t+1})]\) on the full latent. On Push-T this rung underperforms A2 by roughly \(14\) points: applying the straightening cost to the full \(\boldsymbol{z}\) overlaps with SIGReg’s isotropic-Gaussian requirement on \(\boldsymbol{z}^{\mathrm{cont}}\), an instance of the double-counting that Prop. 2 predicts. The same regulariser is the central contribution of [2], who use it under stop-gradient anti-collapse where the conflict does not arise (Sec. 6). (ii) A5 adds polar \((\sin\theta, \cos\theta, r)\) conditioning of the predictor’s AdaLN on top of A4 and yields a further \({\sim}2\)-point drop, suggesting that exposing the angular readout to the predictor at training time over-specialises the content channel without compensating gains. (iii) A6 additionally engages the angular planning cost \(\gamma_\theta\,(1 - \cos(\hat{\theta} - \theta_g)) + \delta_r\,(\hat{r} - r_g)^2\) on \(\boldsymbol{z}^{\mathrm{prog}}\) at evaluation time. We also tested an auxiliary first-anchor regulariser pulling the genuine first-of-episode frames in each batch toward a fixed reference direction in \(\boldsymbol{z}^{\mathrm{prog}}\), as an attempt to globally calibrate \(\theta_t\) across episodes; this rung returns to the canonical A2 baseline and supports re-engaging the angular planning cost without regression but does not improve over A2 in our setup. We document these explorations to show that the canonical SD-JEPA architecture (A2) does not require any of these extra regularisers, and that adding them tends to interfere with rather than complement the disjoint-supports decomposition.

Additional environments can be further explored such as WorldGym [27], richer learning approaches directly in the latent [28], or totally different modalities such as in neuroscience [29] to observe if we can link such progression metrics to isolate distinct phenomena.

13 Extended latent geometry across environments↩︎

This appendix collects the per-environment latent-geometry figures complementing §5.5. For each environment we ran the latent-compass analysis on the trained checkpoint indicated, on a small set of held-out episodes spanning the dataset’s full episode-id range. Per checkpoint we report (i) the t-SNE projections of \(\boldsymbol{z}^{\mathrm{prog}}\) and \(\boldsymbol{z}^{\mathrm{cont}}\), (ii) the trajectory-overlay heatmap of \(\theta_t\) in raw state space, (iii) per-episode signed-Spearman heatmaps against candidate task-progress proxies, and (iv) one frame strip per env exhibiting an interesting non-monotone or multi-phase \(\theta\) behaviour.

13.1 Push-T at \(k_{\mathrm{prog}}{=}2\) (canonical A2)↩︎

The body’s t-SNE pair (Fig. 4) is from this checkpoint. The Push-T \(r\)-coordinate stays in roughly \([0.90, 1.57]\) with per-episode std up to \(0.49\), \(\theta\)-span around \(3.3\) rad. Per-episode signed Spearman against five candidate proxies (Fig. 5) shows that \(|\rho|\) with step_idx is consistently very high (\(\approx 0.86\)–\(1.00\)) on all but one episode (ep 8000, the regression case), with block_angle_err and block_goal_dist as close runners-up.

13.2 Push-T at \(k_{\mathrm{prog}}{=}8\)↩︎

The \(k=8\) checkpoint shows a substantially cleaner \(S^1\) progression manifold: \(r \in [0.62, 0.77]\) with per-episode std \(\in [0.04, 0.07]\), angular span \(6.1\) rad (effectively full-circle coverage). The clock advantage seen at \(k=2\) also narrows: mean \(|\rho|\) vs.step_idx drops from \(0.91\) to \(0.56\), with block_angle_err and block_goal_dist reaching \(|\rho|\approx 0.56\) as well. At \(k=8\) \(\theta\) is less time-like and more cyclic-phase-like.

13.3 Reacher at \(k_{\mathrm{prog}}{=}8\)↩︎

Mean \(|\rho|\) on Reacher is \(0.53\) for the clock and \(0.49\) for ee_target_dist; on individual episodes the latter beats the clock. joint_speed is essentially uncorrelated with \(\theta\) (\(|\rho|\approx 0.08\)): \(\theta\) is not encoding “how fast is the arm moving.” The \(+3.3\) pp full-z planning lift on Reacher (Tab. 3) is consistent with the angular-cost-on-\(\boldsymbol{z}^{\mathrm{prog}}\) adding usable information when the latent’s \(\theta\) is well-coupled to the end-effector-to-goal geometry.

13.4 Two-Room at \(k_{\mathrm{prog}}{=}8\), the cyclic-phase example↩︎

The mean Two-Room \(|\rho|\) is \(0.50\) for the clock and \(0.34\) for agent_target_dist, but the per-episode picture is bimodal: ep 1000 is clock-like (\(\rho=+0.88\)), ep 7000 is fully task-coupled (\(\rho=-0.90\)), ep 3000 (Fig. 8) is the cyclic case where Spearman fails. Mutual information or circular correlation would tighten these estimates; we leave that to follow-up work.

13.5 OGBench-Cube at \(k_{\mathrm{prog}}{=}8\)↩︎

OGBench-Cube is the only environment in our suite where a task-physical signal beats the elapsed-time clock in correlation with \(\theta_t\) (mean \(|\rho|\) with block_target_dist is \(0.59\) vs \(0.47\) with step_idx). It is also the environment with the most varied per-episode behaviour: spatial-distance encoding for reaches, gripper-state encoding for picks, and U-cycle encoding for U-shaped paths. The frame strip on episode 500 (Fig. 9) and the 3D t-SNE projections (Fig. 10) capture this variability.

14 Supplementary perturbation analyses↩︎

14.0.0.1 Brief background on head-direction cells.

For readers unfamiliar with the neuroscience analogy invoked in §7: head-direction (HD) cells are neurons in the rodent hippocampal-entorhinal system that encode the animal’s facing direction in the world, independent of location [25], [26]. Each cell fires preferentially when the head is oriented in a specific allocentric direction; together they form a 1-D circular code where a population peak travels around a circle as the animal turns. Two properties make the parallel with \(\theta_t\) natural: (i) a scalar circular coordinate (a 360\(^\circ\) heading is collapsed to one moving peak), and (ii) re-mapping under perturbation, when the environment is teleported, rotated, or otherwise restructured, HD-cell populations re-anchor their peak to the new visual cues and then track coherently from there, the same kind of structural re-localisation we observe in \(\theta_t\) after our teleport-and-continue perturbation. We do not claim our latent rediscovered the biological circuit; the parallel is qualitative, two systems converging on a similar geometric solution (a single circular phase coordinate \(+\) re-mapping under sensory restructuring) for what is broadly the same problem.

The main text headlines the teleport-and-continue perturbation (Figure 3), in which the “moment” (latent surprise spike) and the “meaning” (post-perturbation relocalisation in \(\theta\)) come apart cleanly. Three complementary perturbation modes confirm and refine that picture. All run on Push-T A2 (canonical, \(k_{\text{prog}} = 2\), seed \(3072\), epoch 10), processed by the same encoder with no test-time fine-tuning.

14.0.0.2 Splice (Figure 11).

We replace all frames of episode \(0\) from step \(50\) onward with frames from episode \(1000\). Unlike the headline teleport-and-continue — which uses an episode-pair chosen to maximise the visual contrast in solving mode, the splice picks an arbitrary substitute scene. Result: a \(\approx 0.17\) rad surprise spike at the boundary (\(\sim 8\times\) baseline mean), followed by sustained elevated noise across steps \(50\)–\(70\) as the encoder processes the unfamiliar trajectory. The post-perturbation \(\theta_t\) diverges from the clean baseline but does not pin to a single new sector; instead it traces a new path consistent with the new episode’s progression. This is a weaker but still positive surprise signal, the magnitude depends on how angularly distinguishable the donor episode is from the original.

14.0.0.3 Single-frame teleport (Figure 12).

We insert a single frame from a structurally different scene at step \(50\) of episode \(0\), then resume episode \(0\). The angular surprise spikes to \(\approx 0.20\) rad at the substituted frame and returns to baseline within a single step, a textbook illustration of localised surprise with immediate recovery. The model registers the anomaly, reports a sharp \(\theta\) jump, and resumes normal localisation as soon as the genuine observation stream returns. This mode tests recovery dynamics: when the perturbation is transient, relocalisation is also transient.

14.0.0.4 Reverse-window perturbation (Figure 13).

We reverse a 10-frame window of episode \(0\) starting at step \(40\). Critically, the set of frames presented to the encoder is unchanged — only the temporal order is altered. Any angular discontinuity in this mode must therefore be a model-dynamics phenomenon, not a scene-novelty artifact. The result: matched spikes of \(\approx 0.11\) rad at the window’s leading boundary (step \(40\)) and \(\approx 0.13\) rad at the trailing boundary (step \(50\)), with elevated noise within the reversed window. This rules out the alternative hypothesis that the surprise spikes in the splice and teleport modes are simply registering “a novel scene appeared”; they are registering “the dynamics violated the model’s local expectations”, which is the stronger and more interesting claim.

14.0.0.5 Animated GIF for talks / online supplementary.

The same teleport-and-continue experiment renders as an animated GIF via analysis/surprise_voe.py --gif --mode teleport_continue. The animation marks the perturbation frame with a red border + “TELEPORT” overlay text and switches the on-screen \(\theta_t\) trace from blue (pre-perturbation) to red (post-perturbation), making the relocalisation visible in real time. The static Figure 3 carries the same information for the paper.

15 Phase-event localisation vs.classical surprise: extended results↩︎

This appendix collects the full empirical evidence for the comparison between \(|\Delta\theta_t|\) and z-MSE as surprise signals on OGBench-Cube (§5.5, Tab. 4, Fig. 2). The headline finding is that the two metrics measure complementary things: \(|\Delta\theta_t|\) is the right tool for semantic-event localisation, z-MSE is the right tool for raw magnitude-anomaly detection, and combining them does not improve detection on either test.

15.1 Phase-alignment summary across all 40 episodes↩︎

Per-episode AUROC distribution at \(\pm 2\)-step tolerance:

Both AUROC distributions sit below \(0.5\) at the tight \(\pm 1\)-step tolerance: z-MSE peaks during the stable-contact phase (when the cube is being translated while gripped) rather than at the transitions themselves, and only \({\sim}11\%\) of all steps are positive labels. \(|\Delta\theta|\)’s peaks are consistently closer to transitions even when not exactly aligned, so widening the tolerance window catches more of those near-aligned peaks and pushes \(|\Delta\theta|\) above \(0.5\) before z-MSE follows. The right reading is the gap, not the absolute level.

15.2 Action-corruption AUROC: the magnitude-anomaly negative control↩︎

To verify that the two metrics genuinely measure different things, we run a complementary test where the corruption is a magnitude anomaly rather than a semantic event. We replace a contiguous \(25\%\) band of each episode’s actions with a random shuffle and ask each metric to localise the corrupted band as a per-step random-segment AUROC.

z-MSE wins this test by construction: the action-shuffle injects a magnitude anomaly that a \(192\)-dimensional MSE detects more sensitively than a \(1\)-dimensional phase scalar; the per-step random-segment label rewards information bandwidth, not semantic alignment. The combined rank-sum score does not beat z-MSE on any cell; the two metrics are not strictly orthogonal but are clearly answering different questions.

15.3 Regime-change segmentation (change-point detection)↩︎

A second formulation of the same hypothesis: at semantic phase events, \(\theta_t\) exhibits a regime shift (different mode of evolution before vs.after) while z-MSE just spikes and returns to baseline. We run BinSeg with \(K = 4\) requested change points on the \(40\) cube episodes and compute precision/recall/F1 against the ground-truth contact transitions with greedy 1-to-1 matching at tolerance \(\pm\)tol.

The regime-change test is a partial vindication: at loose tolerance the observed angular-velocity trace beats both prediction-error metrics, but at tight tolerance z-MSE’s spike-and-return pattern produces sharper change-point boundaries that BinSeg locates more accurately. This is consistent with the AUROC reading: \(|\Delta\theta|\) excels at peak alignment with semantic events, z-MSE excels at sharp magnitude-spike segmentation; both are useful and the choice depends on what the practitioner wants to detect.

15.4 Linear-probe information density: full per-feature table↩︎

We fit Ridge probes from each latent feature to a per-episode-normalised cube-to-target progress target. The pooled regime leaves \(25\%\) of episodes out and fits on the rest (\(30\) bootstraps); the per-episode regime fits within each of \(40\) held-out episodes with leave-one-step-out CV. The body of the paper (Fig. 16) shows the headline pooled-vs-per-episode contrast; the full per-feature table is below.

The interpretation in the body covers the headline. The pooled probe makes one further point clearly: the compass is a per-trajectory phase coordinate, not a globally calibrated absolute-distance estimate. The two regimes ask different questions, and the answers are consistent with the design intent (\(\boldsymbol{z}^{\mathrm{prog}}\) encodes within-episode progression, while episode-specific scene content lives in \(\boldsymbol{z}^{\mathrm{cont}}\)).

16 Cross-machine reproducibility note↩︎

Trainings ran on a mix of H100 GPUs and (for Push-T) RTX 5090 GPUs; all evaluations are on H100. Single-seed checkpoints exhibit non-trivial planning-success variance across these microarchitectures even when training data, config, and seed are identical. Concrete observation: a SD-JEPA \(k_{\mathrm{prog}}{=}4\) checkpoint trained on an H100 at seed \(0\) evaluates to \(100\%\) on Push-T; the analogous re-train on an RTX 5090 at the same seed evaluates to \(94\%\), a 6-point gap on identical configuration. We attribute this to non-bit-deterministic bf16 matmuls across tensor-core kernels, dataloader shuffling sensitivity to num_workers, and library version pinning differences across the two machines. Multi-seed averaging absorbs the variance: the 3-seed Push-T means at \(k_{\mathrm{prog}}{=}4\) and \(k_{\mathrm{prog}}{=}8\) (\(96.0\) and \(97.3\)) sit cleanly within LeWM’s own reported single-environment \(\sigma = 2.83\). We therefore treat \(\geq n=3\) multi-seed means as the reportable comparison granularity throughout the paper.

References↩︎