PaperlayerView exact revision on arXiv

Source-linked AI summary

Rollback: Experience Backtracking in World Models for Policy Improvement

Jiayi Luo, Yang Zhang, Liliang Chen, Nan Wang, Chen Gao, Boxiang Qiu, Zhuoxuan Li, Baorui Peng, Chuankang Li, Yuxiang Yan, Si Liu, Qingyun Sun, Shuicheng Yan, Maoqing Yao, Jianxin Li, Guanghui Ren

arXiv:2610.36068v1cs.ROcs.AIcs.LG

TL;DR

Forward-only world-model learning leaves counterfactual exploration of past decision points unused despite the cost of physical robot interaction. Rollback rewinds to those decisions, branches into alternative actions, and compares outcomes to provide action-level supervision while real experience refines the model. Across three real-world tasks, it improves success over RISE w/ GE-Sim2 and improves more steadily.

  • Problem

    Existing world-model policy learning largely relies on forward rollouts and discards intermediate contexts, leaving past decisions unavailable for counterfactual exploration.

    Source: §1 · p. 1

  • Method

    Rollback restores decision contexts, searches alternative action continuations, aggregates their outcomes into action-level supervision, and refines the world model with real transitions.

    Source: §1 · p. 2

  • Results

    Rollback improves success over RISE w/ GE-Sim2 by 10.67%, 11.11%, and 11.85% on block extraction, ball tossing, and tabletop curling, respectively, with more stable improvement.

    Source: §1 · p. 3, §6.2 · p. 11

  • Takeaways & Limitations

    Rollback concentrates imagination on informative decision comparisons and supplies chunk-level supervision without a learned value model.

    Source: §7 · p. 15

    Rollback relies on a task-specific reward model and world-model fidelity over long contact-rich horizons.

    Source: G Limitations and Future Work · p. 24

Abstract

from arXiv · show

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.

1 Introduction

World models make policy learning less dependent on costly physical trials, but forward-only rollouts leave their ability to revisit decisions and explore alternatives unused. Rollback uses experience backtracking to compare counterfactual actions, produce action-level supervision, and refine the world model with real experience; experiments report higher success and more stable improvement.

1 Introduction

  • Existing world-model policy learning largely uses forward rollouts and discards intermediate contexts, leaving it unable to rewind and reconsider decisions as Rollback does.A few critical decisions can determine outcomes, while trajectory-level feedback does not identify which intermediate actions mattered.Source: §1 · pp. 1–2
  • Rollback turns world models into counterfactual engines by restoring decision contexts, branching into alternative action chunks, and comparing their outcomes.Monte Carlo tree search allocates simulation budget toward decisions whose alternatives lead to divergent outcomes.Source: §1 · p. 2
  • Subtree-level outcome aggregation provides fine-grained, action-level advantage supervision without an additional learned value model.Contrasting sibling-branch outcomes isolates the effect of individual actions.Source: §1 · p. 2
  • A scalable training system couples distributed world-model search with asynchronous robot execution, using real transitions to refine predictions in states the policy visits.This addresses the limitation of offline-trained fixed simulators, which may have seen least the states reached by an improving policy.Source: §1 · pp. 2–3
  • Rollback improves success over the strongest baseline by 10.67%, 11.11%, and 11.85% on block extraction, ball tossing, and tabletop curling, respectively.It also achieves faster, more consistent improvement under matched interaction budgets and surpasses average human success on ball tossing and tabletop curling within 8 hours.Source: §1 · p. 3

2 Related Work

Related work studies robot learning through physical interaction, action-conditioned world models, and policy improvement using imagined experience. Rollback builds on these strands by revisiting decision contexts and grounding world-model predictions with real transitions.

Robot Policy Learning from Interaction

  • Physical-robot reinforcement learning uses execution outcomes for policy improvement, but remains resource-intensive because interactions are costly and resets require labor.Approaches include alternating reinforcement and supervised learning, reward models, sparse-success learning, parallelized sampling, and expert corrections (Guo et al., 2025b; Lu et al., 2025; Tan et al., 2025; Li et al., 2025a; Physical Intelligence et al., 2025a).Source: §2 · p. 3

