One-sentence Summary

Researchers from AgiBot, BUAA, and NUS propose Rollback, a policy improvement paradigm that turns world models into counterfactual engines by rolling back to critical decision points, branching into alternative actions, and contrasting their outcomes for action-level supervision without a separate value model; using distributed search and LoRA-based test-time physics grounding, it improves real-world success rates by 10.67 to 11.85 points over the strongest baseline.

Key Contributions

• This work introduces Rollback, a policy improvement paradigm that turns action-conditioned world models into counterfactual engines by rewinding to critical decision points, branching alternative actions, and contrasting outcomes to attribute success to individual decisions. It distills counterfactual outcomes into fine-grained action-level supervision without requiring a separately learned value model.
• The paper presents a unified training system that couples distributed world-model search with asynchronous real-robot execution, forming a self-reinforcing cycle where real experience sharpens the world model and the improved model drives further policy training. It also introduces test-time physics grounding, using a lightweight LoRA pipeline to correct the world model's dynamics from real transitions generated by the policy being trained.
• Experiments on three real-world robot tasks show that Rollback improves success rates over the strongest baseline by absolute 10.67%, 11.11%, and 11.85% on block extraction, ball tossing, and tabletop curling, respectively. Rollback also achieves more stable improvement during training and surpasses average human success on ball tossing and tabletop curling within 8 hours of training.

Introduction

Action-conditioned world models provide simulated environments for robot policy learning, reducing reliance on costly physical trials by letting policies improve from imagined outcomes. However, prior world-model-based policy learning mostly keeps the trajectory-centric interaction protocol of real-world training: it generates forward rollouts and discards intermediate contexts, so it does not exploit a learned world model’s ability to rewind to earlier decision points and explore alternatives. Existing methods also face limited informative exploration, sparse decision-level supervision, and biased world-model predictions when policies branch into less familiar states. The authors propose Rollback, a policy-improvement paradigm that turns world models into rewindable training environments or counterfactual engines, using experience-backtracking tree search, tree-structured advantage estimation without a separate value model, and world model–policy co-evolution with test-time physics grounding for more stable and sample-efficient real-world robot learning.

Method

The authors propose Rollback, a policy improvement paradigm that exploits experience backtracking in learned video world models. As illustrated in the framework diagram, Rollback revisits decision points, compares alternative actions through their downstream outcomes, and distills the results into the policy, while real-world feedback refines the world model.

Experience-Backtracking Tree Search Unlike physical environments, a world model allows its explicit data structure to be cached and restored, enabling many continuations to branch from a single decision point. The authors refer to this operation as experience backtracking. To organize reusable decision contexts, Rollback constructs a search tree. Given a seed trajectory from offline demonstrations or real trajectories, the world model predicts future multi-view frames and the robot's proprioceptive state. Each non-terminal node serves as a branching point where the policy proposes alternative actions and the world model simulates their consequences. The world-model context at each imagined node is cached, allowing further continuations without regenerating the shared prefix.

To concentrate the imagination budget on promising regions, the authors employ Monte Carlo tree search (MCTS). For each node $v$v, they maintain a visit count $N(v)$N(v) and accumulated return $W(v)$W(v), yielding the empirical search value $Q(v) = W(v) / N(v)$Q(v)=W(v)/N(v). An existing child $u$u is selected according to the PUCT score:

where $P(u \mid v)$P(u∣v) is the prior. For continuous action chunks, progressive widening permits new actions. The policy itself acts as the proposal distribution, sampling chunks from the highest-bin policy to explore plausible alternatives. To ensure informative comparisons, the system allocates further rollouts to sibling subtrees with different provisional value bins until a confidence threshold is met.

Tree-Structured Value Estimation The tree structure connects an action to multiple downstream outcomes, enabling chunk-level quality labels without fitting an additional value function. The authors estimate the value of each action chunk $\mathbf{a}_i$ai​ from the outcomes of its downstream continuations using a fine-tuned world judge. The chunk value is estimated as:

where $\mathcal{L}(u_i)$L(ui​) denotes the leaves in the subtree rooted at $u_i$ui​. To convert these estimates into policy supervision, alternative actions from the same decision context are compared. The chunk value is normalized by the best sibling value $q_{\max}(v)$qmax​(v) and discretized into $K$K bins:

The label $b_i \in \{1, \dots, K\}$bi​∈{1,…,K} indicates relative performance. Following RECAP, the authors train the policy conditioned on these discretized advantage labels using a flow-matching objective:

where $\mathcal{D}_{\mathrm{mix}}$Dmix​ combines tree-derived samples with offline demonstrations and robot-collected trajectories.

World Model-Policy Co-Evolution As the policy evolves, its distribution may shift beyond the world model's training data. To prevent the policy from exploiting inaccurate predictions, the authors continually adapt the world model through test-time physics grounding. Newly collected real trajectories are added to a replay set, and the world model is adapted using an action-conditioned prediction objective:

A LoRA adapter keeps each update small, and distribution matching distillation (DMD) is subsequently applied to obtain a few-step generator for efficient imagined interaction.

