Reinforced Planning with Latent World Models

Humans solve complex problems by constructing plans and mentally simulating their outcomes with an internal model of the world, and this work introduces Reinforced Planning (RP1), a method that learns how to improve multi-step plans by reinforcing good search rules into a neural planner. RP1 is the first model-based planner to fully learn an update rule over a multi-step action plan, significantly outperforming hand-designed search algorithms while using vastly fewer world-model rollouts.

The gist

RP1 learns both how to evaluate imagined outcomes through a critic and how to improve multi-step plans through an optimizer trained fully offline from imagined world-model rollouts.

Background on Planning Limitations

Existing model-based planners typically rely on fixed or inherited planning rules, such as CEM or MPPI, which are either hand-designed or learned from a hand-designed optimizer. These methods often require thousands of world-model evaluations per decision and configurations that must be chosen separately for different tasks and models. Other approaches exist:

1. Amortized model-based control (e.g., Dreamer methods) train an amortized policy but do not plan through the world model at inference time.

2. Planning to inform amortized policies (e.g., IBP and Thinker methods) learn which imagined trajectories to construct or inspect, but they do not iteratively improve a candidate plan itself and necessitate online learning.

3. Iterative next-action optimization (e.g., IAPO) learns an iterative optimizer for the current-state action distribution, which remains one-step improvement rather than learning an update rule over a multi-step plan.

Reinforced Planning Architecture

RP1 follows an actor-critic architecture where a critic scores latent states with respect to a goal, and an actor, the learned planner, optimizes the action plan against the critic’s final state estimate. The process involves three steps per refinement round:

1. Roll out: Generating the next imagined state via the N-step rollout operator, defined as composing N forward rolls of the world model: Hϕ(a, zˆ0):= hϕ(hϕ(· · · hϕ(ˆz0, a0)· · ·, aN−2), aN−1= ˆzN).

2. Evaluate: Computing the terminal value and gradient using the critic: vk = Vψ(zˆN, zg), gk = ∇ak Vψ(zˆN, zg).

3. Improve: Applying the learned update rule to generate the next plan iteration: ak+1 = Fθ(ak, vk, gk).

Learning and Training Objectives

The planner is trained to minimize the terminal cost-to-go predicted by a frozen world model Hϕ and value function Vψ. The objective is formalized as: θ⋆ = arg min θ E(z0,zg)∼D [Vψ(Hϕ(aK, zt), zg)] + C. Updates that produce lower-cost imagined plans are reinforced in the shared parameters of Fθ, while higher-cost plans are suppressed. The planner is trained entirely offline using cached latents and a goal-conditioned critic learned via offline temporal-difference (TD) learning.

Empirical Results and Advantages

RP1 significantly outperforms hand-designed search algorithms across visual navigation (TwoRoom), continuous-control reaching (Reacher), and robotic manipulation (OGBench Cube). Key findings include:

: RP1 exceeds the strongest existing planners while using only 9 world-model rollouts per decision, compared with 9,000 for the strongest competitor method.

In TwoRoom, RP1's learned value function saturates the benchmark across implemented planners. In Reacher, even when latent L2 distance is an adequate surrogate for cost-to-go, RP1 matches or exceeds baselines while using two to three orders of magnitude fewer world-model queries.

Furthermore, the paper demonstrates that a learned neural planner can adapt its update rule across tasks: a learned neural planner can adapt its update rule across tasks, whereas a conventional optimizer uses one fixed configuration throughout the task distribution.

Addressing Model Errors and Speed

The system incorporates mechanisms to handle model inaccuracies. When world models hallucinate contact outcomes, a Dyna iteration is deployed: one Dyna iteration [46] deploys the trained planner, collects its (failure) rollouts, finetunes the world model on them, and retrains the planner. This mitigates exploitation gaps. Regarding speed, RP1 achieves substantial latency reductions: it is up to 67× faster than the strongest alternative under concurrent inference when multiple control loops share one GPU. The reduction in world-model rollouts translates into a "1,000× reduction in world-model rollouts.

Here are specific improvements to AI systems based on the findings presented in the paper, along with what these improved systems can achieve:

Improvements for AI Systems

