我们提出DriftWorld，一种基于漂移的动作条件世界模型，能够在单次前向传播中生成高质量的未来帧预测。这种单步行为使得模型在生成轨迹时比现有的主要基于多步扩散的动作条件世界模型快得多。
问题定义
动作条件世界模型定义为一个函数，它根据当前观测和动作预测未来观测，并通过漂移目标进行学习。
动作条件架构
为了一次性生成多个未来帧，U-Net使用分解的时空卷积：空间卷积独立且相同地应用于每个T时间步，然后时间卷积独立且相同地应用于每个空间位置。
世界建模的视觉评估
在所有四个环境中，DriftWorld在视觉生成质量上匹配或优于动作条件世界模型基线，同时推理时间仅为其一小部分。
Robomimic. 在单视图和双视图设置中，DriftWorld与Ctrl-World和GPC扩散世界模型相比取得了竞争性能，并且速度显著更快。
DriftWorld中的推理时策略改进
与现有的扩散世界模型相比，DriftWorld能够实现更快、更有效的推理时策略改进。我们通过获取预训练策略、在DriftWorld中 rollout 其动作提议，并选择导致最佳未来状态的动作来评估这一能力。
具体而言，我们采用GPC-RANK方法。每一步，基础策略采样K个候选动作提议。每个提议在想象中通过世界模型进行 rollout。然后，奖励函数对预测的未来观测进行评分，选择具有最高奖励的动作块用于实际执行。
DriftWorld作为策略评估的离线模拟器
DriftWorld作为一个准确的离线模拟器用于策略评估，成功预测了策略的绝对性能和相对排名。我们在Robomimic和Push-T任务上验证了这一能力。
Robomimic. 我们进一步在Robomimic任务上验证DriftWorld。为此，我们在额外包含失败案例的视频上进行了后训练。我们在世界模型中 rollout 了九个扩散策略，并记录了50个环境初始化种子下的平均成功率。DriftWorld中的模拟成功率与真实MuJoCo模拟器中的成功率密切相关，相关系数分别为0.9916和0.9250，优于基线Ctrl-World模型。
讨论
局限性。 首先，DriftWorld依赖鲁棒的预训练特征提取器（DINOv2/v3）来保持复杂场景中的视觉清晰度。其次，漂移框架在训练时需要比标准扩散更高的内存使用，因为模型每次前向传播需要生成多个负样本（例如64个）来计算漂移损失。因此，每个负样本的上下文和生成帧数量受限于GPU显存。未来的一个方向是增加上下文和生成窗口的长度以改善长期时间一致性，可能通过使用稀疏历史或具有时间压缩的视频VAE来实现。
结论。 我们介绍了DriftWorld，一种快速、单步、动作条件的世界模型。为了将漂移生成模型适应于条件视频生成，我们提出了关键改进，包括动作条件的U-Net架构、特征空间集成和加强的动作跟随。DriftWorld平均比最先进的扩散基线实现了17倍的加速，并生成准确且视觉一致的 rollout，从而实现了改进的推理时动作搜索和可靠的离线策略评估。
Bridge-V2、RT-1和Language Table上生成视频的比较
DriftWorld在Language Table、Bridge-V2和RT-1数据集上生成的视频与真实情况非常相似。
DriftWorld策略rollout的比较
DriftWorld中的策略rollout与真实rollout非常相似。
单帧与分块级模拟
DriftWorld在单次前向传播中同时支持单帧和分块级模拟。对于Push-T，生成4个未来帧效果最佳。
漂移损失的运动加权
运动加权技术对于DriftWorld在具有复杂背景的真实机器人数据集上的自回归生成至关重要。没有运动加权，生成的机器人夹爪基本保持静止。有了运动加权，生成的夹爪会根据动作移动。
DriftWorld中推理时策略改进的细节
提供了通过GPC-RANK进行推理时策略改进的更多细节。
DriftWorld中策略评估的细节
评估的策略。 对于Push-T，我们取第50、100到300 epoch的六个扩散策略检查点，以及一个更长时间训练的epoch 650检查点。对于Robomimic Lift任务，我们取第2、4、6到18 epoch的九个扩散策略检查点。对于Robomimic Can任务，我们取第50、100到400 epoch以及epoch 75的九个扩散策略检查点。
在失败演示上的后训练。 对于Robomimic Lift和Can任务，DriftWorld和基线最初在多人类数据集上训练，该数据集仅包含成功演示。由于这些初始模型对策略成功率提供了过于乐观的估计，我们在包含大量失败演示的数据集上对模型进行后训练，这些失败演示是通过早期扩散策略检查点的 rollout 创建的。这种后训练使模型能够准确模拟失败。
常见问题
DriftWorld相比基于扩散的世界模型有多快？ DriftWorld通过单次前向传播生成未来帧，而不是多次迭代去噪步骤，平均比最先进的扩散基线实现了17倍的加速。
DriftWorld的架构在动作条件视频生成方面有何独特之处？ DriftWorld使用分解的时空U-Net，结合漂移损失函数、特征空间集成和运动加权，实现单步条件视频生成。
DriftWorld能否用于离线策略评估？ 可以，DriftWorld准确预测了策略的绝对性能和相对排名，在Robomimic任务上与真实模拟器的相关系数分别达到0.9916和0.9250。
DriftWorld方法的主要局限性是什么？ DriftWorld需要鲁棒的预训练特征提取器来保证视觉质量，并且在训练时由于每次前向传播需要生成多个负样本，内存使用较高。