Robot action models improve policy with repair and rejection steps
$R^2$-WAM: Repair-and-Reject Post-Training for World Action Models
Robotics
Summary
Predicting what will happen when a robot takes an action can sometimes look right but actually be wrong, which can cause problems when teaching robots new skills. The authors developed a two-step method called R2-WAM to fix these predictions and remove bad examples so the robot learns better. First, they check that the predicted movements match what the robot really does, making sure predictions reflect the input actions. Then, they reject imagined outcomes that are worse than real examples to fine-tune the model. This method helps robots plan better without needing extra data or changes in how they think.
What this means in practice
- •For robotics engineers: Improve robot action prediction models to enhance skill learning and task success without extra environment trials.
- •For manufacturing automation teams: Deploy more reliable robots for complex tasks like folding clothes by ensuring predicted actions closely match actual movements.
Authors
Ruiyan Xu, Haisheng Su, Sixu Lin, Zhaokun Yue, Chengming Hu, Xin Jin, Guiliang Liu
Abstract
World Action Models (WAMs) emerge as a promising foundation for policy refinement by predicting the consequences of sampled actions. However, visually plausible predictions can mislead policy refinement if they fail to reflect the input actions. To address this mismatch, we introduce $R^2$-WAM, a two-stage repair-and-reject post-training framework that first improves the consistency of predicted futures with input actions, then uses these futures to select inferior action samples for negative fine-tuning. The repair stage grounds imagination in observed robot behavior through a kinematic alignment score that measures agreement between predicted and demonstrated motion, enabling the predicted video to faithfully reflect its input actions. Using the repaired video model, the rejection stage compares imagined outcomes of sampled and demonstrated actions, selectively applying negative fine-tuning to samples whose predicted task progress falls below the demonstrated reference by a prescribed margin. Together, the two stages extend video prediction from representation learning to consequence-based policy refinement without additional environment interaction or changes to the inference procedure. $R^2$-WAM achieves 93.8% average success on RoboTwin 2.0 across clean and randomized settings. On the long-horizon real-world Fold Shirt task, it achieves 87.5% average success, compared with 0% for Fast-WAM.