DiLA: Disentangled Latent Action World Models

Paper: arXiv:2605.15725v1

1. Motivation (研究动机)

现有 Latent Action Models (LAMs) 的核心问题是论文称为 LAM Trade-off 的张力:如果为了让 latent action 更抽象、更可迁移而施加强 bottleneck(例如 VQ 或 variational bottleneck),模型容易过度简化动作流形并损害视频生成质量;如果放松 bottleneck,生成质量变好,但 latent action 往往混入颜色、纹理、背景等 content 细节,不再是纯 dynamics/action 表征。

Figure 1 解读:这张图给出论文的核心假设:latent action learning 和 disentanglement 不是两个分离目标,而是可以共演化。受限的预测 bottleneck 迫使模型把可预测运动压缩进结构通路,把外观细节卸载到内容通路;反过来,结构/内容分离又让 IDM/FDM 更容易学习高层、content-invariant 的 latent action。

论文要解决的具体目标是:在 不依赖真实 action labels、只用 unlabeled video 的前提下,同时得到 (1) 可跨 embodiment、跨场景迁移的抽象 latent actions,和 (2) 高保真未来视频/特征预测。这个问题值得研究,因为 unlabeled video 远比机器人 action-labeled 数据更易扩展;如果 latent action 可以从视频中自监督抽取并用于 rollout、action transfer、MPC planning,就能把 world model 从“需要显式控制信号的预测器”推进到更通用的自监督交互模型。

2. Idea (核心思想)

DiLA 的核心 insight 是:预测 bottleneck 本身可以成为内容-结构 disentanglement 的驱动力。模型不直接要求 latent action 预测完整视觉 embedding,而是让它只预测 dynamics-relevant 的 structure embedding;因此,高熵外观细节无法通过低维动作瓶颈传递,只能进入 content pathway。

关键创新是一个双通路 latent world model:structure pathway 用 temporal difference、IDM 和 FDM 学习连续 latent action;content pathway 用 Mamba 记忆时间不变的外观/纹理/遮挡背景;Fusion Decoder 再把预测结构、content memory 和初始帧 embedding 重组为下一帧 DINOv2 embedding。与 Genie/LAPA/Moto/AdaWorld/villa-X 等 LAM/VLA 方法相比,DiLA 不把 FDM 当成预训练后可丢弃的中间模块,也不把生成交给外部 diffusion world model,而是在单一框架里同时学习 action extraction、latent rollout 和高保真生成。

3. Method (方法)

Overall framework

DiLA 输入视频 clip ,先用 frozen DINOv2 抽取视觉 embedding ,再经过 spatial-temporal Transformer 建模全局空间依赖和因果时间依赖。随后特征被拆成两条路径:structure pathway 只保留 motion-related spatial layout,content pathway 记忆 temporal-invariant appearance。最后 Fusion Decoder 把预测结构、content memory、初始 embedding 融合成下一帧 embedding ,再由 frozen RAE decoder 可视化。

Figure 2 解读:左侧是 DINOv2 + ST-Transformer 的共享视觉特征抽取;上方 structure pathway 用 Structure Encoder 产生 ,IDM 从 推断 ,FDM 预测下一结构状态;下方 content pathway 用 Content Encoder + Mamba 保存静态外观记忆;右侧 Fusion Decoder 通过双 cross-attention 把 组合,输出用于可视化的目标 DINOv2 embedding。

Structure pathway: IDM/FDM 与预测 bottleneck

Structure Encoder 将每帧 token 压缩为 。IDM 不直接看完整图像,而是处理 temporal difference ,用 3D convolution blocks 抽取空间平移不变特征和时间上下文,生成 的全局 latent action。FDM 是轻量 spatial-temporal Transformer,用 AdaLN-zero 以 latent action 条件化当前 structure,预测 residual displacement: 这里的关键不是简单降维,而是“让低维 action 只负责预测结构变化”。如果 携带颜色/纹理等高熵 content,低维 无法稳定预测它,训练损失会迫使这些 nuisance details 从 structure pathway 中被排出。

Content pathway: Mamba 作为慢特征记忆

Content Encoder 把 token 压缩为 ,Mamba memory 聚合历史 content: 论文把这一路径解释为类似 Slow Feature Analysis:在 POMDP 中,物体外观、背景、纹理相对 world state 是静态的,但由于遮挡和视角变化,它们在 belief state 中会逐渐被揭示。content pathway 专门处理这类 belief update,使 structure pathway 不必为遮挡背景和纹理细节牺牲动作抽象性。

Fusion Decoder 与 latent rollout

