Sensorimotor World Models: Perception for Action via Inverse Dynamics

How can we train latent world models from high-dimensional observations without suffering from representation collapse, while ensuring the learned representations are useful for downstream control tasks? The authors investigate whether grounding the latent space in the agent's ability to act—rather than just visual fidelity—provides a more robust and efficient learning objective.

The authors propose the Sensorimotor World Model (SMWM), which learns a latent representation by jointly optimizing three components: an encoder, a forward dynamics model, and an inverse dynamics model. The key innovation is the use of an inverse dynamics loss as a regularizer. While the forward model predicts the next state from the current state and action, the inverse model predicts the action taken between two consecutive states. By forcing the latent space to preserve enough information to recover the action, the model is naturally biased toward capturing the controllable degrees of freedom in the environment while discarding irrelevant visual distractors.

• Anti-Collapse: The inverse dynamics loss acts as a simple, effective mechanism to prevent representation collapse, eliminating the need for complex techniques like frozen encoders or exponential moving averages.
• Controllable Structure: In controlled experiments, the model correctly identifies the intrinsic dimensionality of the environment. It allocates latent capacity only to controllable variables, effectively filtering out random, uncontrollable background motion.
• Spatial Faithfulness: The learned latent space preserves the topology of the environment, meaning that spatial neighbors in the observation space remain neighbors in the latent embedding space.