一句话定位
Iso-Dream 把 dreamer-v2 的单一 RSSM 世界模型拆成三条分支——action-conditioned（可控）、action-free（不可控但可预测，如背景车流、水波）、static（时不变背景）——用逆动力学损失强迫可控分支只编码”能被动作解释”的状态变化,再用一个未来状态注意力机制让策略在想象轨迹里”预览”不可控分支未来若干步的演化(例如提前看到别的车会往哪开),在 DMC 噪声背景控制、CARLA 自动驾驶、BAIR/RoboNet 动作条件视频预测四个基准上全面超过 DreamerV2 等基线;NeurIPS 2022 Spotlight。
背景与定位
planet 提出 RSSM(action-conditioned 循环状态空间模型)、dreamer-v2 在其上做想象式 actor-critic 训练,这条 Dreamer 家族路线假设”整个观测序列的时空动态都能被动作信号解释”。但现实场景(尤其自动驾驶)里视觉动态天然是混合的:一部分完全响应动作(自车运动导致的像素变化),另一部分与动作无关但仍然有规律可循(其他车辆的运动、背景水波、光照变化)。把两者硬塞进同一个隐状态会带来两个问题:①噪声背景污染表征,损害非平稳环境下的泛化;②智能体无法显式利用”不可控但可预测”的未来信息做前瞻决策(比如提前预判别的车会撞过来)。
Iso-Dream 的定位是:显式地把 u₁:ₜ(观测背后的时空动态)分解为可控隐状态 s₁:ₜ 和不可控隐状态 z₁:ₜ 两条独立演化的分支,st+1 ~ p(st+1|st,at)、zt+1 ~ p(zt+1|zt),并且根据任务先验决定策略是否要看 z(纯噪声型环境只用 st;类似自动驾驶这种”不可控但与决策相关”的环境则用 st 和 rollout 出来的 zt:t+τ 一起做决策)。论文明确对比的相关工作是 InfoPower(Bharadhwaj et al. 2022,ICLR),区别在于 InfoPower 只是从视觉输入里优先提取”与任务功能相关”的信息,而 Iso-Dream 显式建模可控/不可控两条状态转移分支,并且提出了新的行为学习算法来利用解耦后的不可控预测。
模型架构
三分支世界模型(图 2a),在 planet/dreamer-v2 的 RSSM 基础上新增两条分支:
- Action-conditioned 分支(可控):沿用 PlaNet 式 RSSM,
p(s̃t|s<t,a<t)=p(s̃t|ht),ht=GRUs(ht-1,st-1,at-1),GRU 保存可控动态的历史信息。 - Action-free 分支(不可控):结构与可控分支相同但不输入动作,
p(z̃t|z<t)=p(z̃t|h't),h't=GRUz(h't-1,zt-1)。 - 后验推断:st
q(st|ht,ot)、ztq(zt|h’t,ot),由共享编码器 Encθ 加各分支专属编码器 Encφ1/Encφ2 从当前观测 ot∈R^(3×H×W) 推断。 - Inverse Cell(逆动力学损失,关键设计):2 层 MLP,ãt-1=MLP(st-1,st),用可控分支相邻两步的后验表征去回归真实动作 at-1。这一约束强迫可控分支只学”能被动作解释”的状态转移,是论文的核心解耦手段之一。
- 防塌缩的图像重建:为避免”可控分支学到几乎所有信息、不可控分支学不到东西”的训练塌缩,重建时用先验 s̃t、后验 zt 分别生成可控视觉分量 ôst∈R^(3×H×W)(配 mask Mts∈R^(1×H×W))和不可控视觉分量 ôzt(配 mask Mtz),再叠加从前 K 帧提取的时不变静态背景 ôb=Decφ3(Encθ,φ3(o1:K)):
ôt = Mts⊙ôst + Mtz⊙ôzt + (1-Mts-Mtz)⊙ôb。 - 奖励建模两种模式:不可控状态若被视为纯噪声,奖励只依赖可控状态 p(rt|st);若不可控状态与决策相关(如自动驾驶),奖励用 p(rt|st,zt)。
- 世界模型总损失(Eq.6):图像/奖励/折扣的负对数似然 + α·ℓ2(at,ãt)(逆动力学动作回归损失) + β1·KL[q(st|ht,ot)‖p(st|ht)] + β2·KL[q(zt|h’t,ot)‖p(zt|h’t)]。论文正文明确给出的一组超参(CARLA 实验,4.3 节):α=1,β1=1,β2=1。
- 行为学习:未来状态注意力(Eq.7,图 2b):想象阶段先用 action-free 分支单独滚动出长度 L+τ 的不可控状态序列 {z̃i},在每个想象时间步用注意力把当前可控状态 s̃t 和滑动窗口内的未来不可控状态 z̃t:t+τ 融合成”更有远见”的表征 et:
et = softmax(s̃t·z̃t:t+τᵀ)·z̃t:t+τ + s̃t,τ 为窗口大小(CARLA 实验中 τ=5)。策略 at~π(at|et)、价值模型 vξ(et)≈E[Σγ^(k-t)rk] 沿用 DreamerV2 的 λ-return 训练。 - 部署阶段:与训练一致,若不可控状态与决策相关,交互时先用 action-free 分支滚动出未来 τ-1 步 z̃t+1:t+τ,再用同一个注意力网络融合 st、zt、z̃t+1:t+τ 得到 et 生成动作。
编码器/解码器配置(附录 Table 4,按环境):
- DMC:Encθ conv3-32;可控分支 Encφ1 conv3-64、GRUs 隐藏维度 200、Decφ1 conv3-4;不可控分支 Encφ2 conv3-64、GRUz 隐藏维度 200、Decφ2 conv3-4;静态分支 Encφ3 conv3-64、Decφ3 conv3-3。
- CARLA:与 DMC 相同的通道配置(Encθ conv3-32,两分支 GRU 隐藏维度均为 200)。
- BAIR/RoboNet:Encθ conv3-64;两分支把 GRU 换成 2 层 ST-LSTM,隐藏维度 64(论文正文明确说明”用 ST-LSTM 单元替换两条分支里的 GRU cell”,因为该任务序列更长、需要更强的时空记忆)。
- 附录 Table 4 中 α/β1/β2 三个 KL/动作损失权重在 pdftotext 提取后的表格因合并单元格错位,除正文明确给出的 CARLA 数值(α=β1=β2=1)外,DMC 与 BAIR/RoboNet 逐项取值未能可靠复原,故此处不臆测,标注未披露。
数据
Iso-Dream 是在线 model-based RL / 视频预测训练范式,无预训练语料,“数据”即环境交互产生、存入回放缓冲区的轨迹,或用于视频预测的机器人数据集:
- DMC Generalization Benchmark(video_easy):在原版 DeepMind Control Suite 基础上把背景替换为真实世界视频(制造不可控噪声动态),4 个任务:Walker Walk、Cheetah Run、Finger Spin、Hopper Stand;环境步数上限 500k。
- CARLA(Town04 高速公路场景):第一人称视角,车顶摄像头 60° 视场、64×64 像素图像;30 辆其他车辆作为不可控但可预测的动态源;单个 episode 最长 1000 步;奖励设计 rt = v_ego·û_h·Δt − ξ1·I − ξ2·|steer|(Δt=0.05,ξ1=1e-4,ξ2=1,I 为碰撞冲量、steer∈[-1,1] 为转向惩罚)。
- BAIR robot pushing:Sawyer 机械臂自监督推物体数据集,静态背景,人为叠加弹跳球制造与动作无关但可预测的不可控动态;所有输入 resize 到 64×64。
- RoboNet:4 个实验室(Berkeley、Google、Penn、Stanford)、7 台机械臂的大规模动作条件交互数据集,同样叠加弹跳球。
- 训练/测试的时间跨度设计(BAIR/RoboNet 视频预测实验):训练阶段只学”从 2 帧观测预测未来 10 帧”;测试阶段外推到远超训练时长——BAIR 上用前 2 帧预测后 28 帧,RoboNet 上预测后 18 帧,用于检验长程预测泛化能力。
- 论文未披露 DMC/CARLA 实验的总环境交互步数以外的其他数据规模聚合数字(如具体 episode 数、总帧数),也未披露 BAIR/RoboNet 用于训练的具体轨迹条数。
训练方法
- 两阶段交替循环(Algorithm 1,与 Dreamer 系一致):① 表征学习——从回放缓冲区采样序列,按 Eq.6 计算世界模型损失并更新;② 行为学习——用 action-free 分支单独滚动出长度 L+τ 的不可控状态,在想象轨迹上按 Eq.7/Eq.8 更新策略与价值模型;③ 环境交互——用融合表征 et 生成动作与环境交互,采集新轨迹存回缓冲区。
- 逆动力学训练是解耦的核心机制:用 Inverse Cell 从可控分支相邻状态回归真实动作,损失项 α·ℓ2(at,ãt) 直接加进世界模型总损失联合优化,而不是单独训练一个逆模型。
- 两种世界模型训练模式(按环境特性选择):若不可控信息只用于防止噪声干扰、不参与行为决策(如 DMC video_easy),action-free 分支可以只用重建损失训练,不参与想象与策略部署;若不可控信息与决策相关(如 CARLA),则维持完整 ELBO 目标以保持 z̃t 的语义。
- 行为学习沿用 DreamerV2 的 λ-return 训练 actor 和 critic,区别仅在于 actor/critic 的输入从纯可控状态 s̃t 换成融合了未来不可控预测的注意力表征 et。
- **视频预测任务(BAIR/RoboNet)**的训练目标只有图像重建损失 + Inverse Cell 的动作重建损失,没有奖励/策略学习部分;网络把 GRU 换成 2 层 ST-LSTM 以增强长程时空建模。
- 未发现蒸馏、离散动作 tokenization 或额外的大规模多阶段预训练/微调流程,训练范式即标准 Dreamer 式”世界模型学习 ↔ 行为学习 ↔ 环境交互”闭环外加逆动力学正则。
Infra(训练 / 推理工程)
论文正文与附录均未披露 GPU 型号、GPU 数量、GPU-hours、并行策略或训练精度;也未披露推理 FPS / control-Hz / 延迟等部署侧指标。论文在结论中承认的唯一工程权衡是:相比 DreamerV2,Iso-Dream 因为想象阶段需要额外滚动不可控状态并做注意力融合,“每个 episode 的训练时间更长”(计算效率是作者自陈的局限之一),但同时指出这一开销换来了更高的样本效率(图 4a 显示 Iso-Dream 用更少环境步数达到更高回报)。其余基础设施细节均未披露。
评测 benchmark
DMC video_easy 视觉控制(Table 1,3 个随机种子 × 5 条测试轨迹,均值±标准差;DBC 表示与原论文不同的实验设置)：*
| 任务 | SVEA | CURL | DBC* | DreamerV2 | Iso-Dream |
|---|---|---|---|---|---|
| Walker Walk | 826±65 | 443±206 | 32±7 | 655±47 | 911±50 |
| Cheetah Run | 178±64 | 269±24 | 15±5 | 475±159 | 659±62 |
| Finger Spin | 562±22 | 280±50 | 1±2 | 755±92 | 800±59 |
| Hopper Stand | 6±8 | 451±250 | 5±9 | 260±366 | 746±312 |
Iso-Dream 在全部 4 个任务上超过 DreamerV2 及其余基线(SVEA、CURL、DBC、SAC)。
CARLA 自动驾驶(图 4,3 个随机种子):Iso-Dream 相比 DreamerV2 及其余基线”大幅领先”(论文未给出精确数值表,仅有曲线图);消融实验显示——去掉 Inverse Cell 性能下降,去掉不可控状态滚动(rolling-out)策略性能显著下降,去掉独立的静态分支后性能下降约 15%。
BAIR / RoboNet 动作条件视频预测(Table 2,PSNR↑/SSIM↑,预测 BAIR 未来 28 帧、RoboNet 未来 18 帧)：
| 模型 | BAIR PSNR | BAIR SSIM | RoboNet PSNR | RoboNet SSIM |
|---|---|---|---|---|
| SVG | 18.12 | 0.712 | 19.86 | 0.708 |
| SA-ConvLSTM | 18.28 | 0.677 | 19.30 | 0.638 |
| PhyDNet | 18.91 | 0.743 | 20.89 | 0.727 |
| Iso-Dream | 19.51 | 0.768 | 21.71 | 0.769 |
Iso-Dream 相比 SVG 的 PSNR 提升:BAIR +7.7%、RoboNet +9.3%。
BAIR 消融(Table 3,同一权重分别预测 18/28 帧)：
| 模型 | 18 帧 PSNR | 18 帧 SSIM | 28 帧 PSNR | 28 帧 SSIM |
|---|---|---|---|---|
| Iso-Dream w/o action-free 分支 | 20.47 | 0.795 | 18.51 | 0.690 |
| Iso-Dream w/o Inverse Cell | 21.42 | 0.829 | 19.34 | 0.759 |
| Iso-Dream(完整) | 21.43 | 0.832 | 19.51 | 0.768 |
去掉 action-free 分支的降幅(28 帧预测 PSNR 从 19.51 降到 18.51,约 5.1%)明显大于去掉 Inverse Cell 的降幅(19.51 降到 19.34,约 0.9%),说明”三分支解耦结构本身”比”逆动力学正则”贡献更大,但两者叠加才达到最优。
创新点与影响
贡献:① 提出三分支世界模型(可控 / 不可控 / 静态)显式解耦视觉时空动态,用逆动力学损失约束可控分支只编码动作可解释的变化,这是与此前”content/motion”或”long-term/short-term”式解耦方法(如 PhyDNet)的关键区别——Iso-Dream 按”是否可控”而非按”频率/时间尺度”来切分;② 提出未来状态注意力机制,让策略在想象阶段显式”预览”不可控分支未来 τ 步的演化再决策,把解耦表征直接接入行为学习环节(而不只是表征学习的正则手段);③ 在 DMC 噪声背景、CARLA 自动驾驶、BAIR/RoboNet 视频预测四个基准上全面验证了解耦对下游控制与预测任务的实际收益,尤其在长程决策任务(CARLA)上收益更明显。
影响:Iso-Dream 是把”可控性”作为世界模型解耦准则、并把解耦结果直接喂回策略学习(而非仅用于表征正则)的早期系统性工作,后续作者团队推出了更强版本 Iso-Dream++(仓库:github.com/panmt/MBRL_with_Isolated_Imaginations),延续了这条”隔离并利用不可控动态”的路线。
作者自述局限:① 计算效率——相比 DreamerV2,Iso-Dream 因为想象阶段更密集的状态转移计算,每个 episode 训练耗时更长(但样本效率更高,即用更少环境步数达到更好效果);② 架构需要针对环境专门设计——作者在初步实验中尝试对所有测试基准用同一套模型架构,发现不同基准对网络结构有特定需求(如是否需要独立静态分支、GRU 还是 ST-LSTM),需要依赖对环境的先验知识来定制,不是一套架构包打天下。
原始链接
- arXiv 摘要页:https://arxiv.org/abs/2205.13817
- arXiv PDF:https://arxiv.org/pdf/2205.13817
- GitHub(官方实现):https://github.com/panmt/Iso-Dream
- 项目主页:https://sites.google.com/view/iso-dream
- 后续工作 Iso-Dream++:https://github.com/panmt/MBRL_with_Isolated_Imaginations
一手源存档(sources/)
- iso-dream—github-readme — GitHub README 快照(含训练/评测命令、引用信息、致谢代码来源)
- iso-dream—project-page — 项目主页快照(摘要、作者单位、NeurIPS 2022 Spotlight 标注)
- arXiv 原文 PDF(2205.13817v3,arXiv 原文 PDF,不入 git)——见上方 arXiv 链接