World Models for Robotics

World-Model-Based Policy Improvement

Physics Grounding of World Simulators

Tree Search for Policy Improvement

  • Tree-search methods allocate simulation budgets across branches and can distill search-improved decisions, but prior approaches leave video-world-model context restoration and counterfactual branching unexplored for real-robot policy improvement (Schrittwieser et al., 2019; Hubert et al., 2021).Go-Explore demonstrates the value of returning to visited states before further exploration.Source: §2 · p. 4

3 Preliminaries

The policy receives language instructions, multi-view images, and proprioception, and generates action chunks for world-model interaction. Because the world model retains an explicit temporal context, it can restore real-experience decision points and branch imagined continuations without regenerating their shared prefixes.

3.1 Problem Formulation

  • The policy maps a language instruction and multi-view images with proprioception to an action chunk, and the world model predicts the next observation and updates its temporal context.Real episodes end in binary success or failure, and the objective is to maximize expected task success while minimizing costly physical interaction.Source: §3.1 · p. 4

3.2 World Models as Rewindable Environments

  • Experience backtracking restores a cached world-model context anchored to a recorded observation, allowing alternative continuations from the same real decision point.The shared prefix is not regenerated, so the cost of a new continuation depends on its own length rather than the full trajectory length.Source: §3.2 · p. 4

4 Proposed Rollback

Rollback revisits decision contexts in a world model, compares alternative actions through their downstream outcomes, and uses the resulting supervision to improve the policy while real trajectories refine the model.

4.1 Experience-Backtracking Tree Search

  • Cached contexts let MCTS branch from selected decisions, reuse shared prefixes, and direct limited imagination toward promising or under-visited alternatives.Policy-guided expansion samples plausible action chunks; disagreement-driven refinement adds rollouts to sibling comparisons before they become training signals.Source: §4.1 · pp. 5–6

4.2 Tree-Structured Value Estimation

  • Subtree outcomes estimate chunk quality, and relative comparisons among sibling actions provide policy labels without fitting an additional value function.Sibling values are normalized against the best explored alternative and discretized; groups with no preference signal are discarded.Source: §4.2 · pp. 6–7

4.3 World Model–Policy Co-Evolution

  • Real trajectories continually adapt the world model as the policy changes, using SFT on action-conditioned transitions followed by DMD for few-step generation.The prediction loss covers multi-view frames and proprioceptive state; recent-policy data targets states and actions likely to appear in upcoming search.Source: §4.3 · p. 7
  • Test-time physics grounding uses lightweight LoRA updates from real transitions to address policy-driven shifts beyond the world model’s training data, a capability specific to learned simulators.The paper notes that search can exploit inaccurate predictions and that hand-specified physics engines lack this learned-simulator adaptation capability.Source: §4.3 · p. 7

5 Asynchronous Real–Imagination Training System

Rollback’s real-robot execution, world-model search, and training run concurrently, exchanging data and versioned snapshots rather than waiting on one another. Its cloud–edge setup, cached search contexts, and ongoing adaptation connect real trajectories with policy supervision from search.

5 Asynchronous Real–Imagination Training System

  • Real execution, search, policy training, and world-model adaptation run concurrently, with each process advancing independently and exchanging data or versioned snapshots.Real episodes, search trees, and learner updates proceed at different rates; none waits for another.Source: §5 · p. 7, §5.3 · p. 9

5.1 Cloud–Edge–Device Infrastructure

  • Cloud workers handle search, storage, and training, while edge sites run policy inference and robots execute returned action chunks.The Experience Fabric separates real trajectories from tree supervision; the trainers use these streams for world-model adaptation and policy updates, respectively.Source: §5.1 · p. 8

