Imagine-RL improves robot manipulation by predicting future outcomes
Imagine-RL: Residual-Confidence-Guided Cross-Attention for World-Model-Augmented VLA Reinforcement Learning
Robotics
Summary
Robot arms need to understand not just what they see now but what will happen next when they touch and move things. The authors made a system called Imagine-RL that imagines future movements and forces without creating full images. This helps the robot guess the best action by combining what it sees now with what it expects to happen. They tested this on real robots doing four different tasks and found it works much better than older methods, using less training data.
What this means in practice
- •For robotic manipulation engineers: Improve real-world robot arm control by evaluating candidate actions through predicted future visual and tactile information without retraining core policies.
- •For industrial automation teams: Enhance precision and success rates in automated assembly tasks by integrating future consequence prediction into existing robot controllers.
Authors
Kejia Hu, Wentong Zhai, Bo Zhao, Shuai Liang
Abstract
Reliable action evaluation in contact-rich manipulation requires looking beyond the current observation to future visual and contact consequences. Existing noise-space reinforcement learning efficiently steers a frozen Vision-Language-Action (VLA) policy, but its critics largely ignore these consequences. We present Imagine-RL, which augments noise-space VLA post-training with action-conditioned visual-torque imagination. For each candidate action chunk, a frozen visual-torque latent world model (VTLWM) autoregressively predicts compact future representations without pixel reconstruction. A current image-state-action query attends to observed histories and predicted futures, while previous-window prediction residuals provide token-wise confidence priors that suppress unreliable future tokens. By combining current evidence with predicted consequences, the action critic better evaluates candidate actions and supervises the actor, while the VLA and VTLWM remain frozen. Across four real-robot tasks with 50 evaluation trials per task, Imagine-RL uses only 100 RL trajectories and improves the average success rate by (23.6%) over DSRL and by (60%) over VLA baselines.