Rollback:在世界模型中进行经验回溯的策略改进 Rollback: Experience Backtracking in World Models for Policy Improvement
让世界模型可回溯到关键决策点,用反事实分支对比生成动作级监督,稳定改进真机策略
前置知识
动作条件化世界模型
给定观测历史与动作、预测下一观测的神经模拟器。本文采用 GE-Sim 2.0:输入时间上下文 $m_t$、动作块和任务指令 $g$,同时输出三个相机视角的未来帧和经由 pose-expert 头预测的本体状态(每臂 7 个关节 + 1 个夹爪),使策略在想象中收到与真机完全一致的观测。时间上下文 $m_t$ 是显式数据结构,可缓存可恢复,这是后续一切回溯操作的前提。
它是整个方法的载体。不理解世界模型如何被当作训练环境使用、其上下文为何可以被缓存和恢复,就无法理解经验回溯为何可行、又为何比物理环境多出 rewind 这一能力。
MCTS 与 PUCT 选择
蒙特卡洛树搜索为每个节点维护访问计数 $N(v)$ 与累计回报 $W(v)$,经验价值 $Q(v)=W(v)/N(v)$,按 PUCT 分数 $U(v,u)=Q(u)+c_{puct}P(u|v)\sqrt{\frac{1+N(v)}{1+N(u)}}$ 在利用高成功率分支与探索低访问分支之间平衡,rollout 结局沿路径备份($N\leftarrow N+1$, $W\leftarrow W+y$)。AlphaGo 证明它能把有限算力集中到关键决策上。
论文的想象力预算分配机制直接建立在其上;连续动作块场景还用到渐进加宽和虚拟子节点,读懂第 4.1 节必须有此背景。
动作块
现代机器人策略(如 $\pi_{0.5}$)一次输出长度为 $H$、维度为 $d_a$ 的动作序列 $a_t\in\mathbb{R}^{H\times d_a}$ 而非单步动作,在 30Hz 控制频率下连续执行,兼顾时间一致性与推理效率。本文的决策点、分支、价值估计和监督标签全部以动作块为最小单元,树深按块边界计。
论文反复说 chunk-level supervision(块级监督),若不理解动作块是策略的原生输出单位,就无法理解为什么树搜索的分支和相对质量标签天然对齐策略的训练目标。
流匹配与 $\pi_{0.5}$
流匹配是训练连续生成模型的目标:在流时间 $s$ 构造插值 $x_s=(1-s)a+s\epsilon$($\epsilon\sim\mathcal{N}(0,I)$),回归速度场,损失为 $\|v_\theta(x_s,s)-(\epsilon-a)\|^2$。$\pi_{0.5}$ 是 Physical Intelligence 的通才 VLA 策略,用流匹配生成动作块。本文以它为策略骨干,在其 prompt 中额外编码 $K=5$ 档优势条件。
论文公式 (5) 的优势条件化训练和推理时的引导采样都建立在流匹配框架上,不懂它就看不懂策略是如何被树监督更新的。
分类器无关引导 (CFG)
条件生成中把条件预测与无条件预测按 $v_{guided}=v_\emptyset+\beta(v_K-v_\emptyset)$ 混合,用强度 $\beta$ 放大目标条件的效应,最初来自图像扩散模型。本文在树搜索和部署时都条件在最高档位 $b=K$ 并施加 CFG,使策略向自己认为最优质的动作分布坍缩。
它解释了训练时随机丢条件(概率 $p_{drop}$)与推理时强化条件之间的配套关系,是理解策略实际采样行为(进而影响搜索树的扩展分布)的关键。
DMD 分布匹配蒸馏
Distribution Matching Distillation 把多步扩散/流模型蒸馏成少步生成器:约束生成器输出分布匹配教师从噪声到数据映射所诱导的分布。本文每轮世界模型适配为 400 步 LoRA SFT 加 100 步 DMD,训练前也要先做 2500 步 DMD,使 GE-Sim 2.0 以极少采样步数生成动作块级别的想象帧。
视频世界模型逐步采样非常昂贵,没有 DMD 加速,在 48 块 GPU 上跑树搜索的想象预算 $B$ 根本撑不起 AlphaGo 式的大量模拟。
LoRA
低秩适配:冻结主干权重,只训练旁路注入的低秩矩阵增量 $\Delta W=BA$(秩远小于原维度),可训练参数和显存开销小一个量级。本文用它让世界模型在策略训练期间被反复小步微调(每轮 400 步)而无需全量重训,也保证适配节奏能跟上真实轨迹的持续流入。
test-time physics grounding 这一核心能力的工程可行性就来自 LoRA:轻量到可以在自增强循环里无限次重复执行。
GRPO 组相对优势
GRPO(DeepSeekMath 提出)对同一 prompt 采样一组回答,用组内统计量作基线计算相对优势,从而免去独立的价值网络。本文的兄弟归一化 $\rho_i=\hat{q}(v,a_i)/q_{max}(v)$ 再离散成 $K$ 档与之同构,只是分组单位从同一 prompt 的多个响应换成了同一恢复上下文下的多个备选动作块。
作者明确援引 GRPO 来论证无需价值模型的合法性;理解这层对应关系,才能真正领会论文用什么换了 RISE/VLA-RL 里的可学习价值模型。
研究动机
真实环境中做机器人强化学习极其昂贵:每条轨迹都要占用机器人、依赖人工复位,而 GPU 算力可随硬件线性扩展,这种不对称促使研究者把策略训练搬进动作条件化世界模型。然而现有方法(WMPO、World-Gymnast、RISE、WoVR 等)基本沿袭真实环境以轨迹为中心的交互协议:从初始状态向前 rollout,中间决策上下文用完即弃。这带来三个具体问题。其一,探索信息量低:一条轨迹中只有少数关键决策真正决定成败,对绝大多数决策改换动作结果雷同,独立 rollout 把大量算力花在反复生成共享前缀上,识别不出关键决策点也无法向其倾斜预算。其二,决策级监督稀疏:轨迹级二元结果 $y\in\{0,1\}$ 指不出哪个中间动作决定了结局,额外学习价值模型又会引入建模假设和估计偏差。其三,世界模型预测有偏:策略改进后探索到训练分布之外,预测误差累积甚至被优化目标当作可利用的漏洞。实验中最好的前向 rollout 基线 RISE w/ GE-Sim2 三任务平均成功率仅 46.72%,且训练曲线波动明显(投球 2 小时后一度跌至 24.44%)。
本文的目标是本文的目标是把世界模型从被动的环境模拟器升级为反事实引擎,在几乎不增加真实交互的前提下稳定提升真实机器人策略的成功率。具体拆成四点:第一,利用世界模型独有的可回溯性,回退到真实成败轨迹中的关键决策点,从同一历史分支出备选动作块并对比其下游结局,把有限的想象预算集中在信息量最大的比较上;第二,通过子树结局聚合得到细粒度、动作块级的相对质量标签,在不引入任何额外可学习价值模型的前提下解决信用分配问题;第三,构建云-边-端可扩展训练系统,把分布式世界模型搜索与异步真机执行统一起来,让真实经验持续校准世界模型、更准的世界模型再反哺策略改进,形成自增强循环,作者将世界模型侧的这种能力命名为 test-time physics grounding;第四,在三个高难度真机任务(积木抽取、投球、桌面冰壶)上验证:不仅终点成功率显著高于模仿学习、真机 RL 和前向世界模型 RL 基线,而且训练过程更平稳,并要在单台真机约 8 小时训练内让投球和冰壶超过人类玩家的平均成功率。
与已有工作不同的是,本文最独特的切入角度是:所有已有工作都把世界模型当作只能向前走的物理环境代理,而作者注意到世界模型有一个任何物理环境都不具备的能力——回溯再考虑。世界模型的时间上下文 $m_t$ 是显式数据结构,可以缓存并在需要时恢复,因此能在同一个决策点分支出任意多条备选未来,且共享前缀永远不需要重新生成,新分支的开销只与其自身长度成正比,而不是整条轨迹的长度。作者把这个原语命名为经验回溯,并围绕它在训练期重新设计整个策略学习协议。第二个独特角度是把世界模型在服役期间的持续校准提炼为一种模拟器自身的能力:手工物理引擎的动力学是写死的,系统辨识只能调固定模型的参数;而学习型世界模拟器可以直接吸收新的真实转移,在被当作训练环境使用的同时修正自身物理。这与 AlphaGo/Expert Iteration 的 search-and-distill 范式一脉相承——策略提议动作、搜索评估动作、策略再从搜索结果中学习——但搜索首次发生在一个学习出来的视频世界模型内部,且每棵树的根都锚定在真实机器人经验上,从机器人真正失败或成功的情境出发做反事实推演。
核心方法
先讲直觉:机器人第 200 步把冰壶推短而输掉比赛时,人类会想'如果当时推重一点会怎样'。Rollback 就是让世界模型替策略回答这个问题——恢复该决策点的缓存上下文,模拟备选动作的下游结局,再用结局差异反推哪个动作更好。技术路线分四层。第一层搜索:以 GE-Sim 2.0 为世界模型(联合预测三个相机视角的未来帧与 pose-expert 头的本体状态),从示范或真实轨迹构建搜索树并缓存各节点上下文,MCTS 用 PUCT 分数在想象预算 $B$ 内自适应分配算力。第二层监督:对任务微调的 world judge 给每条想象 rollout 打二元结局,子树结局平均得到动作块价值,按最优兄弟归一化并离散成 $K=5$ 档相对质量标签,兄弟同档的组丢弃。第三层策略训练:以 $\pi_{0.5}$ 为骨干,用优势条件化流匹配目标在 $D_{mix}$ 上训练,训练时以概率 $p_{drop}$ 丢条件,推理时条件在最高档并施加引导采样。第四层协同演化:世界模型每轮用新增真实轨迹做 400 步 LoRA SFT 加 100 步 DMD。系统由 48 块 H200(40 搜索+8 训练)和 1 台真机异步驱动,把三任务平均成功率从最强基线的 46.72% 提到 57.93%。
核心创新是经验回溯这一原语。世界模型的时间上下文 $m_t$ 可缓存可恢复,实现 $$m_t \xrightarrow{\text{RESTORE}} (\hat{o}^{(i)}_{t+1}, m^{(i)}_{t+1}) \sim W_\phi(\cdot\mid m_t, a^{(i)}_t, g)$$ 与物理环境只能复位到初始状态不同,恢复的上下文可锚定在真实经验上,策略因此能站在自己真实成败的情境里重新考虑'当时本可以怎么做'。第二是树结构价值估计:动作块 $a_i$ 的价值取其子树叶子结局的成功率 $\hat{q}(v,a_i)=\frac{1}{|L(u_i)|}\sum_{\ell} y_\ell$——兄弟共享同一决策上下文,对比下游结局即隔离单个动作的效应,信用分配无需可学习价值函数;奖励模型只需判断完成的 rollout 是否成功。再按组内最优归一化 $\rho_i=\hat{q}/q_{max}(v)$ 并离散化 $b_i=\min(K,1+\lfloor K\rho_i\rfloor)$,与 GRPO 的组相对优势同构,只是分组单位从同一 prompt 的多个采样换成同一恢复上下文的多个备选动作。第三是把世界模型在线校准提炼为 test-time physics grounding:每条真实转移 $(o_{\le t},a_t,o_{t+1})$ 自带监督目标,LoRA 让小步更新可无限重复,DMD 恢复少步采样,专门修正改进中的策略最常造访而预训练数据覆盖最少的状态。
方法步骤详情
流程九步。①构建决策上下文:从示范或真实轨迹池采样种子,在动作块边界保留观测作参考路径,真实结局沿路径备份——真实尝试即树的第一个被评估分支。②MCTS 选择:按 PUCT 分数 $U(v,u)=Q(u)+c_{puct}P(u|v)\sqrt{(1+N(v))/(1+N(u))}$ 选子节点,先验均匀;连续动作块用渐进加宽限制子节点数,并设虚拟子节点竞争新动作。③策略引导扩展:恢复被选节点上下文,从最高档策略 $\pi_\theta(\cdot\mid o_v,g,b{=}K)$ 以新噪声采样动作块,每个新分支采两个动作。④想象 rollout 至成功或深度上限,world judge 返回 $y\in\{0,1\}$,按 $N\leftarrow N+1$、$W\leftarrow W+y$ 备份。⑤分歧驱动细化:兄弟落入不同档位时把父节点推入优先队列,向访问最少的兄弟追加模拟直至各 $n_{conf}$ 次,压低标签噪声。⑥树结构价值估计产出 $(o_v,g,a_i,b_i)$ 存入 $D_{tree}$。⑦优势条件化策略更新:在 $D_{mix}$ 上最小化 $\mathcal{L}_\pi=\mathbb{E}\|v_\theta(x_s,s\mid o,g,\tilde b)-(\epsilon-a)\|^2$,$x_s=(1-s)a+s\epsilon$;成功真实轨迹标最高档、失败标最低档,以 $p_{drop}$ 丢条件。⑧世界模型协同演化:真实轨迹累积回放,每轮 400 步 LoRA SFT 接 100 步 DMD,发布版本化快照。⑨异步集成:搜索作业钉扎启动时模型版本保证树内可比,边缘站点在回合边界切换权重,经验织物是唯一数据交换点。
技术新颖性
与已有方法的本质区别有四。其一,交互协议:前向 rollout 方法(World-Env、WMPO、World-Gymnast、RISE、WoVR)把中间上下文用完即弃,Rollback 把世界模型显式当作可回溯环境,新增分支成本与分支长度而非整条轨迹长度成正比,把 Go-Explore '先回到旧状态再探索'的思想首次带进学习型视频模型。其二,监督信号:RISE 需学进度价值模型、VLA-RL 需过程奖励模型、RIPT-VLA 用 leave-one-out 优势,而 Rollback 用兄弟结局对比直接产生动作块级相对标签,无需拟合任何价值函数;分歧驱动细化(补采档位不同的兄弟至 $n_{conf}$ 次)独立压低标签噪声。其三,物理接地:首次把'世界模型在服役期间持续吸收真实转移'提炼为模拟器自身的能力 test-time physics grounding,区别于把世界模型更新当训练循环附产品的 World-Gymnast 与 WoVR,也区别于只能做系统辨识的物理引擎——不断改进的策略恰恰漂向训练数据覆盖最少的状态。其四,系统:三进程完全解耦互不等待,版本钉扎保证每棵树内部自洽(同一策略提议、同一世界模型模拟、兄弟结局可比),想象力吞吐随 GPU 数量线性扩展。
实验结果
主实验(Table 1,AGIBOT Genie-02 真机):Rollback 在积木抽取、投球、桌面冰壶分别达 69.33%、40.00%、64.44%,平均 57.93%;较 SFT(40.30%)高 17.63%,较 RECAP(44.22%)高 13.71%,较共享同一世界模型的 RISE w/ GE-Sim2(46.72%)高 11.21%(相对 +24.0%);投球相对提升最大(38.5%),因结局由释放时机决定,兄弟对比恰好直接识别更好的释放。训练动态(Figure 5):8 小时内投球从 25.56% 单调升至 40.00%(相对 +56.5%),RISE 复现版波动、终值仅 28.89%;冻结世界模型只到 26.67%,证明持续物理接地是稳定改进的必要条件;冰壶从 51.11% 升至 64.44%,RISE 仅 +1.48%。消融(Table 2,块抽取):完整方法接触率 89.33%、成功率 69.33%、连续抽取 2.60 块;直接奖励监督使成功率跌至 59.33%,冻结世界模型使接触率跌至 78.67%、连续抽取 1.90。人类对比(Figure 6):投球人类 31.48%、遥操作 16.53%,Rollback 6 小时超人类(33.33%)、8 小时达 40.00%;冰壶人类 58.39%、遥操作 38.94%,4 小时即超人类,终值 64.44%。保真度(Table 3/4):开环回放 PSNR 均值 21.6 dB,RGB L1 从 6.79 增至 11.79(1.74 倍),奖励误报 0/40;协同演化使平均 L1 从 19.52 降到 12.96(-33%),误差增长从 1.90 倍降到 1.38 倍。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 积木抽取(Block Extraction,真实机器人) | 任务成功率 (%) | 69.33 | 最强基线 RISE w/ GE-Sim2:58.67 | 绝对 +10.67 |
| 投球(Ball Tossing,真实机器人) | 任务成功率 (%) | 40.00 | 最强基线 RISE w/ GE-Sim2:28.89 | 绝对 +11.11(相对 +38.5%) |
| 桌面冰壶(Tabletop Curling,真实机器人) | 任务成功率 (%) | 64.44 | 最强基线 RISE w/ GE-Sim2:52.59 | 绝对 +11.85 |
| 三任务平均 | 任务成功率 (%) | 57.93 | SFT 40.30;RECAP 44.22;RISE w/ GE-Sim2 46.72 | 较最强基线 +11.21(相对 +24.0%) |
| 投球 vs 人类 | 任务成功率 (%) | 40.00(6 小时 33.33% 首次超越) | 人类平均 31.48;遥操作 16.53 | 超人类 +8.52,超遥操作 +23.47 |
| 桌面冰壶 vs 人类 | 任务成功率 (%) | 64.44(4 小时即超越) | 人类平均 58.39;遥操作 38.94 | 超人类 +6.05,超遥操作 +25.50 |
| 消融:块抽取 | 接触率/成功率/连续抽取 | 89.33% / 69.33% / 2.60 | 冻结世界模型 78.67/64.67/1.90;直接奖励监督 77.33/59.33/1.73 | 两组件均不可缺少,直接奖励监督伤害最大 |
局限与改进
作者承认的局限:其一,想象保真度有限——开环 PSNR 仅约 21.6 dB,RGB L1 误差在 150–275 帧 rollout 内增长 1.74 倍(6.79 到 11.79),这正是他们从真实观测出发分支而非长时程纯想象的原因;其二,评估范围窄——只有三个任务、一台真机、一个边缘站点,多站点配置论文明言尚未评估;其三,奖励模型需按任务在人工标注视频上微调,world judge 质量直接决定标签质量,虽然 40 条失败回放实测误报 0,这种阈值化判定对未见过的结果模式是否稳健并未验证;其四,人类基线仅 3 人平均。我自己的观察:算力与机器人配比极不对称(48 块 H200 服务 1 台机器人),作者也承认 Policy Trainer 与 Experience Fabric 会成为吞吐瓶颈;冰壶含人类对手,胜负判定引入额外方差;$K=5$、$n_{conf}$、$c_{puct}$、$\alpha$、$\beta$、$p_{drop}$ 等超参敏感性均未消融;方法依赖世界模型保留可恢复的显式时间上下文,对潜空间世界模型如何迁移没有讨论;策略初始化依赖 10 万步 RECAP 式热身,未测试更弱初始化。
独立分析的弱点
第一,算力经济性:1 台真机配 48 块 H200(40 块只做树搜索),每条真实数据摊销的成本极高;改进方向是把搜索搬进潜空间或蒸馏后的轻量世界模型、量化压缩缓存上下文、用置信度门控和早停削减无效 rollout。第二,任务定制成本高:世界模型需任务专属两阶段训练(1 万步 SFT + 2500 步 DMD)、策略需 10 万步 RECAP 式热身、奖励模型需逐任务人工标注微调,缺一不可;改进方向是训练跨任务通用 world judge(结合 VLM 零样本判罚),让新任务只用少量真实轨迹冷启动。第三,长时程想象仍会漂移:L1 误差 1.74 倍增长限制树的深度与纯想象链条长度;改进方向是引入不确定性感知分支——在预测方差高的节点强制回真实校验,或以真实观测为锚周期性回插。第四,理论缺位:兄弟对比得到的相对标签在世界模型有偏时并非无偏优势估计,论文未给出标签偏差随世界模型误差的上界,也未分析异步快照版本错位的影响;改进方向是建立反事实对比加在线接地的偏差传播分析。第五,评测规模:三任务、单机器人、3 人人类基线,结论外推到长时程、多物体、双臂协作或移动操作尚待验证。
未来方向
作者明示的方向:把云-边-端系统扩展到多站点多机器人——论文明确说此类配置未评估,需验证经验织物与编排器在真实多机负载下的吞吐;把 Rollback 用于通才基座策略($\pi$ 系列、OpenVLA)的任务级强化而非单任务专用策略。基于成果可延伸:其一,把经验回溯与课程学习结合,用 MCTS 统计量(节点分歧度、访问分布、成功率)自动定位任务难点并生成针对性种子场景;其二,探索潜空间世界模型中的回溯——缓存隐状态而非视频上下文,把每分支成本再降一个量级,同时需解决隐状态不可解释带来的锚定问题;其三,把 test-time physics grounding 形式化为在线系统辨识问题,与物理引擎混合:引擎提供结构先验、学习模型修正残差;其四,在想象树中引入安全性约束与风险敏感备份,让高危动作在想象中被淘汰而非在真机试错,这对接触密集任务尤为重要;其五,从反事实对比出发建立动作块级优势估计的理论框架,分析世界模型误差向策略梯度传播的速率,据此自适应调节 $n_{conf}$ 与细化预算。
复现评估
复现难度很高。论文只给出项目网站 rollback-wm.github.io,正文与附录均未提及开源代码或权重。关键依赖链:世界模型 GE-Sim 2.0(智元 AgiBot 的模型,未见公开权重)、策略骨干 $\pi_{0.5}$(Physical Intelligence)、逐任务微调的 world judge(需人工标注成败的机器人视频)、AGIBOT Genie-02 真机与遥操作示范数据。算力门槛:世界模型任务微调 1 万步(32 块 H200,全局 batch 128)加 DMD 2500 步(16 块 H200,batch 16);策略热身 10 万步(16 块 H200,batch 64);策略改进阶段 48 块 H200 加 1 台真机约 8 小时。即使复用开源的 $\pi_0$ 系权重与可获取的开源世界模型(如 Cosmos-Predict、V-JEPA 2-AC),仍需自建异步训练系统、经验织物与真机闭环,工程量极大;任务成功判定与人类基线协议只在附录给出一部分。可行的降级路线:在仿真中用小规模视频世界模型复现经验回溯加树结构价值估计的核心思想,先消融 MCTS 分配与分歧细化的贡献,再考虑真机。
论文图表
三范式对比图:(a) 真实环境 RL 依赖昂贵且需要人工复位的物理交互;(b) 前向 rollout 世界模型 RL 把大量交互预算花在低信息量探索上,训练过程不稳定(图中标了 × 与波动曲线);(c) Rollback 通过经验回溯树搜索在关键决策点探索备选未来、从结局差异中学习,同时用真实反馈持续精调世界模型。底部展示三个真实任务:积木抽取、投球、桌面冰壶。
一图看懂论文动机与三大挑战(探索信息量低、监督稀疏、世界模型有偏),是理解'为什么要回溯'的最佳入口。