Rollback: Experience Backtracking in World Models for Policy Improvement

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.

한 줄 요약

Rollback은 경험 역추적을 통해 월드 모델을 카운터팩트 엔진으로 전환하여 로봇 정책을 개선하는 새로운 학습 패러다임이다.

핵심 기여도

Rollback이라는 새로운 정책 개선 패러다임을 제안, 월드 모델을 단순 시뮬레이터가 아닌 카운터팩트 엔진으로 전환.
경험 역추적 트리 탐색(Monte Carlo Tree Search)을 도입하여 중요한 결정 지점에 집중적 상상 수행.
트리 구조 가치 추정(Tree-structured value estimation)을 통해 별도 가치 모델 없이 세부 행동 수준의 지도 제공.
월드 모델-정책 공동 진화(LoRA 파이프라인)를 통해 실제 경험으로 월드 모델을 지속적으로 정교화.

핵심 아이디어

기존 월드 모델 기반 정책 학습은 주로 미래 상태를 예측하는 단방향 시뮬레이터로 사용되지만, Rollback은 결정 지점으로 돌아가서 대안 행동을 탐색하는 역추적 기법을 도입한다. 이는 물리적 환경에서는 불가능한 기능을 월드 모델에 적용한 것으로, 정책이 특정 결정의 영향을 분석할 수 있도록 한다. Rollback은 세 가지 핵심 구성 요소로 이루어진다: (i) 경험 역추적 트리 탐색, (ii) 트리 구조 가치 추정, (iii) 월드 모델-정책 공동 진화. 특히, 트리 구조 가치 추정은 대안 분기의 결과를 비교하여 개별 행동의 영향을 추정함으로써 별도 가치 모델 없이도 정확한 지도를 제공한다.

기술적 접근법

**Experience-backtracking tree search**: Monte Carlo Tree Search(MCTS)를 사용하여 중요한 결정 지점에 집중적으로 탐색을 수행.
**Tree-structured value estimation**: 대안 분기 간 결과 비교를 통해 세부 행동 수준의 가치 추정을 제공.
**World model–policy co-evolution**: LoRA 파이프라인을 통해 실제 경험으로 월드 모델을 지속적으로 정교화.
**Test-time physics grounding**: 실제 정책이 방문한 상태에서 월드 모델의 예측을 정교화하여 시뮬레이션 신뢰도를 높임.

주요 결과

Block Extraction, Ball Tossing, Tabletop Curling 세 가지 실제 로봇 작업에서 Rollback은 최강 기준 모델 대비 성공률을 각각 10.67%, 11.11%, 11.85% 개선.
8시간의 학습 후 Ball Tossing과 Tabletop Curling에서 인간 평균 성공률을 초과.
Tree-structured value estimation이 없을 경우 Task Success가 59.33% 감소, Frozen World Model일 경우 Consecutive Extractions가 2.60 → 1.90 감소.

의의 및 한계

Rollback은 월드 모델을 단순 시뮬레이터가 아닌 역추적 가능한 학습 환경으로 전환함으로써 정책 개선의 효율성을 높인다. 특히, 별도 가치 모델 없이도 세부 행동 수준의 지도를 제공하는 점에서 기존 접근과 차별화된다. 그러나 월드 모델의 예측 오류가 누적될 수 있으며, 실제 경험 없이 학습된 모델은 정확도가 떨어질 수 있다. 또한, MCTS 기반 탐색은 계산 자원을 많이 소모할 수 있다.

실용적 활용

Rollback은 실제 로봇 작업에서 정책 학습을 효율적으로 수행할 수 있는 기술로, 특히 반복적이고 복잡한 작업(예: 물체 추출, 볼 던지기, 테이블 컬링)에 적용 가능하다. 또한, 인간 수준의 성능을 빠르게 달성할 수 있어 산업용 로봇의 자율성 향상에 기여할 수 있다.