Fusion Decoder 是带 dual cross-attention 的 spatial Transformer:以预测结构 为 query,先 attend 到 ,再 attend 到初始视觉 embedding 。论文强调 很重要,因为它补充压缩过程中丢失的高频细节: 推理/rollout 时,DiLA 不在像素空间 autoregressive,而是在 latent structure space 迭代: 得到 后再经 Fusion Decoder 重构 ,并更新 content memory 以供下一步使用。这种 latent rollout 是其相比高维 observation-space rollout 的效率和稳定性来源。

Training objective

训练完全自监督,使用 teacher forcing,仅依赖视频序列;DINOv2 encoder 和 RAE decoder 冻结。总损失包含视觉 latent prediction、structure prediction、latent action consistency 和正则项: 其中 过滤静态帧。cosine 项要求正向/反向 transition 的 latent action 近似相反,以 group symmetry 方式塑造有意义的动作流形;norm 和 variance 正则分别约束流形紧致性、避免信息熵塌缩。

Action transfer、rebinding 与 disentanglement 检验

Figure 3 解读:Action transfer 直接检验 latent action 是否摆脱 source content。图中同一段 latent action 可从人类动作迁移到机器人、从虚拟导航迁移到真实导航,也可以跨物体和视角保持“拿起/移动/转向”等语义。若 编码了 source 外观,这类迁移会产生明显 texture/object leakage。

Figure 4 解读:A 部分 rebinding 把 sequence 的 structure 与 sequence 的 content memory/initial embedding 融合,输出保持 的空间动态但继承 的颜色、纹理、地标;B 部分固定 structure 为 ,只让 content memory 随时间更新,生成结果保持静止,说明 content pathway 没有显著 motion leakage。

Figure 5 解读:去掉 IDM+FDM 后,structure pathway 退化成直接压缩当前帧的 autoencoder, 会保留冗余外观细节;rebinding 时出现来自 structure 序列的 texture leakage。该 ablation 支撑论文主张:真正驱动 disentanglement 的不是双通路命名本身,而是 latent action predictive bottleneck。

伪代码与代码状态

代码搜索未找到开源实现;项目页显示 Code (Coming Soon),Web/GitHub 搜索未发现与论文匹配的 public repository。因此以下伪代码是论文机制级 PyTorch-style pseudocode,不是 released code 映射;也无法建立 github_ref

class DiLATrainStep(nn.Module):
    def __init__(self, dinov2, st_transformer, struct_enc, content_enc, idm, fdm, mamba, fusion_dec):
        super().__init__()
        self.dinov2 = dinov2.eval()          # frozen
        self.st = st_transformer
        self.struct_enc = struct_enc
        self.content_enc = content_enc
        self.idm = idm
        self.fdm = fdm
        self.mamba = mamba
        self.fusion_dec = fusion_dec
 
    def forward(self, video):
        with torch.no_grad():
            e = self.dinov2(video)           # [B, T, P, 768]
        h = self.st(e)                       # causal temporal + spatial attention
        s = self.struct_enc(h)               # [B, T, P, 32]
        c = self.content_enc(h)              # [B, T, P, 256]
        c_mem = self.mamba(c)                # historical content memory
 
        ds = s[:, 1:] - s[:, :-1]
        z = self.idm(ds)                     # [B, T-1, 256]
        s_hat = s[:, :-1] + self.fdm(s[:, :-1], z)
        e_hat = self.fusion_dec(s_hat, e[:, :1].expand_as(e[:, 1:]), c_mem[:, :-1])
 
        z_re = self.idm(s_hat - s[:, :-1])
        z_fwd = self.idm(s[:, 1:] - s[:, :-1])
        z_bwd = self.idm(s[:, :-1] - s[:, 1:])
        moving = ((s[:, 1:] - s[:, :-1]).norm(dim=(-1, -2)) > tau).float()
        inv_temporal = ((F.cosine_similarity(z_fwd, z_bwd, dim=-1) + 1.0) ** 2 * moving).sum() / (moving.sum() + 1e-6)
        reg = z.norm(dim=-1).mean() + inv_temporal + 0.5 * torch.exp(-10.0 * z.std())
 
        loss_e = F.mse_loss(e_hat, e[:, 1:])
        loss_s = F.mse_loss(self.fdm(s[:, :-1], z), ds)
        loss_z = F.mse_loss(z_re, z.detach())
        return 2.0 * loss_e + 0.03 * loss_s + 0.03 * loss_z + 0.001 * reg