5.2 Search Scheduling and Context Reuse

  • Cached world-model contexts let branches resume from selected nodes, while version-pinned jobs keep each tree’s policy and world model fixed until completion.Context reuse avoids recomputing shared prefixes; DMD reduces sampling steps, though the Policy Trainer and Experience Fabric eventually limit throughput as search GPUs are added.Source: §5.2 · p. 9

5.3 Training Process

  • The robot uploads labeled real trajectories as the policy trainer uses search-generated tree supervision, while the world model is repeatedly adapted to accumulated real data.Snapshots reach robots between episodes and search workers at the start of their next tree; experiments used one physical robot, and additional sites were not evaluated.Source: §5.3 · pp. 9–10, Figure 4 · p. 9, §5.1 · p. 9

6 Experiments

Rollback is evaluated on three real-world robot tasks, with experiments covering baseline performance, component ablations, training dynamics, human comparisons, and world-model fidelity. The results show gains over baselines and human averages on two tasks, alongside evidence about the effects of model adaptation and tree-structured supervision.

6.1 Experimental Setup

  • The evaluation covers contact-rich block extraction, precise ball tossing, and competitive tabletop curling, using one AGIBOT Genie-02 robot.The setup reports six initial positions per task, nine for curling, and repeated trials at each position.Source: §6.1 · p. 10

6.2 Main Results

  • 10.67%, 11.11%, and 11.85% absolute success-rate gains over RISE w/ GE-Sim2 on block extraction, ball tossing, and tabletop curling, respectively, accompany Rollback’s highest success rate on all three tasks.Rollback reaches 69.33% on Block Extraction, 40.00% on Ball Tossing, and 64.44% on Tabletop Curling.Source: §6.2 · p. 11, Table 1 · p. 11

6.3 Ablation Studies

  • Removing either world-model adaptation or tree-structured value estimation degrades all three Block Extraction metrics; direct reward supervision lowers task success to 59.33%.Freezing the world model reduces consecutive extractions from 2.60 to 1.90; target contact falls from 89.33% to 78.67% with a frozen model and to 77.33% with direct reward supervision.Source: §6.3 · p. 11, Table 2 · p. 11

6.4 Training Dynamics

  • Over eight hours, Rollback rises from 25.56% to 40.00% on Ball Tossing and from 51.11% to 64.44% on Tabletop Curling, while RISE w/ GE-Sim2 ends at 28.89% and 52.59%, respectively.With a frozen world model, Ball Tossing reaches only 26.67%.Source: §6.4 · p. 12, Figure 5 · p. 12

6.5 Qualitative Results

  • On Tabletop Curling, rollback comparisons show that changing the second push can determine whether the red stone reaches the house, while imagined branches distinguish fast from slow releases.The illustrated real attempt and alternative share the context up to the second push; another pair of imagined continuations branches at step 50.Source: §6.5 · p. 12, Figure 7 · p. 13

6.6 Surpassing Human Performance

  • Rollback exceeds human-average success within eight hours on Ball Tossing and Tabletop Curling, reaching 40.00% and 64.44%, respectively.The human averages are 31.48% and 58.39%; Rollback first exceeds them after six hours on Ball Tossing and four hours on Tabletop Curling.Source: §6.6 · p. 12, Figure 6 · p. 12

6.7 World-Model Fidelity and Reward Reliability

  • On 40 open-loop Block Extraction rollouts, RGB L1 error grows from 6.8 in the first quarter to 11.8 in the last, while the reward model reports success on none of the replays of failed episodes.The rollouts replay eight failed real episodes five times each; the error increase is 1.74× over 150–275 frames.Source: §6.7 · pp. 12, 14, Table 3 · p. 14
  • On 64 Tabletop Curling validation episodes, co-evolution lowers average RGB L1 error from 19.52 to 12.96 and reduces rollout spread from 9.39 to 6.12.The error rises from 12.48 to 23.73 across the frozen-model rollout, versus 10.50 to 14.47 for the co-evolved model; the figure plots error over control steps with 95% confidence intervals.Source: §6.7 · p. 14, Table 4 · p. 14, Figure 8 · p. 14

