Under review at ICLR 2027

Rollback

Experience Backtracking in World Models for Policy Improvement
Anonymous Authors
Scroll
+10.44pts over the strongest baseline on Block Extraction
+11.11pts over the strongest baseline on Ball Tossing
+11.85pts over the strongest baseline on Tabletop Curling
> Humansurpasses human players on Ball Tossing and Tabletop Curling
TL;DR

World models can do something the real world cannot: rewind.

Existing world-model-based methods treat world models as forward-only surrogates of the physical environment. Rollback turns them 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. A training system that unifies distributed world-model search with asynchronous real-robot execution closes the loop: real-world experience sharpens the world model, and a sharper world model drives further policy improvement.

From real-world RL to world-model RL to Rollback
From costly real-world interaction, to forward-only world-model rollouts, to experience-backtracking tree search.
Tasks

Three real-world challenges.

Contact-rich manipulation, precise dynamic control, and competitive play, all on a dual-arm AGIBOT robot with OmniPicker grippers.

Block Extraction

Goal
Remove the target block from a stacked tower, one block at a time.
Success
The block is fully extracted and the tower stays standing; a collapse counts as failure.
Tests
Precise, gentle contact under stability constraints.
contact-richlong-horizon

Ball Tossing

Goal
Grasp a ball and throw it into a bucket placed on the floor.
Success
The ball lands and stays inside the bucket; rim bounces and misses count as failure.
Tests
Release velocity and timing; the outcome is decided within a short window.
dynamicprecise control

Tabletop Curling

Goal
The robot and a human opponent each deliver two stones toward the target house.
Success
Scored after all four shots by proximity to the center; later shots may knock earlier stones away.
Tests
Strategic shot selection and precise delivery against an opponent.
competitivemulti-step
Real-robot executions

Three tasks, one physical robot.

Final Rollback policies running on the real robot, recorded from a third-person view.

✓ success3× speed

Block Extraction

Five blocks extracted in a row while the tower stays standing.

69.33% success
✓ success1×

Ball Tossing

A single release decides the outcome; the ball lands in the bucket.

40.00% success
✓ success1×

Tabletop Curling

Two deliveries in sequence, both stones come to rest in the house.

64.44% success
Method

Rewind, branch, compare, improve.

01

Experience-backtracking search

Starting from real trajectories, restore cached world-model contexts and branch into alternative action chunks. MCTS concentrates the imagination budget on decisions whose alternatives diverge.

02

Tree-structured value estimation

Siblings share a decision context, so contrasting their subtree outcomes isolates the effect of each action, yielding action-level labels without a learned value model.

03

World model–policy co-evolution

The robot keeps collecting real experience asynchronously; it refines the world model, which in turn provides more faithful imagination for the next round.

Rollback framework
Overview of Rollback.
Training process
Training process: continual real-robot rollout, rollback search in the world model, and periodic world-model adaptation.
Inside the world model

Same history, different futures.

All clips below are generated by the world model. From one restored context, Rollback branches into sibling actions and rolls each out to its outcome. Comparing siblings tells the policy which decision made the difference.

restored context
Tabletop Curling, t = 0
rollback & branch
✓ branch A
highsubtree value
✗ branch B
lowsubtree value
left: world model · right: real

Is the imagination faithful?

Open-loop world-model prediction versus the real recording under the same action sequence on Block Extraction.

World-model fidelity
Head
23.4 dB
L-wrist
20.6 dB
R-wrist
20.9 dB

PSNR over 40 open-loop rollouts. The reward model flags none of the 40 replayed failures as success.

Results

Best on every task.

Real-world success rates (%). All methods start from the same π0.5 initialization; RISE is reimplemented with our world model for a fair comparison.

BaselinesRollback (ours)

Block Extraction

SFT
48.67
RECAP
56.00
RISE
58.89
Ours
69.33

Ball Tossing

SFT
23.33
RECAP
25.56
RISE
28.89
Ours
40.00

Tabletop Curling

SFT
48.89
RECAP
51.11
RISE
52.59
Ours
64.44
Rollback (ours) RISE Human players Teleoperation

Ball Tossing · real-world success rate during training

Robot vs. human

Surpassing human players.

Three human players perform each task by hand under the same initial positions and success criteria. Although Rollback starts from teleoperated demonstrations far below human level, it surpasses the human average on both tasks.

human player

Ball Tossing

Teleop.
16.53
Human
31.48
Ours
40.00
human player

Tabletop Curling

Teleop.
38.94
Human
58.39
Ours
64.44