Rollback: Experience Backtracking in World Models for Policy Improvement

Rollback. Experience Backtracking in World Models for Policy Improvement. By Jiayi Luo and fifteen others, from AgiBot, Beihang University, and one other institution. Published on the twenty-sixth of October.

Action-conditioned world models offer a scalable alternative to costly real-world interaction when improving robot policies. Existing methods built on them, though, largely treat a world model as a forward-only surrogate for the physical environment. That leaves untapped a capability no physical environment affords: rewinding to a past decision point and exploring a counterfactual future from there. A real robot cannot do that. Once an action has been taken the episode moves on, and when it ends in failure there is no going back to see which of the many decisions along the way was responsible. A world model has no such restriction, yet using it forward-only treats it as though it did.

Rollback is a policy improvement paradigm built on exactly that capability. It turns the world model from a passive simulator into a counterfactual engine. It rolls back to critical decision points, branches into alternative actions, and contrasts their outcomes, so that success can be attributed to individual decisions. The contrast is what does the attributing. When two futures split at one decision point and only one of them succeeds, the difference in outcome belongs to the choice made there.

Imagination is concentrated where it matters most, and the counterfactual outcomes are distilled into fine-grained, action-level supervision, without a separately learned value model.

Around this sits a scalable training system that unifies distributed world-model search with asynchronous real-robot execution. The two halves form a self-reinforcing cycle. Real-world experience continually sharpens the world model, and a sharper world model in turn drives further policy improvement.

The authors call the world-model side of that 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 very policy it trains correct its dynamics, through a lightweight LoRA pipeline. The environment and the learner improve each other. A hand-built simulator stays as accurate as it was built to be, while this one is updated by the robot's own experience as training goes on.

Three challenging real-world tasks put it to the test: block extraction, ball tossing and tabletop curling. Over the strongest baseline, Rollback improves success rates by an absolute ten point six seven percent on block extraction, eleven point one one percent on ball tossing, and eleven point eight five percent on tabletop curling, while improving more stably throughout training.

The starting point was modest. Rollback was initialized from teleoperated demonstrations far below human level. Within eight hours of training, it surpasses the average success rate of human players on both ball tossing and tabletop curling.

What separates a world model from the world, in this account, is the ability to rewind. Rollback makes that difference the whole method, returning to the moments that decided an outcome and asking what else could have happened there. And because real transitions keep correcting the model's dynamics, the futures it imagines stay tied to the physics the robot actually meets.