PaperlayerView exact revision on arXiv

Source-linked AI summary

Sensorimotor World Models: Perception for Action via Inverse Dynamics

Petr Ivashkov, Randall Balestriero, Bernhard Schölkopf

arXiv:2606.20104v1cs.LGcs.AI

TL;DR

Latent world models predict future states from high-dimensional observations, while a useful representation should retain distinctions relevant to action. SMWM trains a predictive latent model with inverse dynamics, and its learned representations track controllable structure, preserve geometry, and filter distractors. The paper identifies action recoverability and single-frame state observability as boundaries of the approach.

  • Problem

    Latent world models predict future states in embedding space, but the representation must retain distinctions relevant to actions and their consequences.

    Source: §1 · pp. 1–2

  • Method

    SMWM jointly trains an encoder and predictive latent model with an inverse-dynamics head that predicts actions from consecutive embeddings.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Inverse-dynamics training yields representations that track controllable degrees of freedom, preserve Cartesian and periodic geometry, and filter uncontrollable distractors.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The paper reports that this structural advantage carries over to competitive planning across 2D and 3D control tasks.

    Source: §1 · p. 3

    The approach assumes actions are recoverable from consecutive observations, which may fail when distinct actions produce the same visible change.

    Source: §6 · p. 11

Abstract

from arXiv · show

Perception for action suggests that representations of the world should be shaped not by visual fidelity alone, but by their relevance for actions. At the same time, latent JEPA-style world models advocate learning compact predictive states from high-dimensional observations to facilitate the prediction of future states, but end-to-end training of these models is nontrivial because representations may collapse if our only goal is to construct a latent state that is easy to predict. We introduce a sensorimotor world model (SMWM): a latent world model trained end-to-end with inverse dynamics regularization. This single regularizer addresses both issues: it prevents representation collapse and induces action-aligned representations. By forcing latent states to preserve information about the action underlying a transition, it biases the model toward the controllable degrees of freedom of the environment while discarding uncontrollable distractors. This yields stable latent world models trained from offline, reward-free trajectories, without frozen encoders, exponential moving averages, or complex latent regularizers. Empirically, SMWM learns compact, interpretable latent spaces and enables competitive planning performance across simple 2D and 3D control tasks.

1 Introduction

Latent world models predict in embedding space, but jointly learning predictive states can produce collapsed representations and does not by itself specify what information those states should retain. SMWM uses inverse dynamics to shape states around action-relevant information while addressing collapse.

1 Introduction

  • World models predict future states given actions, while latent models avoid pixel-space prediction and JEPA-style models predict directly in embedding space.Pixel-space learning is difficult with high-dimensional observations and can emphasize background details irrelevant to control.Source: §1 · p. 1
  • JEPA-style models predict in embedding space, but joint encoder–dynamics training can collapse every observation to one embedding, making prediction trivial and the model useless downstream.Prior approaches include frozen encoders, multi-term embedding regularizers, Gaussian distribution matching, and EMA target encoders.Source: §1 · p. 1
  • Perception-for-action motivates representations that retain distinctions relevant to available actions and their consequences, rather than merely preserving information about observations.The introduction connects this view to causal representations and common coding.Source: §1 · p. 2
  • SMWM is a predictive latent model whose state is shaped by motor commands and sensory change, trained by propagating inverse-dynamics gradients into the encoder.The inverse head predicts the action between consecutive observations from their embeddings, encouraging action-relevant features while discarding the rest.Source: §1 · p. 2
  • SMWM uses inverse dynamics as its sole anti-collapse mechanism and is reported to learn action-aligned representations and achieve competitive planning against SIGReg baselines across 2D and 3D tasks.The reported representations track controllable degrees of freedom, preserve spatial geometry, and filter uncontrollable distractors.Source: §1 · p. 3

2 Background and related work

The paper situates SMWM within latent predictive modeling, collapse prevention, and inverse-dynamics representation learning. It uses inverse dynamics as a standalone anti-collapse mechanism in offline pixel-action JEPA training.

