Prioritized Rollouts for Efficient World Model-based Vision-Language-Action Policy Optimization
Authors: Yifei Sheng, Haoxiang Ren, Zhilong Zhang, Haonan Wang, Runjie Xu, Yihao Sun, Nan Tang, Zhichao Wu, Lei Yuan, Haoxin Lin, Yang Yu
Organizations: National Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China · School of Artificial Intelligence, Nanjing University, Nanjing, China · Cirquar Technologies, Nanjing, China · Mila – Quebec AI Institute · Université de Montréal
Vision-Language-Action (VLA) models have emerged as a powerful paradigm for embodied intelligence, but fine-tuning them with reinforcement learning (RL) remains constrained by the cost of real-world robot interaction. Model-based reinforcement learning (MBRL) reduces this cost by using a learned world model to generate rollouts for policy optimization. However, it becomes computationally expensive as VLA policies and world models scale. Existing methods typically treat states equally, overlooking substantial differences in their utility for policy improvement. In this paper, we show that policy uncertainty helps identify states with greater potential for policy improvement. The policy exhibits high uncertainty at only a small subset of states, often during decision-sensitive stages where small action differences can alter task outcomes, suggesting that policy improvements at these states could be particularly valuable. Building on these findings, we introduce U-GROW, a lightweight, plug-and-play sampling layer that directs more model rollouts to these informative states. By modifying only the branched-start distribution, U-GROW can be integrated into existing MBRL pipelines without changing the policy optimization objective. Experiments in both simulated and real-world manipulation tasks demonstrate the efficiency and effectiveness of U-GROW, supporting the use of policy uncertainty to guide experience generation.
Figures & tables
Figure 1: Overview of U-GROW
Figure 2: (a) Success rate when a stronger policy replaces the behavior policy at high- or low-uncertainty states. (b) Uncertainty across 3,532 visited states; the dashed line marks the 80th percentile.
Figure 3: Policy uncertainty across manipulation phases. The curve shows policy uncertainty U(st,zt) along a representative trajectory. Shaded regions mark low- and high-uncertainty segments, with corresponding snapshots shown at the top and bottom, respectively. Uncertainty remains low as the robot approaches the red block but rises during placement and release, which require precise spatial alignment and timing.
Figure 4: Policy uncertainty in a post-deviation failure tail. The curve shows policy uncertainty U(st,zt) along a representative failed trajectory. Snapshots illustrate a failed grasp that displaces the red block outside the camera’s field of view, making recovery difficult. The shaded region marks the resulting post-deviation failure tail, in which the policy exhibits high uncertainty at several states.
Figure 5: Success rate versus the cumulative number of training samples on RoboTwin 2.0 and LIBERO. Each training sample is a transition generated by the world model. Both methods use the same world model and total rollout budget; only the rollout-start distribution differs.
RoboTwin 2.0
LIBERO
Method
BH
BR
HB
MC
Avg.
Spatial
Object
Goal
Long
Avg.
πbase
57.0
46.0
17.0
43.0
40.8
78.4
94.8
87.8
59.4
80.1
πRL
60.0
50.0
21.0
48.0
44.8
84.2
96.6
88.4
65.4
83.7
IDQL
59.0
48.0
21.0
44.0
43.0
79.0
95.4
88.6
61.0
81.0
VLA-MBPO
66.0
68.0
30.0
53.0
54.3
88.8
97.8
95.0
69.0
87.7
U-GROW (Ours)
71.0
72.0
37.0
55.0
58.8
90.4
98.4
96.0
71.4
89.1
Table 1: Main Results.
Figure 6: Real-world manipulation tasks on the Arx-X5 bimanual platform. Blocks Ranking requires arranging three colored blocks according to a language instruction, while Blocks Stacking requires stacking them in a specified order. Fold Towel involves bimanual folding of a deformable object, and Plug Cable requires sub-centimeter alignment to insert a cable into a socket.
Figure 7: Success rate over the course of MBRL training on the four real-world tasks, evaluated on 50 trials per task.
Figure 8: Success rates on the four RoboTwin 2.0 tasks in unseen randomized environments. Policies are trained in the clean setting and evaluated over 100 rollouts per task under variations in textures, distractor objects, lighting, and tabletop geometry.
Figure 9: U-GROW ablation on Handover Block .
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Representative manipulation sequences across the three experimental settings.
Figure 11: Clean-to-randomized evaluation on RoboTwin 2.0. The top row shows the clean training setting, and the bottom row shows representative randomized test scenes. From left to right, the columns correspond to Beat Block Hammer, Blocks Ranking RGB, Handover Block, and Move Can Pot. Policies trained in the clean setting are evaluated under visual and geometric shifts without additional training in the randomized environments.
Figure 12: Arx-X5 platform for real-world data collection and policy evaluation.
Figure 13: Objects used in the real-world experiments. Left (seen): three towels for Fold Towel , colored blocks for Blocks Ranking and Blocks Stacking , and a power strip and plug adapter for Plug Cable . Right (unseen): two held-out towels for evaluating Fold Towel on unseen object instances.
Vision-Language-Action (VLA) models show strong generalization for robotic control, but finetuning them with reinforcement learning (RL) is constrained by the high cost and safety risks of real-world interaction. Training VLA models in interactive world models avoids these issues but introduces several challenges, including pixel-level world modeling, multi-view consistency, and compounding errors under sparse rewards. Building on recent advances across large multimodal models and model-based RL, we propose VLA-MBPO, a practical framework to tackle these problems in VLA finetuning. Our approach has three key design choices: (i) adapting unified multimodal models (UMMs) for data-efficient world modeling; (ii) an interleaved view decoding mechanism to enforce multi-view consistency; and (iii) chunk-level branched rollout to mitigate error compounding. Theoretical analysis and experiments across simulation and real-world tasks demonstrate that VLA-MBPO significantly improves policy performance and sample efficiency, underscoring its robustness and scalability for real-world robotic deployment.
Zhilong Zhang, Haoxiang Ren, Yihao Sun +6
National Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China · School of Artificial Intelligence, Nanjing University, Nanjing, China · Mila - Quebec AI Institute +1
Vision-Language-Action (VLA) models map observations to actions with no objective that accounts for how the world responds, so their robustness is bounded primarily by data coverage. World models carry precisely that missing objective and are better grounded for it, yet rolling the future forward costs seconds per decision and rules them out of the control loop. We show the two can be separated. What a world model knows about physical scenes lives in its internal features; generating the future is merely the objective that produced them, so the grounding can be inherited while the generative machinery is left behind. We add one feature-alignment term to ordinary VLA training: a frozen world model is run over the training frames once and cached, and the student learns to agree with that cache. No teacher is loaded during training, the projector is discarded after it, and the deployed policy is identical to the undistilled baseline, running in 32ms and 1.86GB on a consumer RTX5090, so every gain is attributable to the representation rather than to added capacity or test-time compute. A 0.8B student reaches 97.9% on LIBERO, improves from 48.2% to 50.5% on RoboCasa-GR1 humanoid manipulation, and the same objective carries over to real hardware, on both a single-arm and a bimanual platform. The gain survives changes of student scale, backbone, alignment layer, and teacher, indicating a broad representational prior rather than a fragile alignment between two particular networks. Project page: https://thaw-vla.trung-dt.com/.
Trung Dao, Sankalp Yamsani, Jaden Park +2
University of Wisconsin–Madison, USA. · University of Illinois Urbana-Champaign, USA.
Recently, vision-language model (VLM) agents have shown promising progress in open-world tasks, where successful task completion often requires multiple turns of visual perception and action execution. However, existing methods still rely primarily on Supervised Fine-Tuning (SFT) with expert demonstrations, while the advanced reinforcement learning (RL) algorithm, specifically Group Relative Policy Optimization (GRPO), has not been effectively employed for multi-turn RL in these tasks because standard GRPO requires full trajectories as training samples which leads to excessively long context and noise. To address this issue, we propose GROW, a RL framework for open-world VLM agents that decomposes collected trajectories into state-action samples, and computes advantages between these samples rather than treating a full trajectory as a single entity. We further provide a surrogate analysis indicating that, even though the grouped samples are conditioned on different local states rather than an identical prompt context, the objective can preserve the core relative policy optimization signal of GRPO under simplifying assumptions. Experiments on more than 800 Minecraft tasks show that our method achieves state-of-the-art (SOTA) performance, demonstrating the effectiveness of our proposed RL framework for open-world VLM agents.
Xiongbin Wu, Zhihao Luo, Shanzhe Lei +7
Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory · East China Normal University +2