Neural Motion Simulator: Pushing the Limit of World Models in Reinforcement Learning

TL;DR

This work introduces MoSim, a neural motion simulator that delivers state-of-the-art long-horizon prediction of an embodied system's physical state by combining a physics-informed predictor with learned correctors within a Neural ODE framework. By modeling $\dot{\boldsymbol{s}}(t)=\boldsymbol{f}(\boldsymbol{s}(t),\boldsymbol{a}(t))+\boldsymbol{\epsilon}(\boldsymbol{s}(t),\boldsymbol{a}(t))$ and decomposing $\boldsymbol{f}$ into a rigid-body component with $\ddot{\boldsymbol{q}}=M(\boldsymbol{s})[\boldsymbol{b}(\boldsymbol{s})+\boldsymbol{\tau}(\boldsymbol{a})]$, MoSim achieves robust, long-horizon predictions that enable zero-shot model-based RL and easy integration with any model-free RL algorithm. The authors demonstrate strong raw and latent-space prediction performance, show zero-shot and few-shot RL improvements, and introduce techniques to handle distribution shifts (e.g., residual-flow penalties) and to quantify horizon requirements for zero-shot learning. Overall, MoSim offers a practical path to decouple environment modeling from RL algorithm development, improving data efficiency and generalization for embodied systems.

Abstract

An embodied system must not only model the patterns of the external world but also understand its own motion dynamics. A motion dynamic model is essential for efficient skill acquisition and effective planning. In this work, we introduce the neural motion simulator (MoSim), a world model that predicts the future physical state of an embodied system based on current observations and actions. MoSim achieves state-of-the-art performance in physical state prediction and provides competitive performance across a range of downstream tasks. This works shows that when a world model is accurate enough and performs precise long-horizon predictions, it can facilitate efficient skill acquisition in imagined worlds and even enable zero-shot reinforcement learning. Furthermore, MoSim can transform any model-free reinforcement learning (RL) algorithm into a model-based approach, effectively decoupling physical environment modeling from RL algorithm development. This separation allows for independent advancements in RL algorithms and world modeling, significantly improving sample efficiency and enhancing generalization capabilities. Our findings highlight that world models for motion dynamics is a promising direction for developing more versatile and capable embodied systems.

Figures (8)

• Figure 1: This figure demonstrate the long-horizon precise prediction by the Neural Motion Simulators. In each of the three pictures, the first row shows the ground-truth states and the second row shows the predicted states with the same initial condition and actions sequence. Humanoid predicts for 30 steps with rendering every 3 steps; Panda predicts for 200 steps with rendering every 20 steps; myohand predicts for 400 steps with rendering every 40 steps.
• Figure 2: (a) Structure of predictor for a sub-step. (b) Structure of MoSim using Neural ODE to integrate many sub-steps to make a prediction for the next state.
• Figure 3: Ablation study of inductive bias on Hopper-Hop.
• Figure 4: Ablation study of training method on Hopper-Hop.
• Figure 5: Policy learning with different prediction horizen