Latent world models

  • Dreamer and PlaNet learn latent dynamics with pixel decoders; TD-MPC trains an encoder and latent dynamics end-to-end, while JEPAs predict in embedding space without reconstruction or reward.The paper places its work in the JEPA family and offline, reward-free regime.Source: §2 · p. 3

Representation collapse

  • Collapse remedies include contrastive negatives, stop-gradient architectures, and embedding-distribution regularizers; related JEPA models freeze encoders, add auxiliary losses, match Gaussian embeddings, or use EMA targets.SMWM shares the JEPA framework and offline pixel-action setting but uses inverse dynamics as its standalone anti-collapse mechanism.Source: §2 · p. 3

Inverse dynamics

  • Prior inverse-dynamics work targets controllable features or, under stated assumptions, recovers control-relevant state; those theoretical results do not transfer directly to SMWM’s continuous-state, continuous-action single-step setting.The paper treats these works as motivation for action prediction as an inductive bias.Source: §2 · p. 3

Other directions

  • Action-free world models infer pseudo-actions from video, visual RL shares encoders without explicit world models, and bisimulation and causal abstraction address related state-representation questions.The paper distinguishes these adjacent directions as addressing different problems.Source: §2 · p. 3

3 Method

SMWM learns an encoder and forward dynamics model from offline observation-action transitions, then adds an inverse-dynamics objective so embeddings are both predictive and informative about actions. Joint optimization uses both losses to train the encoder.

Dataset and setting

  • The method assumes offline transitions of consecutive observations and continuous actions, without requiring a behavior policy or reward function.The encoder maps observations to compact embeddings, and the forward model predicts the next embedding from the current embedding and action.Source: §3 · p. 3

Forward dynamics in embedding space

  • Forward prediction alone admits a zero-loss collapsed solution in which every observation and predicted embedding is constant.The forward objective matches the predicted next embedding to the encoded next observation using mean-squared loss.Source: §3 · p. 4

Inverse dynamics regularization

  • The inverse head predicts the action from consecutive embeddings; under squared loss, collapse reduces it to a constant-action predictor, so better prediction requires action-predictive information in the embeddings.Unlike an isotropic-Gaussian prior, inverse dynamics does not prescribe embedding geometry.Source: §3 · p. 4

Joint training objective

  • The joint objective adds the inverse loss to the forward loss, with λ controlling its strength, and sends gradients from both objectives into the encoder.The forward model and inverse model receive gradients from their respective losses alongside the shared encoder.Source: §3 · p. 4

4 Learned latent structure

On a testbed with known state and action geometry, the learned representations are low-dimensional and spatially faithful, with structure that follows controllability and action-induced dynamics.

Recovered intrinsic dimension

  • In the 2D dot world, two principal components capture essentially all variance in a 64-dimensional embedding while preserving spatial neighborhoods without state supervision.The spectrum drops sharply after the true intrinsic dimension of 2; the latent grid is near-square, with mild deformation near the boundary.Source: §4 · p. 4, Figure 2 · p. 5

Forward model

  • In a five-step rollout, latent predictions track encoded ground-truth embeddings, and actions produce approximately state-independent displacements in latent space.The forward model behaves approximately as ga(z) ≈ z + ρ(a), although its training loss does not enforce that structure.Source: §4 · p. 5, Figure 3 · p. 6

Equivariance and latent action composition

  • Exact equivariance implies that action composition is represented by composition of latent interventions; experiments suggest the learned maps are approximately additive translations.The homomorphism result holds on the encoded data manifold under exact equivariance, while the learned-model identity is approximate.Source: §4 · pp. 5–6

Controllable degrees of freedom

  • Across four configurations, significant principal-component counts match controllable dimensions (4, 2, 2, and 6), while uncontrollable distractor directions effectively collapse.The random dot is ignored because its state is neither needed to recover actions nor predictable from past latent state.Source: §4 · pp. 6–7, Figure 4 · p. 7

