Reinforced Planning with Latent World Models
RP1 learns multi-step plan improvement rules via reinforcement, achieving near-perfect success with 1/1000 model rollouts and 67× faster inference in robotic tasks.
Key Findings
Methodology
RP1 employs an actor-critic framework, where the critic is a goal-conditioned quasimetric value function trained via offline TD learning from imagined world-model rollouts. The neural planner iteratively refines action sequences by learning a residual update rule (Fθ), which takes current plans, terminal state values, and gradients to produce improved plans. This process avoids reliance on traditional optimizers, enabling fully learned, adaptive multi-step plan improvement. Experiments across visual navigation, robotic arm reaching, and manipulation tasks demonstrate that RP1 significantly outperforms hand-crafted search algorithms, achieving success rates near 100% with only 9 model rollouts per decision, compared to thousands for baselines, and reduces inference latency by up to 67×.
Key Results
- In the TwoRoom navigation task, RP1 reaches 100% success with only 9 rollouts, outperforming CEM and MPPI which require 9000 rollouts, and reduces planning latency by 67× in multi-control settings.
- In the Reacher task, RP1 achieves 82% success at a tight tolerance of 0.05 radians, surpassing traditional methods by over 20%, with only 3000 rollouts, demonstrating high efficiency.
- In OGBench Cube manipulation, RP1 attains 89% and 82% success at 25 and 100 steps respectively, far exceeding hand-designed algorithms, with only 9 rollouts per decision.
- Dyna fine-tuning further boosts success rates by 12-19%, alleviating model hallucination issues and enhancing robustness.
Significance
This work pioneers fully learned multi-step plan optimization within model-based control, overcoming sample inefficiency and slow inference of prior methods. By learning search rules directly, it enables rapid, adaptive decision-making suitable for real-time robotics and navigation. Its generality and efficiency address longstanding bottlenecks, advancing autonomous systems’ capability to operate in complex, dynamic environments. The approach opens new avenues for end-to-end reinforcement learning in planning, promising broader impact across AI and robotics.
Technical Contribution
The paper introduces a novel framework combining a goal-conditioned critic trained via offline TD learning with a neural residual planner that iteratively refines plans. It departs from traditional fixed or imitation-based optimizers, enabling fully learned, task-adaptive search rules. The theoretical analysis demonstrates that latent space distances cannot reliably reflect temporal reachability, emphasizing the importance of learned value functions. Empirically, the method achieves high sample efficiency, significantly reducing model evaluations, and enabling real-time inference in complex tasks. The residual update structure ensures training stability and performance gains.
Novelty
This is the first work to fully learn multi-step plan improvement rules in a model-based setting, moving beyond fixed or hand-designed search algorithms. It integrates offline value learning with iterative plan refinement, allowing the system to adapt its search strategy across tasks. Unlike prior methods limited to single-step or imitation learning, RP1's end-to-end training achieves a flexible, task-aware search policy, representing a significant leap in autonomous planning research.
Limitations
- Model hallucinations, especially contact predictions, can mislead planning, requiring additional correction mechanisms such as Dyna fine-tuning. The approach still depends on accurate world models, which may not generalize well to unseen environments.
- Training relies on large offline datasets, limiting online adaptability and real-time learning capabilities. Extending to online or continual learning remains an open challenge.
- While inference speed is greatly improved, high-dimensional or multi-agent scenarios may still pose computational challenges, necessitating further model and hardware optimization.
Future Work
Future directions include integrating online learning for continuous adaptation, improving model robustness against hallucinations, and extending the framework to multi-agent coordination. Exploring multi-modal inputs and hierarchical planning could further enhance performance in complex, real-world environments. Additionally, developing more efficient architectures and hardware acceleration techniques will be crucial for deploying RP1 in large-scale autonomous systems.
AI Executive Summary
Humans excel at solving complex problems by mentally simulating future outcomes using internal models of the world. This cognitive ability enables flexible planning and decision-making. Inspired by this, machine learning has advanced the development of world models that predict future states based on action sequences, supporting tasks like visual navigation and robotic manipulation. However, leveraging these models for effective planning remains challenging. Traditional approaches rely heavily on hand-crafted search algorithms or fixed optimization routines, which are computationally expensive and lack adaptability across different tasks.
This paper introduces Reinforced Planning (RP1), a novel framework that fully learns how to improve multi-step action plans using reinforcement learning principles. At its core, RP1 employs an actor-critic architecture: a critic estimates the cost-to-go from imagined future states conditioned on goals, trained offline via temporal-difference learning; a neural planner iteratively refines action sequences by learning a residual update rule, which is applied repeatedly during inference. Unlike conventional methods, RP1 does not depend on fixed optimizers like CEM or MPPI, but instead learns a flexible, task-adaptive search policy that can generalize across environments.
The key innovation lies in the integration of a goal-conditioned critic with a learned plan optimizer, enabling the system to evaluate and improve plans efficiently. Theoretically, the authors demonstrate that latent space distances cannot reliably encode temporal reachability, emphasizing the importance of learned value functions. Empirically, RP1 achieves near-perfect success rates in visual navigation, robotic arm reaching, and manipulation tasks, using only 1/1000 of the model rollouts required by traditional methods. It also reduces inference latency by up to 67×, making real-time deployment feasible.
These results mark a significant step toward autonomous systems capable of rapid, sample-efficient planning in complex environments. The approach’s flexibility and efficiency open new avenues for scalable, adaptive AI in robotics, autonomous driving, and beyond. Future work will focus on online adaptation, robustness against model errors, and extending to multi-agent scenarios, pushing the boundaries of autonomous decision-making.
Deep Dive
Abstract
Humans solve complex problems by constructing plans and mentally simulating their outcomes with an internal model of the world. Machine learning has produced world models that similarly predict the outcomes of action sequences, but the improvement of candidate plans still isn't fully learned. Current planners are either hand-designed, distilled from a hand-designed optimizer, or learned only to inform an amortized policy rather than to revise the plan itself. We introduce the Reinforced Planning, a method based on the idea that search can be learned by reinforcing good search rules into a neural planner. Our implementation RP1 learns both how to evaluate imagined outcomes through a critic, as well as how to improve multi-step plans through an optimizer trained fully offline from imagined world-model roll-outs. To our knowledge, RP1 is the first method to fully learn how to improve multi-step plans. Furthermore, it can be trained independently of and attached to any pretrained latent world model. Across visual navigation, arm reaching, and robotic manipulation on two world-model backbones, RP1 substantially outperforms hand-designed search algorithms, reaching near-perfect success in several settings while using $1,000 \times$ less world-model rollouts and being up to $67 \times$ faster than the strongest alternative under concurrent planner inference.