WEAVER, Better, Faster, Longer: An Effective World Model for Robotic Manipulation
About
The potential impacts of world models (WMs, i.e., learned simulators) on robotics are far-reaching -- policy evaluation, policy improvement, and test-time planning -- all with limited real-world interaction. To unlock these downstream capabilities, a WM needs to jointly satisfy three desiderata: $\textit{(i)}$ fidelity (i.e., producing simulated trajectories that correlate with reality), $\textit{(ii)}$ consistency (i.e., producing simulated trajectories that are coherent over long horizons), and $\textit{(iii)}$ efficiency (i.e., producing simulated trajectories quickly). We propose WEAVER (World Estimation Across Views for Embodied Reasoning): a WM architecture that simultaneously achieves all three desiderata, providing state-of-the-art results on robotic manipulation tasks. WEAVER is a multi-view WM trained to predict future latents and reward values via a flow-matching loss. We distill the key design decisions across model architecture, memory, and prediction objectives required to unlock the kinds of long-horizon dynamic manipulation tasks that have confounded prior world modeling approaches. We apply WEAVER in robotic hardware, demonstrating its effectiveness at policy evaluation ($\rho$=0.870 correlation with real-world success rate), policy improvement (real-world success rate improvement of $38\%$ on top of the $\pi_{0.5}$ robot foundation model), and test-time planning (real-world success rate improvement of $14\%$ with a $5-10\times$ speedup over prior WMs). WEAVER also demonstrates better performance than prior WMs when evaluated on out-of-distribution scenarios. Code, models, and videos at: https://arnavkj1995.github.io/WEAVER/ .
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Video Generation | DROID Wrist | LPIPS0.364 | 16 | |
| Trajectory Generation | Task Data Exterior OOD (test) | LPIPS0.116 | 8 | |
| Video Generation | DROID (Exterior) | LPIPS0.117 | 8 | |
| Video Generation | New dataset Exterior OOD | LPIPS0.154 | 8 | |
| Trajectory Generation | Task Data Wrist OOD (test) | LPIPS0.303 | 8 | |
| World Modeling | DROID (val) | Exterior FID9.51 | 4 | |
| World Modeling | Task data OOD (val) | Exterior FID23.48 | 4 | |
| Reward Prediction | Out-of-distribution task dataset (test) | RMSE0.188 | 3 |