Visualizing control-dependent perception

  • With a triangular sprite, the representation preserves controlled pose variables and averages out uncontrolled ones, producing an orientation-averaged reconstruction when rotation is uncontrolled.When x/y/θ are all controlled, the decoded reconstruction is sharply oriented.Source: §4 · p. 7

5 Planning in latent space

The section evaluates latent-space planning and the learned representations that support it. SMWM matches SIGReg on three 2D tasks and outperforms it on the 3D task, while its embeddings organize controllable state compactly.

Setup

  • Planning encodes current and goal observations, scores candidate action sequences through latent rollouts, and uses CEM within a receding-horizon MPC loop.Evaluation uses a 50-step budget and goals 25 steps ahead.Source: §5 · pp. 7–8, Figure 5 · p. 8

Architecture

  • SMWM uses LeWM’s encoder and predictor with a single-step inverse model that predicts actions from consecutive embeddings.The added inverse model is a two-layer MLP mapping the concatenated embeddings to an action prediction.Source: §5 · p. 8

Planning performance

  • SMWM matches SIGReg on TwoRoom, Reacher, and Push-T, and reaches 84% success versus SIGReg’s 59% on OGBench-Cube.Both regularized methods substantially outperform Forward-only and Random across all four environments.Source: §5 · pp. 8–9, Figure 5 · p. 8

Probing physical state

  • Both regularized methods recover physical state near-perfectly with MLP probes, while SMWM and SIGReg have different linear-probe advantages across quantities.SMWM leads on Push-T agent position and all three Cube quantities; SIGReg leads on Reacher joint angles and Push-T block pose.Source: §5 · p. 9, Table 1 · p. 9

Latent geometry

  • SMWM embeddings concentrate in a low-dimensional subspace, with Cartesian coordinates forming approximately linear directions and Reacher’s periodic angles forming a torus.The geometry is visible in the PCA spectra and projections; the Reacher torus appears as a cylindrical structure in 3D projections.Source: §5 · p. 9, Figure 6 · p. 10

6 Discussion

The discussion presents inverse dynamics regularization as addressing collapse and action alignment together. It also identifies assumptions and evaluation boundaries that constrain the present results.

Discussion

  • A single-step inverse-dynamics head stabilizes end-to-end JEPA-style training while biasing representations toward compact, action-sufficient structure.The learned representations track controllable degrees of freedom, preserve Cartesian and periodic geometry, and filter uncontrollable distractors.Source: §6 · p. 10
  • SMWM matches or outperforms SIGReg across four control tasks, with its largest margin on the 3D manipulation task.The discussion characterizes this as evidence for a perception-for-action bias in latent world-model learning.Source: §6 · p. 10

Limitations and future work

  • The approach assumes actions are recoverable from consecutive observations, and evaluation is limited to moderate-scale simulated control tasks.Single-frame encoders also cannot capture quantities such as velocity that are not identifiable from one frame; offline dataset coverage and compounding long-horizon model error constrain planning.Source: §6 · p. 11

A.1 Training objective

SMWM trains its encoder with forward prediction and inverse dynamics losses. The appendix describes its architecture, optimization, evaluation protocols, and compute.

A.1 Training objective

  • The training objective combines forward prediction loss with a weighted inverse-dynamics loss, and gradients from both reach the encoder.The inverse head predicts the action between consecutive observations from their embeddings.Source: §A.1 · p. 15

A.2 Architecture

  • The architecture uses LeWorldModel’s encoder and forward model, adding a two-layer MLP inverse head that predicts actions from consecutive embeddings.The encoder projects a ViT-Tiny CLS token to a 192-dimensional latent; the forward model predicts from one latent-action pair.Source: §A.2 · p. 15

