Rollback: experience backtracking in world models for policy improvement

What this paper shows and why it matters

According to the authors, action-conditioned world models offer a scalable alternative to costly physical interactions when training robot policies (arxiv.org). However, existing methods treat world models only as forward-only simulators, missing a unique capability: the ability to rewind to previous decision points to explore alternative futures (arxiv.org). The authors introduce Rollback, a policy improvement paradigm that turns world models into counterfactual engines that roll back to critical decision points, test alternative actions, and contrast outcomes to supervise individual decisions (arxiv.org). This method 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 (arxiv.org).

The problem

The authors state that existing world-model-based policy learning encounters three main challenges (arxiv.org). First, limited informative exploration occurs because independent rollouts cannot identify consequential decisions or adaptively allocate exploration among them (arxiv.org). Second, sparse decision-level supervision creates a credit-assignment problem, while learning an additional value model introduces extra modeling assumptions and potential estimation bias (arxiv.org). Third, biased world-model predictions can accumulate errors when branching into states outside the world model's training distribution (arxiv.org).

The idea

To address these challenges, the authors present Rollback, which comprises three core components (arxiv.org). First, an experience-backtracking tree search rolls back to intermediate decision points from successful and failed real-world trajectories, using Monte Carlo tree search to allocate simulation budgets (arxiv.org). Second, tree-structured value estimation contrasts sibling branches sharing a decision context to isolate the effect of individual actions, yielding fine-grained advantage supervision without a separately learned value model (arxiv.org). Third, world model and policy co-evolution unifies distributed world-model search with asynchronous real-robot execution, forming a self-reinforcing cycle where real-world experience continually sharpens the world model through a lightweight LoRA pipeline in a process termed test-time physics grounding (arxiv.org).

The results

The authors report experiments on three challenging real-world robotics tasks: block extraction, ball tossing, and tabletop curling (arxiv.org). Rollback improves success rates over the strongest baseline by an absolute 10.67% on block extraction, 11.11% on ball tossing, and 11.85% on tabletop curling, while achieving more stable improvement throughout training (arxiv.org). Furthermore, 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 (arxiv.org).

Why it matters

By turning passive world models into active counterfactual engines that can rewind, branch, and learn from alternative actions without a separate value network, Rollback offers a scalable way to improve robot policies (arxiv.org). The integration of asynchronous real-world execution with distributed world-model search provides a self-reinforcing loop that corrects simulator dynamics during training (arxiv.org). This helps bridge the gap between simulated imagination and real-world execution performance (arxiv.org).

Caveats

This is a preprint from arXiv, meaning it has not yet undergone peer review (arxiv.org). The authors note limitations inherent to world-model-based methods, such as prediction errors potentially accumulating when branching from policy-selected contexts beyond the model's initial training distribution (arxiv.org).