7 Conclusion

Rollback uses experience backtracking and counterfactual outcome aggregation to produce chunk-level policy supervision without a learned value model. It improves success rates and training stability over competing baselines, including surpassing human averages on two tasks after eight hours.

7 Conclusion

  • Rollback improves success rates over RISE w/ GE-Sim2 by 10.67%, 11.11%, and 11.85% on block extraction, ball tossing, and tabletop curling, and surpasses human averages on the latter two after eight hours.The conclusion also reports more stable improvement than competing baselines.Source: §7 · p. 15

A Detailed Training Algorithm

Rollback runs real-robot execution, tree search, policy training, and world-model adaptation as concurrent processes that share data and versioned model snapshots. Search restores cached contexts to branch actions and distills tree outcomes into policy supervision.

A Detailed Training Algorithm

  • Four concurrent processes handle real-robot execution, search, policy training, and world-model adaptation, exchanging data through shared real and tree buffers and models through versioned snapshots.Each search tree is pinned to the snapshots current when it starts.Source: A Detailed Training Algorithm · p. 18
  • Each search worker pins policy and world-model snapshots, samples a demonstration or real trajectory, runs backtracking tree search, and appends extracted labels to the tree buffer.The policy trainer samples from tree, demonstration, and real-data buffers; the world-model trainer adapts on current real data and distills the adapted model.Source: A Detailed Training Algorithm · p. 18
  • Every expansion restores the cached state before sampling an action chunk, so sibling branches share the same decision context without regenerating their common prefix.The highest-bin policy is rolled out to success or the trajectory horizon while intermediate nodes are cached.Source: A Detailed Training Algorithm · p. 18
  • Search counts simulated action chunks against a soft budget, completes initiated rollouts and sibling pairs, and restricts later widening to mixed-outcome groups.The first imagined sibling pair is always admitted; later widening is capped at a task-specific maximum width.Source: A Detailed Training Algorithm · p. 18

B Experiment Details

The experiment details specify a shared real-robot evaluation interface and a search procedure that branches from real trajectories and turns imagined outcomes into relative action supervision. Hyperparameters are shared across tasks except where task-specific values are given.

B Experiment Details

  • Policies use three RGB camera views and robot proprioception, with a 16-dimensional deployed action and a 30 Hz low-level control loop.Block extraction and ball tossing use 50-step predicted action chunks; compared policies share observation processing, action interface, and controller settings.Source: B Experiment Details · pp. 18, 20
  • Rollback builds a reference spine from a labeled real trajectory, then uses PUCT with progressive widening to expand imagined continuations within an imagination budget.Successful seed trajectories are searched to their spine depth, while failed trajectories may be searched to a maximum depth.Source: B Experiment Details · p. 19
  • At branching points, the procedure estimates each action’s success rate from its leaf outcomes and bins it relative to the best observed action.Groups with no successful continuation or no difference in relative bins are omitted from tree-derived supervision.Source: B Experiment Details · p. 19
  • The evaluation resets policies from a shared pool of task-specific configurations, randomizes object configurations, and applies task-specific trial limits and termination rules.Block extraction is capped at 700 control steps and ball tossing at 120; curling ends after each player takes two shots.Source: B Experiment Details · p. 20
  • Real-robot outcomes receive binary task-specific success labels, while imagined rollouts are judged by GE-Sim 2.0, which does not determine reported real-world success rates.Block extraction requires removal with the tower intact, and ball tossing requires landing in the bucket.Source: B Experiment Details · p. 20

C Hyperparameters

  • Hyperparameter values are task-specific for block extraction, ball tossing, and tabletop curling, while all remaining values are shared across tasks.The appendix identifies Table 5 as the hyperparameter listing.Source: C Hyperparameters · p. 20, Table 5 · p. 20

D Per-Position Results