A.3 Optimization details

  • The inverse-dynamics weight λ is environment-specific and requires tuning; the SIGReg comparison uses the same optimizer and architecture with its own regularizer.The appendix reports an example where λ = 0.1 is optimal for TwoRoom but yields near-random planning performance on Reacher.Source: §A.3 · pp. 15–16, Figure 7 · p. 16, Table 2 · p. 16

A.4 Planning protocol

  • Planning encodes current and goal observations, optimizes candidate action sequences with CEM over latent rollouts, and replans in an MPC loop.Evaluation samples held-out starts with goals 25 steps later and measures success within a 50-step budget.Source: §A.4 · p. 16

A.5 Probing protocol.

  • State probes are trained on frozen embeddings and evaluated using held-out R2, with linear ridge regression and a two-layer MLP probe.The protocol uses 25,000 training embeddings per environment and method and a disjoint 5,000-sample validation split.Source: §A.5 · pp. 16–17

A.6 Compute.

  • Each world model was trained in under four hours using a single NVIDIA H100 GPU.Source: §A.6 · p. 17

B Environments and datasets

The appendix describes the offline datasets and simulated environments, tests planning at longer goal offsets, compares SIGReg’s latent geometry, and details a controlled dot-world testbed.

B Environments and datasets

  • The experiments use four offline environments: TwoRoom navigation, Push-T manipulation, OGBench-Cube manipulation, and Reacher control.The datasets include scripted, expert-demonstration, and SAC-generated trajectories, with held-out validation episodes used for reported metrics.Source: B Environments and datasets · p. 17

C Planning performance on extended goal offsets

  • At longer goal offsets, SMWM remains stable on TwoRoom and OGBench-Cube, while both regularized methods decline on Push-T and remain comparable on Reacher.SIGReg degrades sharply on TwoRoom or trails consistently on Cube; the evaluation budget scales to twice the goal offset.Source: C Planning performance on extended… · pp. 17–18, Figure 9 · p. 18

D SIGReg latent geometry

  • SIGReg’s embeddings lack a clear PCA elbow and spread state-dependent structure across many components, making the representation less compact and interpretable than SMWM’s.The appendix attributes this pattern to SIGReg’s isotropic-Gaussian matching regularizer.Source: D SIGReg latent geometry · p. 18, Figure 10 · p. 19

E.1 Environment

  • The dot-world testbed varies controllable motion types and action dimensions, including random dots that move without contributing to the action vector.Its configurations include independent dots, coupled pairs, distractor dots, and a combined setting.Source: §E.1 · pp. 18–19

E.2 Architecture and training

The appendix specifies the model’s network and training setup, tests collapse without inverse loss, and introduces a triangular-agent experiment that varies which pose changes are exposed as actions.

E.2 Architecture and training

  • The encoder uses three stride-2 convolutions, max-pooling, and a 64-dimensional tanh embedding; the forward and inverse models are two-layer MLPs.All networks train jointly for 100 epochs on 250,000 transitions, with λ = 10 and a held-out split of 1,000 transitions.Source: §E.2 · p. 19

E.3 Collapse without the inverse loss

  • With λ = 0, the encoder maps dot-world probe observations to essentially the same embedding, so PCA collapses to a point at the origin.The forward loss alone can be minimized by a constant, trivially predictable representation.Source: §E.3 · p. 20, Figure 11 · p. 20

F.1 Environment and control settings

  • The triangular agent has pose (x, y, θ), while the experiment varies whether none, translation only, or translation and rotation are exposed as actions.Uncontrolled degrees of freedom continue to move randomly, so the visual world remains fixed while the action repertoire changes.Source: §F.1 · p. 20

F.2 Training and reconstruction probe

The triangular-agent probe trains the same inverse-dynamics-regularized world model under different action interfaces, then uses a decoder to visualize information in frozen embeddings. Reconstructions preserve the pose components that are controllable and average over uncontrolled orientation.

F.2 Training and reconstruction probe

  • For each control setting, the same encoder, forward model, and inverse model are trained; with no action dimensions, training reduces to the forward-only objective.The post-hoc decoder is not used to train the representation.Source: §F.2 · pp. 20–21