Asynchronous Real-Imagination Training System Search in a video world model is computationally heavy, while real-robot interaction is slow. To avoid idle resources, Rollback runs real-robot execution, world-model search, policy training, and world-model adaptation as concurrent processes. The infrastructure is divided into three tiers: Cloud, Edge, and Device.

All GPU-intensive work, including search, data storage, and training, runs in the cloud. The Distributed Search Pool hosts policy replicas, world-model instances, and context caches on GPU workers. Each worker builds a tree for a seed trajectory independently, and supervision is sent only after the tree is finished. The Experience Fabric stores real trajectories and tree supervision in separate streams, allowing learners to prefetch batches without waiting for producers. The Co-Evolution Learners consist of the Policy Trainer and the WM Trainer, which optimize their respective objectives continuously using distributed data parallelism. Edge sites run policy inference for the robots, buffering trajectories for upload and caching policy snapshots.

The training process operates entirely asynchronously, exchanging information only through the Experience Fabric and versioned model snapshots.

The robot continually executes the latest policy and uploads labeled trajectories to $\mathcal{D}_{\mathrm{real}}$Dreal​. Meanwhile, GPU workers continually run experience-backtracking search, and the Policy Trainer updates the policy on $\mathcal{D}_{\mathrm{mix}}$Dmix​. In parallel, the WM Trainer repeatedly adapts the world model to accumulated real trajectories. Each search tree is pinned to the latest policy and world-model snapshots available when it starts, ensuring internal consistency. A policy snapshot reaches the robot at its next episode boundary, and each search worker picks up new snapshots at the start of its next tree, creating a continuous loop of improvement.

Experiment

Rollback is evaluated on a real robot across three manipulation tasks, comparing it with imitation, real-world RL, and a world-model baseline that shares the same model. The experiments show that Rollback outperforms all baselines, with experience backtracking and tree-structured value supervision providing the largest gains, while ablations confirm that both world-model adaptation and subtree-aggregated labels are important for stable improvement and task success. Training dynamics and qualitative analyses also show that Rollback improves steadily, surpasses human performance on two tasks by branching from real intermediate states, and benefits from a co-evolved world model that reduces drift and preserves reliable reward labels.

Rollback achieves the highest success rates across all three real-world tasks. Compared with the strongest prior world-model baseline, it improves by roughly eleven percentage points on each task and by about 24 percent in average relative terms. The gain from experience backtracking exceeds the combined gains from real-world reinforcement learning and standard forward world-model rollouts. Rollback reaches the top success rate on Block Extraction, Ball Tossing, and Tabletop Curling, with the largest gains on Ball Tossing and Tabletop Curling. Experience backtracking with the same world model adds more average improvement than real-world reinforcement learning and forward world-model rollouts combined.

An ablation on Block Extraction compares the full framework with a frozen world model and with direct reward supervision replacing tree-structured value estimation. The full framework achieves the best target contact, task success, and consecutive extraction count. Direct reward supervision records the lowest values across all three metrics, indicating that both test-time world-model adaptation and tree-structured value estimation support more reliable extraction. The full framework outperforms both ablations on target contact, task success, and consecutive extractions. Direct reward supervision trails the frozen world model on all reported metrics, with the lowest task success and consecutive extraction count. Both ablations fall below the full framework, showing that test-time physics grounding and subtree-aggregated value estimation each contribute to extraction performance.

On Block Extraction, open-loop rollouts in the co-evolved world model show moderate visual fidelity, with the head camera predicted more accurately than the wrist cameras. Prediction error increases over the rollout horizon, indicating that longer imagined rollouts drift further from reality. Across 40 replayed failed episodes, the reward model produces no false positive success judgments. The head camera view is predicted more accurately than the left and right wrist views. Prediction error grows from the first quarter to the last quarter of the rollout, supporting the use of short imagined branches. The reward model reports zero success events across the 40 replayed failed episodes, avoiding false positives.

The co-evolved world model predicts Tabletop Curling validation episodes more faithfully than the frozen model. It shows higher PSNR and lower RGB L1 error across camera views, and its average error is lower on every episode. Prediction drift also grows more slowly over rollouts, with less spread among repeated rollouts. The co-evolved model improves PSNR and reduces RGB L1 error across head, left wrist, and right wrist views. Error growth from the first quarter to the last quarter is much smaller for the co-evolved model than for the frozen model. Repeated rollouts from the co-evolved model show lower spread, indicating more consistent open-loop predictions.

Rollback is evaluated on three real-world manipulation tasks, where it achieves the highest success rates and shows that experience backtracking contributes more than real-world reinforcement learning and forward world-model rollouts combined. An ablation on Block Extraction indicates that both test-time world-model adaptation and tree-structured value estimation improve target contact, task success, and consecutive extractions, while direct reward supervision performs worst. Analysis of the co-evolved world model shows moderate but drifty visual prediction on Block Extraction with no false positive reward judgments on failed episodes, and on Tabletop Curling it predicts validation episodes more faithfully and consistently than a frozen model.