Sensorimotor World Models: Perception for Action via Inverse Dynamics
Introduces inverse dynamics regularization in sensorimotor world models (SMWM), preventing representation collapse and improving control in offline, reward-free settings.
Key Findings
Methodology
The approach combines an end-to-end trained latent world model with an inverse dynamics head, enforcing the preservation of action information in the latent space. The encoder fθ maps high-dimensional observations to a compact latent space, while gϕ predicts future states and hψ predicts actions between states. The training minimizes forward prediction loss and inverse prediction loss simultaneously, ensuring the latent space encodes controllable dynamics. Experiments on 2D and 3D control tasks demonstrate the model learns low-dimensional, interpretable, and control-relevant representations without complex regularizers or frozen encoders.
Key Results
- In a 2D dot environment, PCA analysis shows the learned latent space captures the true 2D structure, with the top two components explaining nearly all variance. The latent grid preserves spatial relationships, indicating faithful topology recovery. The model achieves over 80% success in control tasks, outperforming baseline methods by 20%.
- In 3D environments like OGBench-Cube, the model attains success rates above 85%, significantly better than models without inverse dynamics regularization. Ablation studies confirm that removing the inverse regularizer causes representation collapse and reduces planning success by over 15%.
- The learned representations effectively filter out uncontrollable distractors, focusing on controllable degrees of freedom, as shown in experiments with multiple dots and distractors. The model's ability to recover intrinsic dimensionality matches the controllable factors, demonstrating robustness across environments.
Significance
This work advances the development of stable, controllable latent world models trained solely from offline data, addressing the longstanding challenge of representation collapse. Its simplicity and effectiveness make it a promising foundation for autonomous agents operating in complex, real-world scenarios without reliance on reward signals or complex regularizers. The approach bridges cognitive theories of perception for action with modern deep learning, offering a scalable and interpretable framework for future research and applications in robotics, simulation, and control systems.
Technical Contribution
The key innovation is the integration of inverse dynamics as a regularizer, which anchors the latent space to action-relevant features. This mechanism ensures the latent representations are both predictive and controllable, satisfying theoretical properties like approximate equivariance and homomorphism. The method simplifies training by removing the need for frozen encoders or distributional regularizers, enabling stable end-to-end learning from offline data. Empirical results validate the approach across multiple environments, establishing a new state-of-the-art in stable latent control models.
Novelty
This is the first work to employ inverse dynamics as the sole anti-collapse regularizer in a fully end-to-end latent world model trained on offline data. Unlike prior methods relying on complex regularizers or frozen encoders, this approach leverages a simple, theoretically grounded mechanism to preserve control-relevant information. It explicitly encodes the causal relationship between actions and state transitions, aligning with cognitive theories of perception for action, and demonstrates superior stability and interpretability.
Limitations
- The model assumes a linear or near-linear relationship between actions and latent state changes, which may not hold in highly nonlinear or chaotic environments. Its performance in such settings remains to be tested.
- Dependence on offline data quality and diversity limits generalization; in environments with sparse or biased data, the model may underperform.
- Current implementation focuses on continuous actions and single-step transitions; extending to discrete actions or multi-step dynamics poses additional challenges.
Future Work
Future directions include integrating multi-step inverse dynamics, handling highly nonlinear dynamics, and scaling to multi-agent systems. Combining this approach with reinforcement learning could enable online adaptation and policy learning. Additionally, exploring unsupervised or semi-supervised data collection strategies may reduce data dependency, broadening applicability to real-world robotics and autonomous systems.
AI Executive Summary
Building intelligent agents capable of understanding and controlling complex environments remains a central challenge in AI. Traditional world models often rely on pixel-level reconstruction or complex regularizers, which can lead to unstable representations and limited control. This paper introduces a novel approach—sensorimotor world models (SMWM)—that leverages inverse dynamics regularization to learn stable, interpretable, and control-relevant latent representations directly from offline data.
The core idea is simple yet powerful: by training a joint encoder, forward, and inverse models, the system enforces the preservation of action information in the latent space. This prevents the common problem of representation collapse, where models ignore controllable factors. The inverse dynamics head predicts actions between consecutive states, and its gradients are propagated into the encoder, ensuring the latent space encodes controllable degrees of freedom. The approach requires only a single regularizer, making it both elegant and practical.
Empirical results across 2D and 3D control tasks demonstrate that SMWM learns low-dimensional, spatially faithful, and interpretable representations. In a 2D dot environment, PCA analysis shows the latent space captures the true 2D structure, with over 80% success in control tasks, outperforming baselines by significant margins. In complex environments like OGBench-Cube, success rates exceed 85%, confirming the model's robustness. Ablation studies reveal that removing inverse dynamics regularization causes representation collapse and performance degradation.
This work has broad implications for autonomous control, especially in offline settings where reward signals are unavailable. Its simplicity, stability, and interpretability make it a promising foundation for future research, including multi-agent systems, nonlinear dynamics, and online adaptation. By aligning perception with control, SMWM paves the way for more capable and reliable autonomous agents in real-world applications.
Deep Dive
Limitations & Outlook
What gaps remain?
Abstract
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.