Going back in a world model
Most world‑model‑based methods only run the world model forward, just like the real world. A world model can also go back. Rollback returns to earlier decision points of real and simulated attempts, tries other actions from there, and compares how they end. The comparison tells which actions were better, without training a value model. Real‑world data is used along the way to keep correcting the world model.
Three real‑world games
All games are played by the dual‑arm AGIBOT Genie‑02 robot. The videos show the final Rollback policies.
Jenga Game
- Goal
- Push the target block out of the tower, one block at a time, without knocking the tower over.
Ball Tossing Game
- Goal
- Pick up the ball and throw it into a bucket on the floor. The ball has to stay in the bucket.
Tabletop Curling Game
- Goal
- The robot plays second against a human, and each side slides two stones. After the four shots, the side with the stone closest to the center wins. Later shots can knock earlier stones away.
How Rollback works
The world model, GE‑Sim 2.0, predicts the three camera views and the robot's joint states, so the policy gets the same kind of input in simulation as in the real world.
Experience‑backtracking search
Starting from real trajectories, the search restores saved world‑model states and tries other action chunks. MCTS spends more of the budget on decisions where different actions lead to different outcomes.
Tree‑structured value estimation
Sibling branches start from the same state, so comparing their outcomes shows how good each action is. These comparisons become action‑level labels; no value model is trained.
World model–policy co‑evolution
The robot keeps collecting data while the policy trains. The new data is used to fine‑tune the world model, so later searches run on a more accurate simulator.
A search tree from the Tabletop Curling Game
The blue path is a game recorded in the real world, where the robot plays second. The violet branches are rollouts in the world model, started from earlier states of that episode. Each node is one action chunk of 50 control steps at 30 Hz. Hover over a node to play its video; otherwise the tree replays the search step by step.
Frozen vs. co‑evolved world model
We take real‑world Tabletop Curling Games that the robot won (the robot plays second) and replay their recorded actions in two world models: the initial one, kept frozen, and the one updated together with the policy during training. Both start from the same frame and get the same actions.
Averaged over 40 open‑loop rollouts.
Best on all 3 games
Real‑world success rates (%). All methods start from the same π0.5 policy. RISE w/ GE‑Sim2 is our reimplementation of RISE on the same world model, so the gap comes from backtracking search and tree‑structured supervision.
Surpassing human players
Three people played each game by hand, judged by the same success criteria as the robot. Rollback starts from a policy trained on teleoperation data, which is well below human level, and ends above the human average in both games. Each video shows a human player first, then the robot with Rollback.