The core improvement proposed is a paradigm shift from hand-designed search rules to a learned, adaptive planning procedure—specifically, the introduction of Reinforced Planning (RP1). This involves learning both how to evaluate imagined outcomes (via a critic) and how to iteratively improve multi-step action plans (via an optimizer).

1. Learning an Adaptive Update Rule for Multi-Step Plans

The system moves beyond fixed search rules (like CEM or MPPI) by training a neural planner, denoted as the actor, to produce plan updates based on the current plan, its predicted terminal value from a critic, and the gradient of that value with respect to the plan.

2. Goal-Conditioned Critic Learning via Offline Temporal-Difference (TD) Learning

The system learns a goal-conditioned quasimetric critic function, trained entirely offline from imagined world-model rollouts. This critic estimates the cost-to-go (temporal distance) from a latent state to a goal, which is superior to simple Euclidean latent distance because it captures temporal reachability.

3. Reinforcement through Plan Improvement (The RP1 Loop)

Instead of relying on a single optimization method, the planner is trained by reinforcing good planning rules. The system iterates:

1. Roll out the current plan using a frozen world model.

2. Evaluate the resulting terminal state using the learned critic to get a value and gradient (which indicates which parts of the plan are beneficial).

3. Apply a learned residual update network to modify the current plan based on this evaluation, producing an improved candidate plan for the next iteration.

4. Model Error Exploitation via Dyna Fine-Tuning

The system incorporates a Dyna loop where, if the frozen world model makes contact or transition errors (hallucinations), the planner can exploit these errors by performing a single iteration of model finetuning on failure rollouts, and then retrain the planner.

5. Task-Specific Optimization via Interface Composability

The theoretical framework suggests that a single learned planner architecture can adapt its update behavior across different tasks (e.g., navigation vs. manipulation) by exposing task-dependent information through a fixed interface, leading to strictly better performance than any fixed configuration for a given task distribution.

Capabilities of the Improved AI System

The resulting improved AI systems will demonstrate superior performance across complex visual control and manipulation domains:

1. Superior Visual Navigation (TwoRoom Domain)

The system can solve navigation tasks where the geometric proximity of states does not correlate with temporal reachability (e.g., passing through a doorway). By using the learned value function instead of latent L2 distance, the system correctly prioritizes temporally reachable paths, leading to near-perfect success rates (up to 98.2% on PLDM) while drastically reducing world-model rollouts (using 9 rollouts per decision versus thousands for competitors).

2. Robust Robotic Manipulation (OGBench Cube Domain)

The system excels in contact-rich manipulation where the objective is discontinuous and multimodal. It can accurately determine the optimal sequence of actions to pick up an object and place it at a goal, even when early plan changes are critical for success, by learning to navigate this complex, non-smooth objective landscape.

3. High-Speed Inference and Control

The system achieves significant speedups in real-time planning (up to 67× faster under concurrent inference) because its planning step involves only a small number of world-model rollouts per decision (e.g., 9 instead of thousands), making it viable for high-frequency control loops, such as those required for multiple robotic arms operating in tandem.

4. Generalization Across Diverse Tasks

Due to the learned update rule and the task-specific interface, the system can adapt its planning strategy across different environments (visual navigation, reaching, manipulation) without requiring a complete retraining of the search algorithm for each new task type.

Humans solve complex problems by constructing plans and mentally simulating their outcomes with an internal model of the world. Machine learning has made substantial progress in learning world models that predict the consequences of action sequences, yet the procedures used to plan with these models remain largely hand-designed. Most planners rely on fixed search or optimization rules; approaches that learn aspects of search typically imitate a predefined optimizer or use planning to inform an amortized policy, rather improving multi-step plans. We introduce Reinforced Planning, a method that learns the plan-update itself by reinforcing update rules that produce better plans, using gradients propagated through a differentiable world model. We instantiate Reinforced Planning in RP1, which learns a critic over imagined outcomes via temporal-difference learning and a neural plan-improvement operator trained via imagined rollouts with a pretrained world model. RP1 can be trained fully offline without environment interaction; environment episodes are used only for checkpoint selection. Across visual navigation, arm reaching, and robotic manipulation on two world-model backbones, RP1 matches or exceeds existing planners, achieving near-perfect success in several settings while using 1, 000 times fewer world-model rollouts than the strongest alternative (CEM) and planning up to 67 times faster under concurrent planners inference.