Technical Summary: Sensorimotor World Models (SMWM)

Problem Statement

The paper addresses the challenge of learning latent world models from high-dimensional observations (e.g., pixels) without reward signals or task labels. While Joint Embedding Predictive Architectures (JEPAs) offer a promising framework by predicting future states directly in an embedding space (avoiding pixel reconstruction), they suffer from representation collapse. When an encoder and a forward dynamics model are trained jointly to minimize prediction error, the encoder may learn to map all observations to a constant embedding. This trivial solution minimizes the forward loss but renders the model useless for downstream control or planning.

Existing solutions to prevent collapse often rely on complex mechanisms such as:

• Freezing pretrained encoders (e.g., DINO-WM).
• Adding distributional regularizers (e.g., matching embeddings to an isotropic Gaussian via SIGReg).
• Using Exponential Moving Average (EMA) target networks with stop-gradients.
• Employing multi-term loss functions involving variance-covariance constraints.

The authors posit that these approaches may not fully address the core question of what a useful world model representation should retain. Drawing on the concept of "perception for action" and enactive perception, the paper argues that representations should be shaped by their relevance to actions, preserving distinctions that matter for controllable dynamics while discarding uncontrollable distractors.

Methodology: Sensorimotor World Model (SMWM)

The authors propose SMWM, a latent world model trained end-to-end from offline, reward-free trajectories $(o_t, a_t, o_{t+1})$. The core innovation is the use of inverse dynamics regularization as the sole mechanism to prevent representation collapse.

Architecture

The model consists of three components trained jointly:

1. Encoder ($f_\theta$): Maps observations $o_t$ to a compact latent embedding $z_t$.
2. Forward Dynamics Model ($g_\phi$): Predicts the next latent state $\hat{z}_{t+1}$ given the current state $z_t$ and action $a_t$.
3. Inverse Dynamics Model ($h_\psi$): Predicts the action $\hat{a}_t$ taken between two consecutive embeddings $z_t$ and $z_{t+1}$.

Training Objective

The model is optimized by minimizing a joint loss function:
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• Forward Loss ($\mathcal{L}_{fwd}$): Mean Squared Error (MSE) between the predicted next embedding $\hat{z}_{t+1} = g_\phi(z_t, a_t)$ and the actual encoded next state $z_{t+1} = f_\theta(o_{t+1})$.
• Inverse Loss ($\mathcal{L}_{inv}$): MSE between the predicted action $\hat{a}_t = h_\psi(z_t, z_{t+1})$ and the ground-truth action $a_t$.
• Regularization Weight ($\lambda$): A single hyperparameter controlling the strength of the inverse dynamics signal.

Mechanism of Action

The inverse loss acts as an anti-collapse mechanism. If the encoder collapses to a constant embedding, the inverse model receives identical inputs for all transitions and can only predict a constant action (the mean of the action distribution). To reduce the inverse loss below this constant-predictor risk, the encoder must preserve information about the action $a_t$ contained in the transition $(o_t, o_{t+1})$. This forces the representation to retain action-relevant information (controllable degrees of freedom) while naturally filtering out uncontrollable distractors (exogenous noise) that do not influence the action recovery.

Unlike distributional regularizers that prescribe the geometry of the embedding space (e.g., forcing it to be Gaussian), inverse dynamics anchors the representation to a task-grounded quantity (the action) without imposing specific geometric priors.

Key Contributions

1. SMWM Architecture: A simple, end-to-end trainable sensorimotor world model that uses inverse dynamics as the sole anti-collapse mechanism, requiring only one additional hyperparameter ($\lambda$).
2. Representation Properties: Empirical evidence that the learned representations are not merely non-collapsed but specifically track controllable degrees of freedom, preserve spatial geometry (both Cartesian and periodic), and filter out uncontrollable distractors.
3. Downstream Performance: Demonstration of competitive planning performance across 2D and 3D control tasks, showing that the structural advantages of the representation translate to effective goal-conditioned planning.

Experimental Results

Controlled Testbeds (Dot World & Sprite World)

• Intrinsic Dimension Recovery: In a "dot world" environment, the model successfully recovered the true 2D intrinsic dimension of the state space from pixels alone, even when embedded in a 64-dimensional latent space.
• Distractor Filtering: When uncontrollable "distractor" objects were introduced (moving randomly), the encoder ignored them. The effective latent dimension matched the controllable degrees of freedom, not the total observation dimension.
• Control-Dependent Perception: In a sprite world experiment, the model learned to preserve pose variables (position, orientation) only when they were controllable. When orientation was uncontrolled, the representation averaged it out, resulting in a blurred reconstruction of orientation.

Planning Benchmarks

The authors evaluated SMWM on four offline control tasks: TwoRoom (2D navigation), Reacher (continuous control), Push-T (2D manipulation), and OGBench-Cube (3D manipulation).

• Baselines: Compared against SIGReg (a strong JEPA baseline using Gaussian regularizers), Forward-only (no inverse loss), and Random action sampling.
• Performance:
  • SMWM matched SIGReg performance on the three 2D tasks.
  • SMWM outperformed SIGReg on the 3D OGBench-Cube task (84% success vs. 59%).
  • Both regularized methods significantly outperformed the Forward-only ablation, confirming the necessity of an anti-collapse mechanism.
• Latent Geometry: Probing experiments revealed that SMWM embeddings form low-dimensional subspaces mirroring the physical state geometry (e.g., cylindrical structures for periodic joint angles in Reacher, Cartesian maps for navigation).

Significance and Claims

The paper claims that SMWM provides a principled and simple instantiation of "perception for action" in latent world model learning. By forcing latent states to preserve information sufficient to recover the agent's actions, the model naturally biases the representation toward the controllable degrees of freedom of the environment.

Key takeaways include:

• Simplicity: Inverse dynamics alone is sufficient to stabilize end-to-end JEPA training without reconstruction, reward signals, frozen encoders, or complex distributional priors.
• Structural Bias: The method induces representations that are interpretable and aligned with the causal structure of the environment (actions and their effects), rather than just predictive correlations.
• Practical Utility: The learned representations support effective latent planning, particularly in complex 3D environments where other regularized baselines struggle.

The authors acknowledge limitations, such as the assumption that actions are recoverable from consecutive observations and the current inability to capture quantities not identifiable from single frames (e.g., velocity) without observation histories. However, they position SMWM as a robust, reward-free approach to learning actionable world models from offline data.