Vision language action policies improve with prioritized model rollouts
Prioritized Rollouts for Efficient World Model-based Vision-Language-Action Policy Optimization
Machine Learning
Summary
Training robots to understand vision and language together is hard and slow because real-world practice takes a lot of time. The authors found that some parts of a robot's decision-making are more uncertain and important for learning than others. They created a method called U-GROW that focuses training on these critical moments, making the robot learn faster and better using computer simulations. Their tests showed this method helps both in virtual environments and with real robot tasks.
What this means in practice
- •For robotics engineers: Improve robot training efficiency by focusing computational efforts on uncertain decision points in simulated environments for manipulation tasks.
- •For autonomous system developers: Enhance policy optimization for vision-language-action systems by prioritizing informative states learned through uncertainty-guided rollouts.
Authors
Yifei Sheng, Haoxiang Ren, Zhilong Zhang, Haonan Wang, Runjie Xu, Yihao Sun, Nan Tang, Zhichao Wu, Lei Yuan, Haoxin Lin, Yang Yu
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.