Langevin Rollout Optimization for Modelic Reinforcement Learning

Planning-driven model-based (modelic) reinforcement learning has achieved impressive success in continuous control tasks but predominantly relies on zero-order optimizers like Model Predictive Path Integral (MPPI). While robust for global exploration, MPPI updates actions solely through sampling and neglects the smooth return gradients inherent in structured dynamics that guide fine-grained search. To complement MPPI’s robustness with gradient-guided precision, we first propose \textbf{La}ngevin \textbf{R}ollout \textbf{O}ptimization (LaRO), which leverages return gradients to refine actions via Langevin dynamics, achieving reliable local convergence without sacrificing multimodal exploration. This is supported by a score-augmented world model that jointly learns dynamics and a score function within a unified latent space, facilitating efficient and accurate gradient estimation for real-time planning. Second, we combine MPPI and LaRO through a simple yet effective choice mechanism, termed \textbf{M}aximum \textbf{L}ook-\textbf{A}head \textbf{P}lanning (MLAP). Finally, we instantiate MLAP within the latest BOOM algorithm, replacing its MPPI-only planner and yielding BOOM-L. Empirical results on the DeepMind Control Suite and Humanoid Bench demonstrate that BOOM-L consistently outperforms strong baselines in both sample efficiency and final performance.

Modern robots and AI agents learn to make decisions by trial and error, which can take an enormous amount of time. A promising shortcut is to give the agent a "world model" — an internal simulator that predicts what will happen if a certain action is taken — so the agent can mentally rehearse decisions before acting in the real world. However, current methods for searching through possible actions using this world model rely purely on random sampling, essentially throwing darts and seeing which ones land well. This ignores a valuable signal: the world model is a mathematical function, and mathematics tells us the direction in which actions can be improved. We propose a new planning method called LaRO (Langevin Rollout Optimization) that follows this mathematical gradient while also injecting controlled randomness, so the search moves purposefully toward better actions without getting stuck. To make this efficient enough for real-time control, we train an auxiliary network that memorizes the optimization landscape of the world model, replacing expensive on-the-fly computations with a fast lookup. We combine LaRO with traditional random sampling in a parallel framework that always picks the better of the two. On 35 challenging continuous control tasks involving humanoid and quadruped robots, our method consistently outperforms strong baselines, with especially large gains on the hardest tasks where random search breaks down entirely.