def latent_rollout(s0, e0, c0, z_seq, fdm, fusion_dec, mamba_update):
    s_hat, c_mem = s0, c0
    e_preds = []
    for z_t in z_seq.unbind(dim=1):
        s_hat = s_hat + fdm(s_hat, z_t)
        e_hat = fusion_dec(s_hat, e0, c_mem)
        e_preds.append(e_hat)
        c_mem = mamba_update(e_hat, c_mem)
    return torch.stack(e_preds, dim=1)

Code-to-paper mapping:无 public released code,无法给出源文件/类/函数映射;论文级概念对应为 DINOv2 encoder、ST-Transformer、Structure Encoder、IDM、FDM、Content Encoder、Mamba memory、Fusion Decoder、RAE decoder。

4. Experimental Setup (实验设置)

Datasets

  • 训练主数据:Something-Something-v2 (SSv2)、RT-1、RECON、LoopNav;训练使用 16-frame clips。SSv2 原始规模为 220,847 video clips,论文过滤静态片段和快速相机运动以强调清晰物理操作。RT-1 包含超过 130k robot episodes,覆盖 picking、placing、drawer opening 等任务。
  • Latent action analysis:OmniObject3D 合成数据,使用 5,911 objects、216 categories;每个物体绕竖直轴 、每 渲染一次,共 72 views,并保存 segmentation mask 用于 scaling/translation 合成。
  • OOD linear probing:Franka Kitchen、Block Pushing、Push-T、LIBERO Goal;每个数据集采样 800 pairs 训练 probing network,200 pairs 测试。
  • Visual planning: benchmark,包括 Robosuite tabletop pushing 和 RoboDesk tasks;action adaptation 用每环境 100 trajectories 训练 action-to-latent MLP,随后用 Robosuite 5,000 trajectories / RoboDesk 35,000 perturbed scripted trajectories fine-tune。

Baselines 与 metrics

Baselines 包括 LAPA、Moto、AdaWorld(FDM/LAM component)、AdaWorld、villa-X,以及 DiLA ablations:w/o content、w/o IDM+FDM、Discrete 、Gaussian 。视频生成使用 16-frame autoregressive rollout,在 SSv2 和 RT-1 上报告 SSIM 与 LPIPS 。Action transfer 使用 cycle transfer LPIPS increase;latent action quality 用 OOD linear probing MSE ;planning 用 success rate 和 normalized aggregate。

Training config

配置来源说明:以下训练/规划超参数来自论文正文与 appendix 的实验描述;截至本笔记代码搜索时未找到 released code,因此这些数值不是从实际 launch script 或 experiment config 校验得到。

DiLA 约 123M trainable parameters,另有约 500M frozen parameters(DINOv2 base with registers + RAE pretrained ViT-XL decoder)。输入 resize 到 ,DINOv2 embedding 维度 768,patch grid ,clip 长度 16。ST-Transformer hidden 768,spatial depth 4,temporal depth 4,8 heads;Structure Encoder per-patch dim 32、depth 2;

IDM 为 4 个 3D conv blocks,stride ,空间从 ,channels ,global latent action ;FDM hidden 256、depth 4;Content/Mamba dim 256,Mamba depth 2、state dim 32、1D conv kernel 4、expand ratio 2;Fusion Decoder hidden 768、context dim 256、spatial depth 3。

训练硬件为 4× NVIDIA A100 80GB。优化器 AdamW,learning rate ,weight decay ,global batch size 32。训练分两阶段:teacher forcing 30k iterations(约 24 小时),再用 latent rollout fine-tune 1k iterations。

loss coefficients 为 。Planning 里 action MLP 是两层 SiLU,hidden dim 等于 latent action size;

先用 SGD lr=0.01、MSE、batch size 10 训练 3,000 epochs,再替换 IDM 并 fine-tune 1,000 steps,batch size 32,AdamW lr ,只优化

MPPI planning 公式

中,历史上下文长度 ,当前时刻 ,horizon 。MPPI 每轮 采样 条 action sequence,按 cost 加权更新: \begin{split} a_{i,t_0:t_0+H-1} &= \sum_{k=1}^{N} w_i^k\cdot \mu_{i,t_0:t_0+H-1}^k, \\ w_i^k &= \frac{\exp[-\gamma\cdot C(\mu_{i,t_0:t_0+H-1}^k)]}{\sum_{j=1}^{N}\exp[-\gamma\cdot C(\mu_{i,t_0:t_0+H-1}^j)]}. \end{split} 候选序列从上一轮解加 temporally correlated noise 得到: \begin{split} \epsilon_{i,t_0}^k &\sim \mathcal{N}(0,\sigma_{t_0}^2\mathbf{I}), \\ \epsilon_{i,t}^k &= \beta\cdot\epsilon_{i,t-1}^k + (1-\beta)\cdot\tilde{\epsilon}_{i,t}^k. \end{split} Robosuite cost 为 ;RoboDesk cost 为 其中 。论文实验设 、momentum 、temperature ;open slide/open drawer 用 ,其他任务用

