Action-conditioned world models offer a scalable alternative to costly real-world interaction for improving robot policies. However, existing world-model-based methods largely treat world models as forward-only surrogates of the physical environment, leaving untapped a capability no physical environment affords: rewinding to past decision points to explore counterfactual futures. We introduce Rollback, a policy improvement paradigm that turns world models from passive simulators into counterfactual engines: it rolls back to critical decision points, branches into alternative actions, and contrasts their outcomes to attribute success to individual decisions. This concentrates imagination where it matters most and distills counterfactual outcomes into fine-grained, action-level supervision without a separately learned value model. We build a scalable training system that unifies distributed world-model search with asynchronous real-robot execution, forming a self-reinforcing cycle: real-world experience continually sharpens the world model, and a sharper world model in turn drives further policy improvement. We term the world-model side of this cycle test-time physics grounding, a capability that hand-built physics simulators lack: while the world model serves as the training environment, real transitions from the policy it trains correct its dynamics through a lightweight LoRA pipeline. Experiments on three challenging real-world tasks show that Rollback improves success rates over the strongest baseline by an absolute 10.67%, 11.11%, and 11.85% on block extraction, ball tossing, and tabletop curling, respectively, while achieving more stable improvement throughout training. Although initialized from teleoperated demonstrations far below human level, Rollback surpasses the average success rate of human players on both ball tossing and tabletop curling within 8 hours of training.
Figures & tables
Figure 1: From real-world interaction to learning in world models. (a) Real-world RL incurs costly physical interaction. (b) Forward-rollout world-model RL spends much of its interaction budget on less informative exploration. (c) Our proposed Rollback enables high-value exploration and dense supervision through experience-backtracking tree search, while real-world feedback refines the world model.
Figure 2: Overview of Rollback. (a) Experience-backtracking tree search restores cached world-model contexts to explore informative branches without regenerating shared prefixes. The world model predicts both multi-view observations and the robot’s proprioceptive state. (b) Tree-structured value estimation aggregates downstream outcomes into relative action-quality labels for policy training without an additional value model. (c) World model–policy co-evolution continually incorporates real trajectories to refine the world model through test-time physics grounding (LoRA SFT followed by DMD) while the policy keeps improving.
Figure 3: Rollback training infrastructure. The cloud runs tree search on GPU workers, stores real trajectories and tree supervision in the Experience Fabric, and trains the policy and the world model in parallel. Edge sites run policy inference for their robots, buffer trajectories for upload, and cache policy snapshots. Robots send observations to their edge site and execute the returned actions. Policy snapshots are sent to search workers and edge sites; world-model snapshots are sent only to search workers.
Figure 4: Training process of Rollback. The robot continually executes the latest policy and uploads real trajectories Dreal . The policy is continually updated with tree supervision Dtree from rollback search, while the world model is repeatedly adapted to real trajectories via LoRA SFT and DMD; each search tree uses the latest versions of both.
Method
Category
Block Extraction
Ball Tossing
Tabletop Curling
Average
π0.5 + SFT
Imitation learning
48.67
23.33
48.89
40.30
π0.5 + RECAP
Real-world RL
56.00
25.56
51.11
44.22
π0.5 + RISE w/ GE-Sim2
World-model RL
58.67
28.89
52.59
46.72
Rollback (Ours)
World-model RL
69.33
40.00
64.44
57.93
Gain over strongest baseline
+10.67
+11.11
+11.85
+11.21
Table 1: Performance on three real-world tasks. Success rates (%) of Rollback and imitation-learning, reinforcement-learning, and world-model-based baselines.
Variant
Target Contact (%)
Task Success (%)
Consecutive Extractions
Frozen World Model
78.67
64.67
1.90
Direct Reward Supervision
77.33
59.33
1.73
Rollback (Full)
89.33
69.33
2.60
Table 2: Ablation study on Block Extraction. We evaluate world-model adaptation and tree-structured value estimation via target contact rate, task success rate, and the average number of consecutive extractions; the full framework performs best on all metrics.
Figure 7
Figure 7: Experience-backtracking search on Tabletop Curling. (a) Search tree. The blue dashed path is a real attempt and colored paths are imagined continuations; edges are green, red, or purple when the continuations below all succeed, all fail, or are mixed. Purple rings mark branch points at intermediate decisions, each followed by successful and failed continuations. (b) Two world-model continuations from a branch point inside imagination; (c) the real attempt and a world-model alternative from a rollback point. Each panel shows stone tracks overlaid on the frame at the branch (left) and time-aligned frames (right). In both cases the continuations differ in how fast the second stone is released: a fast release reaches the house, a slow one stops short.
PSNR (dB) ↑
RGB L1 ↓
Reward false positives
Head
Left wrist
Right wrist
First quarter
Last quarter
23.41
20.63
20.85
6.79
11.79
0 / 40
Table 3: Open-loop world-model fidelity on Block Extraction (40 rollouts).
World model
PSNR (dB) ↑
RGB L1 ↓
Rollout spread ↓
Head
Left wrist
Right wrist
First quarter
Last quarter
Frozen
21.50
17.07
15.01
12.48
23.73
9.39
Co-evolved
22.27
17.79
16.37
10.50
14.47
6.12
Table 4: Open-loop world-model fidelity with and without co-evolution on Tabletop Curling (64 validation episodes, three rollouts each).
Figure 8: World-model fidelity with and without co-evolution on Tabletop Curling. (a) Open-loop predictions of the frozen world model (before co-evolution) and the co-evolved world model under the recorded actions of one validation episode, with the real recording on top. Each cell shows the head camera cropped to the sheet; insets magnify the house, and red boxes mark hallucinated stones. Steps are control steps at 30 Hz. (b) RGB L1 error over control steps, averaged over the 64 validation episodes, with 95% confidence intervals shaded.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Component
Hyperparameter
Value
World model
Initialization
GE-Sim 2.0
Initial SFT (steps / GPUs / batch)
10K / 32 / 128
Initial DMD (steps / GPUs / batch)
2,500 / 16 / 16
Per adaptation
400 LoRA SFT + 100 DMD steps
Frames per action chunk / frame rate
25 / 16 fps
Policy
Initialization
π0.5
Appendix
Table 5: Hyperparameters of Rollback.
Method
P1
P2
P3
P4
P5
P6
Overall
Consec. Extr.
π0.5 + SFT
64.00
80.00
44.00
40.00
44.00
20.00
48.67
1.60
π0.5 + RECAP
48.00
84.00
40.00
48.00
72.00
44.00
56.00
1.78
Frozen World Model
72.00
76.00
48.00
68.00
60.00
64.00
64.67
1.90
Direct Reward Supervision
76.00
64.00
48.00
64.00
64.00
40.00
59.33
1.73
π0.5 + RISE w/ GE-Sim2
60.00
80.00
32.00
68.00
52.00
60.00
58.67
2.00
Rollback (Ours)
72.00
56.00
56.00
64.00
84.00
84.00
69.33
2.60
Appendix
Table 6: Per-position results on Block Extraction (success rate, %; 25 trials per position).
Method
P1
P2
P3
P4
P5
P6
Overall
π0.5 + SFT
13.33
60.00
0.00
0.00
46.67
20.00
23.33
π0.5 + RECAP
0.00
60.00
20.00
13.33
46.67
13.33
25.56
Frozen World Model
0.00
46.67
26.67
13.33
53.33
20.00
26.67
π0.5 + RISE w/ GE-Sim2
13.33
73.33
13.33
13.33
20.00
40.00
28.89
Rollback (Ours)
26.67
60.00
33.33
46.67
53.33
20.00
40.00
Appendix
Table 7: Per-position results on Ball Tossing (success rate, %; 15 trials per position).
Method
P1
P2
P3
P4
P5
P6
P7
P8
P9
Overall
Score
π0.5 + SFT
60.00
86.67
33.33
26.67
60.00
0.00
40.00
80.00
53.33
48.89
0.35
π0.5 + RECAP
60.00
66.67
13.33
26.67
60.00
0.00
60.00
73.33
100.00
51.11
0.44
Frozen World Model
53.33
53.33
13.33
6.67
33.33
46.67
80.00
100.00
86.67
52.59
0.51
π0.5 + RISE w/ GE-Sim2
20.00
86.67
20.00
26.67
40.00
80.00
20.00
80.00
100.00
52.59
0.43
Rollback (Ours)
86.67
80.00
60.00
20.00
20.00
80.00
80.00
80.00
73.33
64.44
0.79
Appendix
Table 8: Per-position results on Tabletop Curling (success rate, %; 15 trials per position).
Figure 9: Search tree on Tabletop Curling with recorded leaf outcomes. Besides branching from the restored context, Rollback rolls back to an intermediate decision point on the real trajectory and explores alternative continuations from there (purple; here color marks where a branch starts, not its outcome).
Figure 10: Sibling continuations on Ball Tossing. Left: ball trajectories of four sibling branches expanded from the same context. Right: close-ups around the bucket for each branch.
Figure 11: Search tree on Block Extraction. (a) Search tree with leaf outcomes, read from the imagined head-camera videos. The blue dashed path is a real attempt and the other paths are imagined continuations; edges are green, red, or purple when the continuations below all succeed, all fail, or are mixed. The ring marks a branch point inside imagination at step 300, whose three children reach 0/3, 2/2, and 1/1 successful continuations. (b) Continuations A and B from this branch point at the same control steps (30 Hz): in A the gripper pulls the red block out of the tower and withdraws with it, while in B it is still in the tower at the end of the horizon.
Figure 12: Open-loop world-model predictions on Block Extraction. Real recording (top) and world-model prediction under the same recorded actions (bottom) for one episode, from the head and both wrist cameras. Steps are control steps at 30 Hz.
Figure 13: Real-world executions. Stroboscopic composites overlay the ball (top) and sliding stones (bottom) over time onto a single frame. Top: ball tossing by the robot and by a human player. Bottom: tabletop curling by human players, with successful and failed deliveries.
Figure 14: Third-person views of Rollback on the real robot. Top: five consecutive extractions without collapsing the tower. Middle: a successful toss, with the ball trajectory overlaid on the final frame. Bottom: two stone deliveries that leave the red stones in the house.
Figure 15: Typical failures on Ball Tossing. Stroboscopic composites of unsuccessful throws sampled from real-robot data collected during training.