Source-linked AI summary
Back on Track: Deployment-time Policy Improvement via World Model Planning
Colton Crosby, Hao Wang, Mac Schwager, Somil Bansal
TL;DR
Visuomotor policies can fail during deployment, while improving them conventionally can require costly data, retraining, or policy changes. This paper uses a black-box liveness monitor and short-horizon world-model planner to intervene selectively; it reports higher success on LIBERO-Goal and Franka Panda, alongside limitations from value and model errors and required training data.
Problem
Visuomotor policies can fail during deployment without predicting or recovering from failures, while conventional improvement can require costly data and access to policy parameters and training pipelines.
Source: §I · p. 1Method
A policy-conditioned liveness value monitors a black-box policy and guides a short-horizon planner in a compact world model to make corrective interventions before returning control.
Source: §I · p. 2Results
On LIBERO-Goal, π0.5 improves from 72.6% to 81.5%, while on Franka Panda the diffusion policy improves from 40% to 60%.
Source: §VI · p. 8Takeaways & Limitations
Brief corrective filtering aims to return a capable nominal policy to its basin of success without replacing or retraining it.
Source: §I · p. 2Inaccurate liveness triggers or exploitable world-model errors can disrupt trajectories, and the framework requires labeled policy rollouts plus auxiliary transition data.
Source: §VI · p. 8
Abstract
from arXiv · showhide
State-of-the-art visuomotor policies are increasingly capable of executing complex tasks, yet can fail during deployment in ways they cannot predict or recover from on their own. Improving their performance and robustness typically requires collecting more data, retraining, or modifying the policy architecture, all of which are time- and resource-intensive in practice. We instead propose an offline-trained monitoring-and-intervention framework that improves a policy during deployment using only black-box query access. From a limited set of labeled on-policy rollouts, we learn a policy-conditioned liveness value that measures the nominal policy's progress toward task completion; we additionally train a compact latent world model using these rollouts together with play data. At deployment, the liveness value continuously monitors execution. When it detects a loss of progress, a short-horizon planner searches the world model for corrective actions using the same value as its planning objective, after which control returns to the nominal policy. We evaluate our method on the LIBERO-Goal task suite and on a Franka Panda manipulator, using the vision-language-action model $\pi_{0.5}$ and a diffusion policy as nominal policies, and demonstrate significant improvements in task success rates, while supporting real-time control. Additional results and videos are available on the project website (https://policy-improvement-wm.pages.dev/).
I. INTRODUCTION
The paper proposes an external, black-box wrapper that monitors a deployed visuomotor policy and briefly intervenes to restore progress, rather than modifying or retraining the policy. Its liveness-guided planning is designed for real-time control and improves success across simulation and hardware evaluations.
I. INTRODUCTION
- Despite their growing capabilities, visuomotor policies can fail at deployment without predicting or recovering from those failures, while conventional scaling requires costly data, compute, or architecture changes.Those approaches also require access to policy parameters and the training pipeline.Source: §I · p. 1
- The setting requires an external wrapper learned from few labeled on-policy rollouts, with monitoring and planning fast enough to keep pace with the control loop.The policy is accessed only through black-box queries; its parameters, architecture, and action-generation procedure are unknown.Source: §I · pp. 1–2
- The framework filters for task completion: brief corrective actions aim to return a faltering policy to its basin of success, rather than replacing or retraining it.The intervention seeks a state from which the nominal policy can finish the task.Source: §I · p. 2
- A policy-conditioned liveness value monitors progress and guides a short-horizon planner in a compact latent world model toward states where the nominal policy can recover.The value is learned from labeled on-policy rollouts, and the world model also uses play data.Source: §I · p. 2
- On hardware, a planning cycle takes 108 ms and the monitor–planner loop runs above 10 Hz by amortizing planning over an action chunk.The planner emits an action chunk rather than a single action.Source: §I · p. 2
- On LIBERO-Goal, the method improves π0.5 from 72.6% to 81.5%, and on Franka Panda it improves a diffusion policy from 40% to 60%.The LIBERO-Goal result outperforms offline-RL and safety-filtering baselines.Source: §I · p. 2
II. RELATED WORK
Prior work improves fixed policies through internal steering, candidate reranking, or world-model planning, but differs in required access, supervision, and deployment cost. This paper instead uses selective black-box intervention guided by a learned liveness value and a compact model.
II. RELATED WORK
- Internal-steering methods rely on policy representations or sampling procedures, so they do not apply when the nominal policy is a black box.Examples include Latent Policy Barrier [15], VLS [16], and [17].Source: §II · p. 2
- RoboMonkey [18] and Steer Your Generalists [19] rerank policy-sampled candidates, requiring multiple policy queries per decision; this work queries πb once per intervention and supplies diversity through its own proposal distribution.A diffusion-policy query takes 180 ms in the hardware setting described.Source: §II · p. 2
- Unlike Latent Policy Steering [20], which searches continuously, this method invokes planning selectively and returns control after a short correction when the nominal policy is progressing.This limits unnecessary computation and intervention, as described by the authors.Source: §II · p. 2
- DreamAvoid [12] and VLA-in-the-Loop target predefined critical events and use additional supervision or policy access that this method does not assume.Their reported intervention times are 2 s and 20 s, respectively, compared with this method’s 108 ms planning cycle.Source: §II · p. 2
- Prior liveness filtering can use exact reachability, but that is intractable for visual manipulation; this work uses a learned value both to detect lost progress and guide corrective planning.The learned value builds on [10], which used it for offline policy evaluation.Source: §II · pp. 2–3
- The compact feature-prediction world model is chosen to evaluate many imagined action sequences within a real-time budget, unlike pixel-generating video models with seconds-scale latency.The passage also distinguishes discrete latent-code prediction in DreamerV3 [21] and patch-level feature prediction in DINO-WM [22].Source: §II · p. 3
III. PROBLEM FORMULATION
The formulation seeks to improve a fixed visuomotor policy’s completion probability during execution without modifying or retraining it. An external monitor and temporary fallback planner are learned from sparse rollout outcomes, optionally augmented with unlabeled transitions.
III. PROBLEM FORMULATION
- The objective is to increase a fixed policy’s task success rate during execution without modifying or retraining it.The formulation concerns episodic tasks with measurable completion outcomes and focuses here on visuomotor manipulation.Source: §III · p. 3
- The nominal policy is treated as a black box: its parameters, architecture, and internal representations are unknown.The available task-specific data are rollouts generated by executing that policy.Source: §III · p. 3
- Each policy rollout contains a state-action trajectory and a binary episode-level completion outcome, the method’s only assumed task-specific supervision.Auxiliary state-action transitions may be used to learn dynamics without outcome labels or a specified generating policy.Source: §III · p. 3
- During execution, the nominal policy acts by default; an external monitor decides when a fallback policy should briefly override it and then return control.Neither component accesses or modifies the nominal policy’s internals.Source: §III · p. 3
- The monitor and fallback policy are designed from rollout data and optional auxiliary transitions to maximize completion probability within the deployment control-loop latency budget.Source: §III · p. 3
IV. METHOD
- In the concrete method, one policy-conditioned liveness value detects lost progress and scores the short-horizon planner’s search in a learned world model.The planner acts as the fallback policy and returns control after executing a short correction.Source: §IV · p. 3
A. Monitor
The monitor uses a policy-conditioned liveness value to detect stalled progress, with ensemble averaging and temporal smoothing to address noisy predictions. A short-horizon planner then uses that value to find corrective actions before handing control back.
A. Monitor
- The value is learned from rollout outcomes, and its min-based update propagates goal attainment backward through the nominal policy’s trajectories.Successful episodes label the final state as a goal; other states receive the nong_goal label.Source: A. Monitor · p. 3
- The liveness value is policy-conditioned: it estimates progress toward completion under the nominal policy, falling toward −1 near a goal and rising toward +1 otherwise.A min update propagates goal attainment backward through nominal-policy rollouts rather than summing discounted rewards.Source: A. Monitor · p. 3
- To reduce noise, especially with small rollout datasets, the monitor averages an ensemble’s predictions and smooths them over a temporal window.The smoothing window averages the current prediction with the previous W −1 predictions.Source: A. Monitor · p. 3
- The monitor triggers when the smoothed value exceeds a time-varying threshold that decays from +1 toward 0; the schedule delays intervention during initial configuration.The experiments use W = 3, and the threshold schedule is shared across all ten tasks.Source: A. Monitor · p. 3
B. Planner
- The planner optimizes an action sequence over horizon H but executes only its first h ≤ H actions before returning control to the nominal policy.This short horizon serves correction rather than independent task completion.Source: B. Planner · p. 3
1) Dynamics Model:
A compact latent dynamics model predicts next states from state–action pairs, while a shared liveness value guides planning toward states where the nominal policy can finish the task.
1) Dynamics Model:
- The world model predicts the next compact state from a state–action pair, prioritizing low-latency trajectory evaluation over model capacity.It is trained on policy rollouts and auxiliary play data, which broadens action coverage beyond the nominal policy's behavior.Source: 1) Dynamics Model · p. 4
- Autoregressive training gradually increases rollout length from 1 to 16, moving from local transitions to longer predictions that include the model's own errors.The rollout-length curriculum is used to improve multi-step prediction with a compact model.Source: 1) Dynamics Model · p. 4
2) Planning Objective:
- Planning scores imagined trajectories with the same policy-conditioned liveness value used by the monitor, favoring sustained progress toward states where the nominal policy can complete the task.The planner seeks recoverable states rather than completing the task on its own.Source: 2) Planning Objective · p. 4
3) Planning Algorithm:
The planner uses parallelizable CEM to search for corrective action chunks, sampling around the nominal policy's actions with a data-informed proposal intended to limit implausible candidates.
3) Planning Algorithm:
- CEM searches the non-smooth learned-value objective with parallelizable sampling, while a proposal centered on nominal actions helps avoid implausible candidates that can exploit model errors.The planner refits the proposal to elite candidates, executes the first h ≤ H actions, then returns control to the monitor.Source: 3) Planning Algorithm · p. 4
- The proposal uses empirical action-chunk correlations from rollout data, sampling in leading covariance directions rather than relying on a generic smoothness prior.The selected eigenvectors explain 98% of measured variance; the gripper dimension remains unperturbed and uses the nominal command.Source: 3) Planning Algorithm · p. 4
- At each trigger, the planner queries the nominal policy for an action chunk, initializes CEM around it, and scores candidates through the world model.After a fixed number of iterations, it executes a short prefix of the best sequence and returns control to the monitor.Source: 3) Planning Algorithm · p. 4
V. EXPERIMENTAL RESULTS
On perturbed LIBERO-Goal trials, the intervention improves the nominal policy's aggregate success rate, with task-level results shaped by both recoveries and harmful interventions. The evaluation uses paired trials and compares the method with policy-replacement and filtering approaches.
A. Experimental Setup
- The evaluation uses π0.5's LIBERO checkpoint across ten tasks, with 300 paired episodes per task and shared initial conditions and policy sampling seeds.To create room for intervention, the study perturbs initial manipulator configuration and object locations; the nominal policy averages 72.6% success under this protocol.Source: A. Experimental Setup · pp. 4–5
B. Deployment-time Improvement Across the Suite
- 81.5% versus 72.6% success (+8.9pp) across 3,000 paired LIBERO-Goal trials; interventions rescued 369 nominal failures and broke 102 nominal successes.The method nets 267 additional successful trials; Table I reports full-suite success rates and rescued and broken outcomes.Source: B. Deployment-time Improvement Across… · p. 5, Table I · p. 5
- Table II compares success rates for methods that replace the nominal policy, retrain it with white-box access, or filter it.TD3+BC and IQL replace πb, Filtered-BC retrains it, and the remaining methods filter it.Source: Table II · p. 5
- Task 0 had the largest gain (+27.0pp), while Task 2 regressed by −1.0pp and Task 7, already at 98.0%, reached 98.3%.For Task 2, the world model sometimes hallucinates grasps on the thin wine bottle, which can tip it over; Task 9 improves by +20.7pp.Source: B. Deployment-time Improvement Across… · p. 5
C. Comparison Against Other Policy Improvement Methods
Among black-box policy-improvement methods, the proposed filter achieves the highest mean success rate, while remaining comparable to a white-box fine-tuning reference. The comparison also highlights why replacing a capable nominal policy or using a reachability-based fallback can be less effective.
C. Comparison Against Other Policy Improvement Methods
- 80.9% mean success for our method exceeds LatentSafe (59.2%), IQL (42.2%), and TD3+BC (39.6%) among black-box methods.The alternatives either train replacement policies from narrow rollout data or use a fallback that does not reliably recover lost progress.Source: C. Comparison Against Other Policy… · pp. 5–6
- In Task 0, the proposed method guides the robot toward the middle drawer, while the nominal policy targets the top drawer.This example illustrates the visual difference between the approaches on the open-middle-drawer task.Source: Figure 2 · p. 6
- TD3+BC and IQL discard the nominal policy’s capabilities, while narrow action coverage in nearly deterministic rollouts limits their ability to infer better actions.TD3+BC is also exposed to world-model error and performs particularly poorly on Task 3 (7.3%), where drawer-state modeling matters.Source: C. Comparison Against Other Policy… · p. 6
- LatentSafe breaks 80 nominal successes while rescuing 7 failures on Task 9, leaving mean success below the nominal policy (59.2% versus 65.2%).Its fallback preserves reachability, but the passage notes that recovering after progress is lost requires more than remaining within the reachable set.Source: C. Comparison Against Other Policy… · p. 6
- 80.9% mean success is comparable to Filtered-BC’s 78.2%, although Filtered-BC fine-tunes all 3.3B policy parameters and requires white-box access and H100-scale compute.Filtering on top of Filtered-BC raises mean success to 79.8%, with Task 0 improving from 81.0% to 95.3% and Task 3 falling from 75.3% to 66.0%.Source: C. Comparison Against Other Policy… · pp. 5–6
D. Robustness Analysis and Component Ablations
The method generally withstands increased perturbation severity, while ablations show that monitor design, the liveness objective, and the planner’s action proposal all contribute to performance. Removing play and demonstration data still leaves a substantial Task 0 gain, alongside reduced action-space coverage.
D. Robustness Analysis and Component Ablations
- 18.0pp is the largest mean improvement when configuration perturbations increase by 50%, while object perturbations yield +11.5pp.The stress test increases one perturbation magnitude at a time, holds the other fixed, and leaves the learned value and world model unchanged.Source: D. Robustness Analysis and Component… · p. 6, Table III · p. 6
- The method improves over πb in seven of eight stress-test conditions, with one exception: Task 4 at the largest ±15 cm object displacement is 1.0pp worse.The table describes conditions where one perturbation magnitude is increased by 50% relative to the main setting.Source: D. Robustness Analysis and Component… · p. 6, Table III · p. 6
- Removing the sigmoid threshold schedule costs 8.5pp, using a single value network costs 8.3pp, and removing temporal smoothing costs 4.3pp in mean success.Even the weakest monitor variant still improves over πb by 6.0pp.Source: D. Robustness Analysis and Component… · p. 6, Table IV · p. 7
- Replacing the liveness value with the distributional Monte Carlo value reduces mean success by 13.1pp, the largest ablation effect.The alternative rescues 152 nominal failures versus 211 for liveness, while breaking 120 nominal successes versus 38 for liveness.Source: D. Robustness Analysis and Component… · pp. 6–7
- Replacing the learned action basis with iCEM-colored noise costs 4.7pp, while replanning after two rather than five planned actions costs 1.3pp.The reported comparisons assess the planner’s action proposal and replanning frequency.Source: D. Robustness Analysis and Component… · p. 7
- Using only 200 policy rollouts still raises Task 0 success from 66.7% to 81.3% (+14.7pp), retaining 54% of the full gain and exceeding the strongest baseline.Without play data, 23 principal components are needed to capture 99% of rollout-action variance, compared with 15 with play data; the fixed 13-component basis therefore spans less.Source: D. Robustness Analysis and Component… · p. 7
E. Hardware Experiment
The hardware evaluation tests lunchbox opening under initial-configuration perturbations, where the nominal policy succeeds on 40% of trials and the filter improves success. Interventions recover some difficult runs, including a complete grasp-and-open maneuver in the illustrated trial.
E. Hardware Experiment
- The nominal diffusion policy was trained on 105 expert demonstrations for opening a deformable cloth lid without spilling the lunchbox contents.Testing perturbed the initial manipulator configuration by 0–9°; training data for the monitor and world model included 120 on-policy rollouts plus demonstrations and play data.Source: E. Hardware Experiment · p. 7
- 60.0% success (+20pp) with the filter, versus 40% for the nominal policy under the perturbed test distribution.The largest observed difference is in the [6.5, 9]° band: 70.6% with the filter versus 22.7% for the nominal policy.Source: E. Hardware Experiment · p. 7
- All 16 trials without monitor intervention succeeded, while 14 of 34 trials with at least one trigger succeeded; every observed failure followed a trigger.The authors report that detection is not what limits the result.Source: E. Hardware Experiment · p. 7
- In the illustrated run, the planner intervenes after the policy closes the gripper without grasping the lid, then executes a predicted grasp-and-open maneuver before handing control back.The intervention sequence begins around 1 s and the planner predicts a complete grasp-and-open plan around 6.5 s.Source: E. Hardware Experiment · p. 7
F. Deployment Cost
The framework improves success across LIBERO-Goal and the Franka Panda lunchbox task, while fitting corrective planning into real-time control. Its interventions add computation and can still fail when the liveness value or learned world model is inaccurate.
F. Deployment Cost
- A full planning cycle takes 108 ± 3 ms beyond the nominal policy’s 180 ± 3 ms query, adding 27 ms per executed control step versus the policy’s 45 ms.Inference costs are amortized over four-step action chunks, yielding an effective intervention control rate of approximately 13.9 Hz.Source: F. Deployment Cost · p. 7
- A batch of 256 imagined 16-step rollouts takes 44.2 ms, and the filtered system executes 20% fewer control steps yet achieves higher success.A single 16-step rollout takes 9.5 ms; the authors attribute affordable candidate evaluation to parallel processing in the compact world model.Source: F. Deployment Cost · p. 7
VI. LIMITATIONS AND CONCLUSION
- LIBERO-Goal success rises from 72.6% to 81.5%, and Franka Panda lunchbox success rises from 40% to 60%.The LIBERO-Goal result outperforms offline-RL and alternative filtering baselines.Source: §VI · p. 8
- The method does not guarantee each intervention improves execution: liveness errors can trigger disruptive interventions, and world-model errors can mislead planning.The framework also requires labeled nominal-policy rollouts and auxiliary transition data, so it is not zero-shot.Source: §VI · p. 8