cs.LGSep 19, 2026

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, 

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 · 

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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)U(s_{t},z_{t}) 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)U(s_{t},z_{t}) 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.0LIBERO
MethodBHBRHBMCAvg.SpatialObjectGoalLongAvg.
πbase\pi_{\mathrm{base}}57.046.017.043.040.878.494.887.859.480.1
πRL\pi_{\mathrm{RL}}60.050.021.048.044.884.296.688.465.483.7
IDQL59.048.021.044.043.079.095.488.661.081.0
VLA-MBPO66.068.030.053.054.388.897.895.069.087.7
U-GROW (Ours)71.072.037.055.058.890.498.496.071.489.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 5050 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 100100 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.
SettingRoboTwin 2.0LIBEROReal-World Tasks
Image resolution256×256256\times 256256×256256\times 256256×256256\times 256
Context frames555
Video sequence length251325
Denoising steps555
Guidance scale1.01.01.0
Noise-schedule shift5.05.05.0
Appendix
Table 2: World model training settings.
SettingRoboTwinLIBEROReal-World Tasks
Learning rate2×10−52\times 10^{-5}2×10−52\times 10^{-5}2×10−52\times 10^{-5}
Weight decay0.10.10.1
Adam (β1,β2)(\beta_{1},\beta_{2})(0.9, 0.95)(0.9, 0.95)(0.9, 0.95)
Adam ϵ\epsilon10−810^{-8}10−810^{-8}10−810^{-8}
Learning-rate scheduleCosineCosineCosine
Warmup fraction0.050.050.05
Appendix
Table 3: Reward model training settings.
SettingRoboTwinLIBEROReal-World Tasks
Action chunk20820
Action dimension14714
Denoising steps535
Rollout chunk232
Actor learning rate2.5×10−62.5\times 10^{-6}5×10−65\times 10^{-6}2.5×10−62.5\times 10^{-6}
Critic learning rate5×10−55\times 10^{-5}10−410^{-4}5×10−55\times 10^{-5}
Appendix
Table 4: Policy optimization settings.
More in this topic

RL for VLMs

Reinforcement learning for vision-language models aligns multimodal reasoning, tool use, and action with task-specific feedback.

See all 476 papers

Explore similar work

CardsList

  1. Towards Practical World Model-based Reinforcement Learning for Vision-Language-Action Models

    Mar 21, 2026Zhilong Zhang, Haoxiang Ren, Yihao Sun +6RL for VLMsVision-Language-Action Models

    National Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China · School of Artificial Intelligence, Nanjing University, Nanjing, China · Mila - Quebec AI Institute +1

  2. Think Like a World Model, Act Like a VLA: Distilling World-Model Representations into Compact Robot Policies

    Sep 21, 2026Trung Dao, Sankalp Yamsani, Jaden Park +2World ModelsRobotic Manipulation

    University of Wisconsin–Madison, USA. · University of Illinois Urbana-Champaign, USA.

  3. GROW: Aligning GRPO with State-Action Modeling for Open-World VLM Agents

    May 18, 2026Xiongbin Wu, Zhihao Luo, Shanzhe Lei +7Reinforcement LearningGroup Relative Policy Optimization

    Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory · East China Normal University +2