5. Experimental Results (实验结果)

Video generation 与 latent-action ablation

ModelSSv2 SSIM ↑SSv2 LPIPS ↓RT-1 SSIM ↑RT-1 LPIPS ↓
LAPA
Moto
AdaWorld(FDM)
AdaWorld
villa-X
DiLA w/o content
DiLA

DiLA 在 RT-1 两项指标上最佳,在 SSv2 LPIPS 上最佳;SSv2 SSIM 略低于 AdaWorld,但 LPIPS 大幅更好,说明 perceptual similarity 与机器人域生成质量是主要优势。w/o content 在 RT-1 LPIPS 仍不错,但整体低于 full DiLA,说明 content pathway 对高保真生成不是附加模块,而是解决 trade-off 的必要结构。

ModelRollouts LPIPS ↓Cycle transfer LPIPS ↓MSE ↓
DiLA w/o content
Discrete
Gaussian
DiLA

Discrete VQ 与 Gaussian prior 都能带来一定 abstraction,但生成质量、cycle transfer 和收敛速度均不如 DiLA。论文的解释是:固定 VQ/KL prior 对 latent space 的形状约束过硬,会破坏动作的连续低维流形;DiLA 的 temporal-difference + symmetry regularization 更贴近“物理 transition 可逆/连续”的几何结构。

Latent action manifold 与 probing

Figure 6 解读:A 显示 translation、scaling、in-plane rotation 分别形成连续 UMAP manifold;B 的 quantitative decoding 验证 latent space 可预测真实 action 参数;C 展示组合动作中 Translation+Scaling 会形成带 scaling topology 的复合簇,Translation+Rotation 则因平移视觉信号更强而部分重叠;D 展示导航场景中 RECON 的 relative yaw 连续谱和 LoopNav 的 forward/left/right 离散簇。

MethodFranka Kitchen ↓Block Pushing ↓Push-T ↓LIBERO Goal ↓
Discrete
Gaussian
DiLA

OOD probing 表明 DiLA 的 与真实 continuous control 更线性相关;这支撑了“连续且有几何正则的 latent action 比离散/高斯 bottleneck 更适合细粒度 transferable action”的结论。

Figure 7 解读:附录可视化展示 OmniObject3D 上单一 transformation 和组合 transformation 的 rollout。它补充说明 UMAP 结构并非只是在 embedding 上可视化好看,而是能映射回实际可生成的平移、缩放、旋转序列。

Figure 8 解读:组合动作 rollout 展示 latent action 可在一定程度上组合 translation、scaling、rotation;这对应主文 Figure 6C 中复合 manifold 的语义解释。

Visual planning

MethodRobosuite push ↑Open slide ↑Blue button ↑Green button ↑Red button ↑Upright block ↑Aggregate ↑
AdaWorld
DiLA

DiLA 在 6 个任务中的 5 个超过 AdaWorld,aggregate 从 21.54 提升到 41.44;唯一落后的 Upright block 差距较小。论文强调 planning fine-tuning 阶段不再做 latent action learning,而是用 action MLP 把真实动作投到 DiLA latent action space,这说明预训练得到的结构化 可以作为 MPC 的高层动作接口。

Additional qualitative results and limitations

Figure 9 解读:RT-1 rollout 的 baseline 对比展示了 DiLA 在较长 16-frame 预测中维持 object layout 和外观一致性的优势;这与 Table 1 中 RT-1 SSIM/LPIPS 最优一致。

Figure 10 解读:附录 action transfer 可视化提供更多跨 embodiment、跨场景例子,说明 latent action 能驱动目标场景中“可承载该动作语义”的实体或视角变化。

作者讨论了三类局限。第一,latent action 的高层抽象会牺牲 fine-grained control precision,单视角视频本身难以恢复精细控制信号;多视角或 proprioception 可能更适合注入细粒度控制。第二,DiLA 当前 disentanglement 主要是 spatial layout vs visual details,并不是 multi-object dynamics 的完整 object-centric decomposition。第三,长程生成仍会受 autoregressive compounding errors 影响。论文还观察到,当目标场景不支持源动作时,transfer rollout 会尽量复现源 dynamics,可能产生物理上不自然的结果,例如把“向前走”转移到前方有墙的场景时让墙逐渐放大/模糊,或把“throwing”动作转移到空 gripper 时把 gripper 本身当成被抛物体。