The appendix reports per-position results for all three tasks, with trial counts and additional task-specific measures. The supplied captions identify the tables’ scope but do not provide their result values.

D Per-Position Results

  • Tables 6–8 report success rates by initial position, using 25 trials per position for block extraction and 15 for ball tossing and tabletop curling.Block extraction also reports blocks extracted in a row before failure, while curling reports average score against the human opponent.Source: D Per-Position Results · p. 20, Table 6 · p. 21, Table 7 · p. 21
  • Figure 9 depicts a tabletop-curling search tree with recorded leaf outcomes and branches rolled back from the real trajectory to probe a late decision point.Its caption marks the purple branches by where branching starts, not by outcome.Source: Figure 9 · p. 21
  • Figure 10 shows four ball-tossing sibling branches expanded from the same context, with separate trajectory views and close-ups around the bucket.The caption specifies the visual content, not which branch performs best.Source: Figure 10 · p. 22

E Additional Qualitative Results

The qualitative results illustrate Rollback’s counterfactual branching, its learned action-level distinctions, and examples of real-world execution. They also show one open-loop world-model prediction alongside the recording it follows.

E Additional Qualitative Results

  • Across 12 imagined block-extraction continuations rolled back from a real attempt, 9 extract the red block while leaving the tower intact.Figure 11’s search tree also includes a real attempt and branches within imagination.Source: E Additional Qualitative Results · p. 22, Figure 11 · p. 22
  • At block-extraction step 300, three sampled action chunks from the same imagined state yield subtrees with 0/3, 2/2, and 1/1 successful continuations.The contrasted continuations distinguish a chunk that extracts the red block by step 500 from one that leaves it in the tower through step 650; subtree values of 0 versus 1 provide the relative training signal.Source: Figure 11 · p. 22, E Additional Qualitative Results · p. 23
  • On one block-extraction episode, the world model reproduces the arm approach and gripper closure over 300 control steps across three camera views, with the tower remaining intact as in the recording.Figure 12 compares the real recording and prediction under the same recorded actions.Source: E Additional Qualitative Results · p. 23, Figure 12 · p. 23
  • Figure 13 depicts successful and failed robot tosses and human curling deliveries, with curling successes bringing red stones into the house and displacing opponent stones.Failed human deliveries leave stones outside the scoring area.Source: E Additional Qualitative Results · pp. 23–24, Figure 13 · p. 23
  • Figure 14 shows five consecutive block extractions, one successful toss, and two curling deliveries that leave red stones in the house.These are examples of final-policy real-robot executions.Source: E Additional Qualitative Results · p. 24, Figure 14 · p. 24

F Failure Modes

Ball-tossing failures are concentrated around the release, where errors in velocity, direction, or gripper timing determine the throw’s outcome.

F Failure Modes

  • Short throws and overshoots arise from release-velocity errors, while rim bounces arise from small release-direction errors; these three types account for most failures.Late gripper opening causes less frequent failed releases that drop the ball near the robot.Source: F Failure Modes · p. 24, Figure 15 · p. 24
  • All four observed failure types are decided within a short window around release, so a single action chunk largely determines the episode outcome.The four types include short throws, overshoots, rim bounces, and failed releases.Source: F Failure Modes · p. 24

G Limitations and Future Work

Rollback depends on task-specific reward models and accurate world-model predictions over long contact-rich horizons. The reported policies are limited to single-task training with one physical robot.

G Limitations and Future Work

  • Rollback depends on task-specific reward models and world-model fidelity over long contact-rich horizons.Prediction error grows along a rollout.Source: G Limitations and Future Work · p. 24
  • Prediction error grows along a rollout, constraining the fidelity of imagined outcomes over extended prediction horizons.Source: G Limitations and Future Work · p. 24
  • Each policy in this work is trained for a single task with one physical robot.Future work includes multi-task policies and scaling to multiple robots sharing one search cluster.Source: G Limitations and Future Work · p. 25

Related papers

Loading 2610.36068v1…