F.3 Interpretation

  • Reconstructions collapse without actions, preserve position but average orientation under x/y control, and preserve full pose under x/y/θ control.The decoder is trained after the encoder is frozen and serves only to visualize information already present in the embeddings.Source: Figure 12 · p. 21, §F.3 · p. 21, §F.2 · p. 21
  • The reconstruction probe provides a visual counterpart to the PCA results: learned representations retain controllable degrees of freedom and average out uncontrolled ones.Source: §F.3 · p. 21

G Equivariance and latent action composition

The appendix relates approximate equivariance to composition of latent interventions and argues that inverse dynamics makes actions identifiable from latent transitions. It situates this design among prior world models and clarifies that existing finite-state inverse-dynamics guarantees do not transfer to this setting.

G Equivariance and latent action composition

  • Approximate equivariance makes composed physical actions correspond to composed latent interventions on the encoded data manifold, under a nontrivial semigroup-closure assumption.The homomorphism result applies on-support; action sequences may not be equivalent to a single primitive action.Source: G Equivariance and latent action… · p. 22
  • Equivariance permits both uncompressed identity encodings and total collapse, so useful latents also need faithful, identifiable action effects.The inverse model imposes identifiability by recovering the action from consecutive latent states.Source: G Equivariance and latent action… · p. 22

H Extended related work

G Equivariance and latent action composition

  • In simple displacement domains, experiments suggest latent interventions approximately act as additive translations, but additivity is empirical rather than imposed.The translation-like pattern need not hold for contact-rich dynamics, boundary effects, periodic variables, strongly state-dependent effects, or noncommutative action laws.Source: G Equivariance and latent action… · pp. 22–23, §H.1.1 · p. 23

H.1.1 Intrinsic Curiosity Module (ICM)

  • Unlike ICM, SMWM sends forward-loss gradients into the encoder and uses its encoder as the downstream world-model representation, not only to compute intrinsic reward.The paper frames inverse dynamics as the anti-collapse mechanism for end-to-end JEPA training without reconstruction, EMA, or distributional regularizers.Source: §H.1.1 · pp. 23–24

H.1.2 Multi-step inverse identifiability theory

  • Finite-state multi-step inverse-dynamics guarantees do not transfer to SMWM’s continuous state and action spaces, single-step inverse objective, and lack of an output-cardinality bottleneck.The paper presents that literature as conceptual motivation rather than a guarantee for its objective.Source: §H.1.2 · p. 24

H.2 Inverse-dynamics-regularized world models

  • Related world models differ in supervision and anti-collapse design: WAM uses reconstruction and rewards, while EB-JEPA combines four regularizers; SMWM uses inverse dynamics alone.The paper also distinguishes DynaMo’s jointly inferred latent actions from its own executed-action supervision.Source: §H.2 · pp. 24–25

H.3.1 Reconstruction-based

  • Reconstruction-based Dreamer and PlaNet use pixel-reconstruction and reward-prediction heads under a variational objective; reconstruction can encode pixel details irrelevant to control.Source: §H.3.1 · p. 25

H.3.3 Augmentation and contrastive methods

The related-work discussion contrasts model-free representation-learning mechanisms with action-conditioned JEPA approaches. SMWM adds executed-action prediction as a signal anchoring embeddings to controllable structure.

H.3.3 Augmentation and contrastive methods

  • CURL, DrQ, and Proto-RL learn encoders for model-free policies rather than world models, and their anti-collapse mechanisms are described as orthogonal to inverse-dynamics regularization.Source: §H.3.3 · p. 25

H.3.4 JEPA-style regularizers

  • Among JEPA-style approaches, SMWM uses task-grounded action prediction to anchor embeddings to controllable structure, unlike frozen encoders, multi-term regularizers, SIGReg, or EMA targets.Source: §H.3.4 · pp. 25–26
Loading 2606.20104v1…