MotionGPT3 - Human Motion as a Second Modality

论文 arXiv | GitHub

研究背景

文本（Text）是离散符号，动作（Motion）是连续信号，两者性质截然不同。现有方法面临两个核心矛盾：

• 量化误差（Quantization Error）：将动作离散化为 VQ 码本索引以适配 LLM 的 next-token prediction，会引入近似误差，衰减高频细节，破坏语义-物理一致性
• 跨模态干扰（Cross-Modal Interference）：在单流（Single-Stream）骨架中混合文本与动作 token，梯度相互拉扯，导致训练不稳定、收敛慢

研究目标

提出一个统一的动作-语言模型，同时支持：

• Text-to-Motion（T2M）：文本描述 → 生成动作
• Motion-to-Text（M2T）：动作序列 → 生成文本描述

且需要避免量化误差，减少跨模态干扰。

核心概念

Mixture-of-Transformers（MoT）

来自 Liang et al. (2024) 的思想：为每个模态配备独立的 Transformer 分支，各自拥有独立的 Embedding、FFN 和 LayerNorm，仅在 Self-Attention 层共享。这样：

• 各模态保留自身的归纳偏置（Inductive Bias）
• 跨模态信息仅通过 Attention 受控交换
• 新增模态只需新增分支，无需重训全部参数

连续动作表示（Continuous Motion Representation）

使用预训练的 Motion VAE 将动作编码为连续 latent 向量（而非 VQ 离散索引），避免量化误差。

研究方法

整体架构

1
2
3
4
5
6

                Motion VAE                    Bimodal LLM                   Motion VAE
              ┌───────────┐            ┌────────────────────┐           ┌───────────┐
原始动作 m ───→│ Encoder E  │──→ z₀     │ Text Branch T      │           │ Decoder D │──→ 动作
              │ 9层Trans.  │           │ Motion Branch M    │           │ 9层Trans. │
              │ 4头,skip   │           │ + Diffusion Head H │──→ ẑ₀ ──→│ 4头,skip  │
              └───────────┘            └────────────────────┘           └───────────┘

1. Motion VAE

• 来自 Xin et al. (2023)（MLD）
• Encoder $\mathcal{E}$: 将 $N$ 帧动作 $m^{1:N}$ 编码为单个 latent $z \in \mathbb{R}^{256}$
• Decoder $\mathcal{D}$: 将 latent $z$ 解码回动作序列
• 训练目标：重建损失 + KL 正则
• 预训练后冻住，不参与后续训练
• 动作长度信息隐式编码在 latent 中，Decoder 可生成变长输出

2. 双分支 Bimodal Transformer（MoT 架构）

两个并行分支，基于 GPT-2 配置（12 层，维度 768，MLP 维度 3072）：

路由机制：硬路由（非学习），由特殊标记决定——

• <som> / <eom> 界定动作边界
• <motion_in> / <motion_out> 标记 I/O 位置
• 文本 token（$\vartheta_i = 0$）→ Text Branch
• 动作 token（$\vartheta_i = 1$）→ Motion Branch

3. 动作接口模块

由于动作是连续信号，不能复用文本的 Embedding lookup 和 softmax 解码，需要专用接口：

• MUH（Motion Understanding Head）：线性投影，将 VAE latent 映射到 Transformer 输入空间（理解任务）
• MGH（Motion Generation Head）：即 Diffusion Head $\mathcal{H}$，将 Transformer 隐状态映射回 VAE latent 空间（生成任务）

4. Diffusion Head $\mathcal{H}$

轻量级扩散模型（3 层 MLP + ResBlock，隐藏维度 1024），在 VAE latent 空间中做去噪：
$$\mathcal{L}{\text{diff}} = \mathbb{E}{z_0, t, \epsilon}\left[|\epsilon - \mathcal{H}(z_t, h_m)|_2^2\right]$$

• 训练时：对 $z_0$ 加噪得 $z_t$，$\mathcal{H}$ 以 Motion Branch 隐状态 $h_m$ 为条件，学习预测噪声 $\epsilon$
• 推理时：插入 $K$ 个 <motion_out> 占位 token，Motion Branch 输出 $h_m^{v:v+K}$，$\mathcal{H}$ 从纯噪声 $z_T$ 逐步去噪（默认 100 步）得到 $\hat{z}_0$，再由 VAE Decoder 解码

训练策略：三阶段对齐

Stage I: Text-to-Motion 预训练

• 冻住 Text Branch $\mathcal{T}$
• 只训练 Motion Branch $\mathcal{M}$ + Diffusion Head $\mathcal{H}$
• 任务：T2M（文本→动作生成）
• 目的：让 $\mathcal{M}$ 在稳定的语言条件下学会动作语义
• 100k iterations

Stage II: 跨模态对齐（Cross-Modal Alignment）

• 仍冻住 $\mathcal{T}$
• 引入多任务：T2M + M2T + Motion Prediction
• 以指令格式呈现，促进双向对齐
• 300k iterations

Stage III: 联合微调（Joint Fine-Tuning）

• 解冻 $\mathcal{T}$，全参数微调
• 混合文本-动作配对数据 + 纯文本数据
• 50k iterations

主要发现

双流 vs 单流

• 双流架构收敛速度约为单流的 2 倍（训练损失）
• 验证指标（R@3, MMDist）收敛快 4 倍
• 相同损失水平下，双流模型质量更优

连续 VAE vs 离散 VQ

• VQ 方案在 R@3 约 0.5 时即饱和（天花板低）
• VAE 连续表示持续改进，最终质量显著更高

实验结果

在 HumanML3D 上：

• T2M：R@3 = 0.837，MMDist = 2.725，达到 SOTA
• M2T：BertScore = 35.231，超越现有统一模型

讨论

优势

• 避免 VQ 量化误差，保留高频运动细节
• 双分支设计减少梯度干扰，加速收敛
• 三阶段训练抑制跨任务干扰
• 仅需 2 张 3090，训练高效

局限性

• VAE 输出单个 latent，不支持长动作的分段组合生成
• 方向性控制（左/右）有时会失败
• 泛化能力受限于数据覆盖范围

与其他多模态架构的对比

未来方向

1. 分层/分段 latent：用 hierarchical 或 segment-wise latent 表示支持长动作和组合生成
2. 更大数据集和更强 LLM：扩展训练规模，评估效率和鲁棒性
3. 局部语义对齐：支持段级别的文本-动作精细对应

参考文献

• Zhu, B., Jiang, B., Wang, S., et al. (2025). MotionGPT3: Human Motion as a Second Modality. arXiv:2506.24086.
• Xin, T., et al. (2023). MLD: Motion Latent Diffusion.
• Liang, C., et al. (2024). Mixture-of-Transformers (MoT).
• Radford, A., et al. (2019). GPT-2.

3D-Scene
Atlas
CLIP
CV
Chemistry
Contrastive-Learning
DINO
DT
Diffusion
DiffusionModel
Embodied-AI
FL
FPN
FoundationModel
Gated-NN
HRI
Hierarchical
HumanoidRobot
Image-Grounding
Image-Text
Image-generation
Image2Text
ImgGen
ImitationLearning
LLM
LatentAction
ML
MoE
MoT
MotionGeneration
MR/AR
Message-Passing
Multi-modal
Multi-view
MultiModal
NLP
NN
Object-Detection
Open-Vocabulary
Panoptic
Physical-Scene
PoseEstimation
QML
Quantum
RL
RNN
Real2Sim
Reconstruct
Representation-Learning
RobotLearning
Robotics
Scalability
Scene-graph
Scene-synthesis
Segmentation
Semantic
Sim2Real
Subgraph
Survey
Task-Planning
Transformer
Translation-Embedding
VAE
VLA
VLM
VLP
VQ-VAE
ViT
Visual-Relation
WorldModel
Unified-Multimodal

BAGEL: Emerging Properties in Unified Multimodal Pretraining

论文链接 | 项目主页

作者：Chaorui Deng, Deyao Zhu, Kunchang Li 等 (ByteDance Seed)

研究背景

统一多模态理解与生成（Unified Multimodal Understanding and Generation）是当前AI领域的热点方向。GPT-4o、Gemini 2.0等闭源系统展现了强大能力，但开源模型与之仍存在显著差距。现有开源统一模型主要在图文配对数据上训练，缺乏对复杂多模态交错数据（Interleaved Data）的利用。

研究目标

1. 缩小开源统一多模态模型与闭源系统（GPT-4o、Gemini 2.0）之间的性能差距
2. 解决现有模型架构中理解与生成模块之间的信息瓶颈（Bottleneck）问题
3. 构建能够支持复杂多模态推理的大规模交错数据

核心概念

理解与生成之间的瓶颈（Bottleneck）

在采用 External Diffuser 架构的模型中，LLM/VLM 与扩散模型通过轻量级适配器连接：

• 语言模型生成少量潜在token作为”语义条件”
• 这些token被传递给扩散模块生成图像
• 问题：LLM的丰富上下文被压缩到少量token中，导致信息损失，尤其影响长上下文多模态推理

Mixture-of-Transformer-Experts (MoT)

与传统MoE不同，MoT复制整个Transformer层而非仅FFN：

• 理解专家：处理文本和ViT token
• 生成专家：处理VAE token
• 两个专家通过共享自注意力在每层交互

研究方法

架构设计

BAGEL采用无瓶颈的集成Transformer方案：

双视觉编码器：

• 理解编码器：SigLIP2-so400m/14，捕获语义信息
• 生成编码器：FLUX VAE，处理像素级信息

训练范式

损失权重比：$\text{CE} : \text{MSE} = 0.25 : 1$

广义因果注意力（Generalized Causal Attention）

对于交错多图像生成：

• Noised VAE tokens：用于Rectified-Flow训练
• Clean VAE tokens：作为后续生成的条件
• ViT tokens：统一输入格式，提升交错生成质量

采用Diffusion Forcing策略，对不同图像添加独立噪声级别。

数据构���

数据规模

交错数据构建流程

视频数据：

1. 视频预处理（分割、裁剪、质量过滤）
2. 使用蒸馏的小型VLM生成帧间描述
3. 构建时序对齐的交错序列

网页数据：

1. 两阶段主题筛选（LLM + fastText）
2. 质量过滤（分辨率、清晰度、相关性）
3. Caption-first策略：为每张图像生成描述并插入其前

推理增强数据（Reasoning-Augmented Data）

受DeepSeek-R1启发，构建50万条推理增强样本：

• Text-to-Image生成
• 自由形式图像操作
• 概念性编辑

主要发现

涌现能力（Emerging Properties）

论文定义：某能力在早期训练阶段不存在，但在后期训练中出现

不同能力的涌现时间点（达到85%峰值性能所需token数）：

关键发现：

• 理解和生成能力最先收敛
• 编辑能力随后涌现
• 需要复杂推理的智能编辑能力最后涌现
• ViT tokens对智能编辑至关重要（移除后性能下降16%）

架构对比实验

在1.5B模型上对比Dense、MoE、MoT三种架构：

• MoT在生成任务上优势最明显
• 表明理解和生成可能需要不同的参数空间

实验结果

多模态理解（7B参数）

图像生成（GenEval）

智能编辑（IntelligentBench）

讨论

优势

• 无瓶颈架构：理解与生成模块间无损信息交互
• 涌现能力：首次系统揭示统一多模态预训练的涌现规律
• 开源贡献：发布代码、模型权重和数据构建协议
• 推理增强：CoT显著提升复杂任务表现（WISE: +0.18, IntelligentBench: +10.4）

局限性

• 与GPT-4o在智能编辑上仍有差距（55.3 vs 78.9）
• 模型规模相对较小（7B active / 14B total）
• 训练计算成本高（需要大规模交错数据）

相关工作

统一多模态模型：

• Janus-Pro：采用离散视觉tokenizer的自回归方法
• MetaQuery-XL：冻结预训练VLM backbone
• Transfusion：统一AR和扩散的早期探索

视觉生成：

• FLUX.1-dev：当前SOTA扩散模型
• SD3-Medium：Stable Diffusion系列

未来方向

1. 更大规模训练：探索更大模型和更多数据下的涌现行为
2. 视频生成：论文展示了初步的视频生成能力，有待深入
3. 强化学习：无瓶颈架构为多模态RL提供了基础
4. 世界建模：导航、3D操作等世界建模任务的进一步探索

参考文献

• Deng et al. (2025). Emerging Properties in Unified Multimodal Pretraining. arXiv:2505.14683
• DeepSeek-AI (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
• Esser et al. (2024). Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (SD3)

LingBot-VLA: A Pragmatic VLA Foundation Model

论文链接 | GitHub | Checkpoints

研究背景

视觉-语言-动作（Vision-Language-Action, VLA）基础模型是机器人操作领域的新兴方法，通过大规模预训练使机器人能够执行由自然语言指令引导的多样化操作任务。然而，目前存在以下问题：

• 缺乏关于真实机器人性能如何随预训练数据规模增长而变化的系统性实证研究
• 缺乏高效的训练代码库来支持大规模数据的扩展评估
• 缺乏跨多平台、多任务的系统性真实世界评估基准

研究目标

1. 探索 VLA 模型在真实世界机器人数据上的扩展规律（Scaling Law）
2. 建立跨多平台、多任务的系统性真实世界评估基准
3. 开发高效的大规模 VLA 训练代码库

核心概念

Mixture-of-Transformers (MoT) 架构

将预训练的视觉语言模型（VLM）与动作生成模块（Action Expert）结合，通过共享自注意力机制实现跨模态统一建模。视觉-语言和动作模态通过独立的 Transformer 路径处理，既保留 VLM 的语义先验，又避免跨模态干扰。

Flow Matching

一种用于连续动作建模的生成方法，通过学习从噪声到目标动作的向量场来生成平滑的机器人控制信号。

Blockwise Causal Attention

将序列划分为图像-指令块、状态块和动作块，应用因果掩码防止信息泄露，确保动作预测只能访问当前和历史观测信息。

研究方法

模型架构

LingBot-VLA 采用 MoT 架构，整合 Qwen2.5-VL 作为视觉语言骨干网络，配合独立的 Action Expert 模块：

联合建模序列：
$$[O_t, A_t] = [I_t^1, I_t^2, I_t^3, T_t, s_t, a_t, a_{t+1}, \ldots, a_{t+T-1}]$$

其中 $I_t^{1,2,3}$ 为三视角图像，$T_t$ 为任务指令，$s_t$ 为机器人状态，$A_t$ 为动作序列（chunk length = 50）。
类似[[BAGEL-Unified-Multimodal-Pretraining]]

Flow Matching 目标函数

定义概率路径通过线性插值：
$$A_{t,s} = sA_t + (1-s)\epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$

训练目标：
$$\mathcal{L}{FM} = \mathbb{E}{s \sim U[0,1], A_t, \epsilon}|v_\theta(A_{t,s}, O_t, s) - (A_t - \epsilon)|^2$$

深度信息蒸馏

通过可学习查询 $Q_t$ 与 LingBot-Depth 模型的深度 token $D_t$ 对齐，增强空间感知：
$$\mathcal{L}{distill} = \mathbb{E}{Q_t}|Proj(Q_t) - D_t|$$

训练效率优化

• FSDP 分布式策略：采用混合分片数据并行（HSDP），为 Action Expert 模块构建专用分片组
• 算子级优化：使用 FlexAttention 优化稀疏注意力计算，torch.compile 进行算子融合
• 混合精度：reduction 使用 float32，存储和通信使用 bfloat16

主要发现

扩展规律验证

• 预训练数据从 3,000 小时扩展到 20,000 小时，下游任务成功率持续显著提升
• 在 20,000 小时数据量下仍未出现饱和迹象，表明 VLA 性能持续受益于数据量增加
• 首次提供了真实世界机器人学习中有利扩展特性的实证证据

数据效率

• 仅使用 80 个演示即可超越 π0.5 使用 130 个演示的性能
• 随着后训练数据量增加，与基线的性能差距进一步扩大

实验设计

预训练数据

• 规模：约 20,000 小时真实世界操作数据
• 来源：9 种双臂机器人平台（AgiBot G1、AgileX、Galaxea R1Lite/R1Pro、Realman Rs-02、Leju KUAVO、Qinglong、ARX Lift2、Bimanual Franka）

评估基准

• GM-100 基准：100 个操作任务，39,000 个专家演示
• 评估规模：3 个机器人平台，每任务 130 个后训练 episode，共 22,500 次试验
• 对比方法：π0.5、GR00T N1.6、WALL-OSS

真实世界评估结果

仿真评估结果（RoboTwin 2.0）

训练吞吐量

• 实现 261 samples/s/GPU（8-GPU 配置）
• 相比 StarVLA、DexBotic、OpenPI 提升 1.5~2.8 倍
• 在 256 GPU 规模下仍保持接近线性扩展

讨论

优势

• 首次在大规模真实世界数据上验证 VLA 扩展规律
• 显著优于现有 SOTA 方法的多平台泛化能力
• 高效的训练代码库，支持大规模分布式训练
• 开源代码、模型和基准数据

局限性

• 目前仅支持双臂机器人配置
• 评估主要集中在桌面操作任务
• 深度信息蒸馏依赖额外的 LingBot-Depth 模型

相关工作

Foundation VLA

• π0：Vision-language-action flow model for general robot control
• π0.5：VLA model with open-world generalization
• GR00T N1.6：Open foundation model for generalist humanoid robots

Spatial VLA

• SpatialVLA：探索 VLA 模型的空间表示
• Spatial Forcing：通过对齐策略增强 VLA 空间理解
• GeoVLA：赋能 VLA 模型 3D 表示能力

高效训练框架

• OpenPI：支持 JAX 和 PyTorch 的 π 系列模型训练框架
• StarVLA：VLA 和 VLM 联合训练的模块化代码库
• DexBotic：统一高效��� VLA 开发生命周期解决方案

未来方向

1. 扩展机器人类型：整合单臂和移动机器人数据，支持更多样化的操作能力
2. 非约束环境：探索在非结构化环境中的移动操作能力
3. 持续扩展：进一步扩大预训练数据规模，探索扩展规律的上限

参考文献

• Black et al. (2025). π0: A vision-language-action flow model for general robot control. RSS.
• Black et al. (2025). π0.5: A vision-language-action model with open-world generalization. CoRL.
• Bjorck et al. (2025). GR00T N1: An open foundation model for generalist humanoid robots. arXiv.
• Bai et al. (2025). Qwen2.5-VL technical report. arXiv.
• Lipman et al. (2022). Flow matching for generative modeling. arXiv.
• Wang et al. (2026). The Great March 100: 100 detail-oriented tasks for evaluating embodied AI agents.

arXiv:2503.07137

作者：Siyuan Mu (四川农业大学), Sen Lin (休斯顿大学)

研究背景

随着AI基础大模型的快速发展，现代数据集变得越来越多样化和复杂，包含多模态数据（文本、图像、音频）和复杂结构（图、层次关系）。这给大模型发展带来两大挑战：

1. 计算资源消耗巨大：训练和部署大模型的计算成本呈指数增长
2. 异构数据拟��困难：在单一模型中整合冲突或异构知识变得困难，导致训练不稳定和性能次优

混合专家模型（Mixture of Experts, MoE）通过动态选择和激活最相关的子模型来处理输入数据，成为解决这些挑战的有效方案。

研究目标

• 填补现有MoE综述的空白（过时或缺乏关键领域讨论）
• 全面总结MoE的基础设计、算法、理论和应用四大关键组件
• 为研究者提供系统性参考，激发进一步研究

核心概念

MoE基本原理

MoE采用”分而治之”（divide and conquer）策略，与传统密集模型不同：

• 传统模型：对每个输入激活所有参数
• MoE模型：根据输入特征动态选择和激活最相关的参数子集

MoE层数学表示

$$
\text{MoE}(x) = \sum_{i \in \mathcal{I}_D} w_i M_i(x)
$$
其中 $\mathcal{I}_D$ 是被选中专家的索引集，$w_i$ 是第 $i$ 个专家的权重，$M_i(x)$ 是专家网络输出。

研究方法

1. 门控函数（Gating Function）

线性门控（Softmax Gating）

$$
G(x)i = \text{softmax}(\text{TopK}(g(x) + R{noise}, k))i
$$
其中 $g(x)$ 是线性函数计算的门控值，$R{noise}$ 是鼓励专家探索的噪声。

非线性门控

• 余弦门控（GMoE）：
  $$
  G(x) = \text{TopK}\left(\text{softmax}\left(\frac{E^T W_{linear} x}{\tau |W_{linear} x| |E|}\right)\right)
  $$
• 指数族分布门控
• Soft MoE：使用加权平均而非离散分配

2. 专家网络（Expert Network）

3. 路由策略（Routing Strategy）

• Token级路由：基于token表示进行路由决策（最经典）
• 模态级路由：根据数据模态进行路由（多模态任务）
• 任务级路由：根据任务ID确定路由（多任务学习）

4. 训练策略

负载均衡损失（Switch Transformer）

$$
\mathcal{L}{aux} = \alpha \cdot N \cdot \sum{i=1}^{N} f_i \cdot Q_i
$$
其中 $f_i$ 是分配给专家 $i$ 的token比例，$Q_i$ 是路由概率比例。

MoA（Mixture-of-Attention）详解

基本架构

MoA将MoE机制引入多头注意力模块，每个注意力头视为一个”专家”。

工作流程

1. 输入token进入MoA层
2. 门控网络计算每个注意力头的重要性分数
3. 选择TopK个最相关的注意力头
4. 仅计算被选中头的输出并加权求和

代码实现核心

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

class MixtureOfAttention(nn.Module):
    def __init__(self, d_model, num_heads=8, head_dim=64, top_k=2):
        self.attention_experts = nn.ModuleList([
            AttentionExpert(d_model, head_dim) for _ in range(num_heads)
        ])
        self.router = AttentionRouter(d_model, num_heads, top_k)

    def forward(self, x):
        weights, indices, logits = self.router(x)
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            for head_idx in range(self.num_heads):
                mask_k = (indices[:, :, k] == head_idx)
                if mask_k.any():
                    head_output = self.attention_experts[head_idx](x)
                    output[mask_k] += weights[:, :, k][mask_k].unsqueeze(-1) * head_output[mask_k]
        return output

MoA vs 标准多头注意力

主要发现

算法应用领域

应用领域

代表性大模型

讨论

优势

• 计算效率：通过稀疏激活显著降低计算成本
• 模型容量：可扩展至万亿参数而不成比例增加计算
• 专业化学习：不同专家专注于不同知识领域
• 可解释性：通过分析专家分配机制理解模型行为

局限性

• 训练稳定性：动态专家选择可能导致负载不均衡和模型崩溃
• 系统复杂性：All-to-All通信模式增加系统设计难度
• 内存需求：多专家参数存储可能超出单设备容量

未来方向

1. 训练稳定性与负载均衡：开发更鲁棒的训练策略
2. 训练与系统效率：优化硬件-软件协同设计
3. 架构设计：使用元学习或强化学习动态调整专家数量
4. 理论发展：深入理解专家路由决策和聚类机制
5. 定制算法设计：探索MoE与对比学习、自监督学习的结合
6. 新应用领域：医疗、机器人、自动驾驶、教育、金融

相关工作

• Switch Transformer：首个万亿参数MoE模型
• V-MoE：视觉领域MoE应用
• Mixtral 8×7B：高效MoE语言模型
• Soft MoE：软分配MoE新范式
• [[BAGEL-Unified-Multimodal-Pretraining]]

参考文献

• Fedus et al. (2022). Switch Transformers: Scaling to Trillion Parameter Models. JMLR.
• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv.
• Riquelme et al. (2021). Scaling Vision with Sparse Mixture of Experts. NeurIPS.
• Jiang et al. (2024). Mixtral of Experts. arXiv.

UniDiffuser: One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale

论文链接 | GitHub

作者：Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang, Gang Yue, Yue Cao, Hang Su, Jun Zhu
发表于：ICML 2023

研究背景

扩散模型（Diffusion Models）在图像生成领域取得了巨大成功，但现有方法主要基于 U-Net 架构。随着 Transformer 在各领域的成功应用，如何将 Transformer 有效地应用于多模态扩散模型成为一个重要研究方向。现有的多模态生成方法通常需要为不同任务（文生图、图生文、联合生成等）设计不同的模型架构。

研究目标

• 多模态联合建模：用一个统一的框架同时处理图像、文本等多种模态的生成任务
• 任务统一：用单一模型支持文生图、图生文、联合生成、无条件生成等多种任务
• 架构创新：设计适合多模态扩散的 Transformer 架构（U-ViT）

核心概念

扩散模型基础

扩散模型通过前向过程逐步向数据添加噪声，再通过反向过程学习去噪，从而实现生成。

多模态独立加噪

UniDiffuser 的核心思想是对不同模态独立添加噪声，使用不同的时间步 $t$ 和 $s$，通过控制时间步实现任务切换。

U-ViT 架构

将 U-Net 的长跳跃连接（Long Skip Connection）引入 Vision Transformer，在浅层和深层之间建立连接。

研究方法

1. 单模态扩散基础公式

前向过程（加噪）

$$q(x_t | x_0) = \mathcal{N}(x_t; \alpha_t x_0, \sigma_t^2 I)$$
参数说明：

• $x_0$：原始干净数据
• $x_t$：时间步 $t$ 的加噪数据
• $\alpha_t, \sigma_t$：噪声调度参数，满足 $\alpha_t^2 + \sigma_t^2 = 1$（VP-SDE 设定）
• 随着 $t$ 增大，$\alpha_t \to 0$，$\sigma_t \to 1$，数据逐渐变成纯噪声

等价的重参数化表示：
$$x_t = \alpha_t x_0 + \sigma_t \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$
这个形式便于采样和训练，$\epsilon$ 是标准高斯噪声。

反向过程（去噪）

扩散模型学习反向过程 $p_\theta(x_{t-1}|x_t)$，通过神经网络预测噪声 $\epsilon_\theta(x_t, t)$，然后恢复 $x_0$：
$$\hat{x}_0 = \frac{x_t - \sigma_t \epsilon_\theta(x_t, t)}{\alpha_t}$$
解释：由 $x_t = \alpha_t x_0 + \sigma_t \epsilon$ 反解得到。预测出噪声后，即可估计原始数据。

2. 多模态联合扩散框架

联合前向过程

对于图文对 $(x_0, y_0)$，独立地对每个模态加噪：
$$q(x_t, y_s | x_0, y_0) = q(x_t | x_0) \cdot q(y_s | y_0)$$
展开形式：
$$x_t = \alpha_t x_0 + \sigma_t \epsilon_x, \quad y_s = \alpha_s y_0 + \sigma_s \epsilon_y$$
关键点：两个模态使用独立的时间步 $t$ 和 $s$，这是实现多任务统一的核心设计。

联合分布的边缘化

数据的联合分布 $q(x_0, y_0)$ 加噪后变为：
$$q(x_t, y_s) = \int q(x_t, y_s | x_0, y_0) q(x_0, y_0) , dx_0 dy_0$$
解释：这是对所有可能的原始数据对进行积分，得到加噪后数据的边缘分布。

3. 统一多任务的核心机制

核心洞察

通过控制 $t$ 和 $s$ 的取值，可以从联合分布中恢复各种边缘分布和条件分布：

条件生成的数学原理

当 $s=0$ 时，$y_s = y_0$（文本无噪声），此时：
$$q(x_t, y_0) = q(x_t | y_0) \cdot q(y_0)$$
解释：对 $y$ 不加噪声（$s=0$）在数学上等价于以 $y$ 为条件进行生成。这是一个优雅的设计——不需要修改模型架构，只需控制时间步即可切换任务。

4. 训练目标函数

噪声预测目标

模型 $\epsilon_\theta$ 同时预测两个模态的噪声：
$$[\hat{\epsilon}_x, \hat{\epsilon}y] = \epsilon_\theta(x_t, y_s, t, s)$$
完整训练损失：
$$\mathcal{L}(\theta) = \mathbb{E}{t, s, (x_0, y_0), \epsilon_x, \epsilon_y} \left[ \lambda_t |\epsilon_x - \hat{\epsilon}_x|^2 + \lambda_s |\epsilon_y - \hat{\epsilon}_y|^2 \right]$$
各项说明：

• $t, s \sim \mathcal{U}[0, T]$：从均匀分布采样时间步
• $(x_0, y_0) \sim q(x_0, y_0)$：从数据集采样图文对
• $\epsilon_x, \epsilon_y \sim \mathcal{N}(0, I)$：独立采样两个高斯噪声
• $\lambda_t, \lambda_s$：损失权重（通常设为 1）
• $|\cdot|^2$：均方误差损失

简化形式（实际训练中常用）：
$$\mathcal{L} = \mathbb{E}\left[|\epsilon_x - \hat{\epsilon}_x|^2 + |\epsilon_y - \hat{\epsilon}_y|^2\right]$$

5. 采样过程

联合采样（DDPM 形式）

从 $t=T$（纯噪声）开始，逐步去噪到 $t=0$：
$$x_{t-1} = \frac{1}{\sqrt{\alpha_{t|t-1}}} \left( x_t - \frac{1-\alpha_{t|t-1}}{\sigma_t} \hat{\epsilon}x \right) + \tilde{\sigma}t z$$
$$y{s-1} = \frac{1}{\sqrt{\alpha{s|s-1}}} \left( y_s - \frac{1-\alpha_{s|s-1}}{\sigma_s} \hat{\epsilon}_y \right) + \tilde{\sigma}_s z’$$
参数说明：

• $z, z’ \sim \mathcal{N}(0, I)$：采样的随机噪声（引入随机性）
• $\tilde{\sigma}_t$：后验方差，控制采样的随机程度
• $\alpha_{t|t-1} = \alpha_t / \alpha_{t-1}$：相邻时间步的比值

条件采样（文生图）

固定 $s=0$（即 $y_s = y_0$ 为输入文本），只对图像进行去噪：
$$x_{t-1} = \frac{1}{\sqrt{\alpha_{t|t-1}}} \left( x_t - \frac{1-\alpha_{t|t-1}}{\sigma_t} \hat{\epsilon}_x(x_t, y_0, t, 0) \right) + \tilde{\sigma}_t z$$
解释：文本时间步固定为 0，模型接收干净文本作为条件，只更新图像。

6. U-ViT 架构中的跳跃连接

U-ViT 在第 $l$ 层和第 $(L-l)$ 层之间添加跳跃连接：
$$h^{(L-l)} = \text{Block}^{(L-l)}\left( \text{Concat}(h^{(L-l-1)}, h^{(l)}) \right)$$
参数说明：

• $h^{(l)}$：第 $l$ 层的隐藏状态
• $L$：Transformer 总层数
• Concat：沿特征维度拼接
• 拼接后通过线性层降维回原始维度

设计动机：借鉴 U-Net 的成功经验，跳跃连接帮助保留低层的细节信息，有助于生成高质量图像。

7. 分类器自由引导（Classifier-Free Guidance, CFG）

为增强条件生成效果，使用 CFG 技术：
$$\tilde{\epsilon}_x = \epsilon_\theta(x_t, \varnothing, t, 0) + w \cdot \left( \epsilon_\theta(x_t, y_0, t, 0) - \epsilon_\theta(x_t, \varnothing, t, 0) \right)$$
参数说明：

• $\varnothing$：空文本条件（null condition）
• $w$：引导强度（guidance scale），通常 $w > 1$
• 第一项：无条件预测
• 括号内：条件预测与无条件预测的差值（条件信号）

等价形式：
$$\tilde{\epsilon}_x = (1-w) \cdot \epsilon_\theta(x_t, \varnothing, t, 0) + w \cdot \epsilon_\theta(x_t, y_0, t, 0)$$

训练技巧：训练时以一定概率（如 10%）将文本随机替换为空文本，使模型同时学习条件和无条件生成。

模型输入总结

主要发现

实验结果

MS-COCO 256×256 文生图（零样本）：

ImageNet 256×256 类条件生成：

讨论

优势

• 统一框架：单一模型支持多种生成任务，无需为每个任务单独训练
• 优雅设计：通过时间步控制实现任务切换，不需要修改架构
• 强大性能：在多个基准上达到 SOTA
• 可扩展性：在 10 亿参数规模上验证有效

局限性

• 需要大规模图文对数据进行训练
• 文本生成质量依赖于 CLIP 编码器的表示能力
• 推理速度受限于扩散模型的迭代采样

相关工作

• DALL-E：基于 VQ-VAE 的文生图模型
• Stable Diffusion：潜空间扩散模型
• DiT：Diffusion Transformer 架构[[Diffusion-Transformers-DiT]]
• GLIDE：文本引导的扩散模型

公式速查表

参考文献

• Bao, F., et al. (2023). One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale. ICML 2023.
• Ho, J., et al. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020.
• Peebles, W., & Xie, S. (2023). Scalable Diffusion Models with Transformers. ICCV 2023.

Scalable Diffusion Models with Transformers | ICCV 2023

研究背景

扩散模型(Diffusion Models)在图像生成领域取得了显著成功，但其架构设计一直沿用卷积U-Net作为主干网络。与此同时，Transformer架构已经在自然语言处理、视觉识别等多个领域取得了统治地位，并展现出优秀的可扩展性。本文探索将Transformer架构引入扩散模型，研究其在图像生成任务中的可扩展性和性能表现。

研究目标

• 突破架构局限：探索用Transformer替代传统U-Net作为扩散模型主干的可行性
• 验证可扩展性：研究Transformer在扩散模型中的可扩展性规律
• 建立性能基准：在ImageNet等基准数据集上达到SOTA性能
• 揭示计算-质量关系：分析模型计算量(Gflops)与生成质量之间的关系

核心概念

Latent Diffusion Models (LDMs)

在潜在空间而非像素空间训练扩散模型，提高计算效率：

1. VAE编码器：将图像压缩到潜在空间 $z = E(x)$
2. 扩散模型：在潜在空间 $z$ 中训练
3. VAE解码器：将生成的潜在表示解码为图像 $x = D(z)$

对于256×256×3的图像，VAE将其压缩为32×32×4的潜在表示（下采样因子为8）。

注意：这里使用的是标准VAE，输出是连续的潜在表示，而非VQ-VAE的离散codebook索引。

Patchify机制

将潜在表示分解为patch序列：

• 输入：32×32×4的潜在表示
• Patch大小：$p \times p$（$p \in {2, 4, 8}$）
• 输出序列长度：$T = (I/p)^2$

例如，$p=2$ 时，序列长度 $T = (32/2)^2 = 256$。

Classifier-Free Guidance

条件生成的采样技巧，提高生成质量：

$$
\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \emptyset) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset))
$$

其中：

• $c$：条件信息（如类别标签）
• $\emptyset$：空条件（训练时随机dropout）
• $s$：guidance scale（$s > 1$增强条件控制）

研究方法

DiT架构设计

DiT基于Vision Transformer (ViT)架构，包含以下组件：

1. Patchify层：将潜在表示转换为token序列
2. 位置编码：使用正弦-余弦位置编码
3. DiT Blocks：N个Transformer block
4. 解码器：将token序列解码为噪声预测和协方差预测

条件注入机制

探索了四种将时间步 $t$ 和类别标签 $c$ 注入Transformer的方式：

1. In-Context Conditioning

将 $t$ 和 $c$ 的embedding作为额外token添加到序列中：

1

tokens = [t_embed, c_embed, patch_1, patch_2, ..., patch_T]

• 优点：无需修改标准Transformer block
• 缺点：性能较差（FID ~80）

2. Cross-Attention

通过交叉注意力机制注入条件：

1
2
3

# 标准self-attention后添加cross-attention
x = x + SelfAttention(x)
x = x + CrossAttention(x, [t_embed, c_embed])

• 优点：灵活的条件控制
• 缺点：增加15%计算量，性能中等（FID ~60）

3. Adaptive Layer Norm (adaLN)

通过自适应归一化层注入条件：

1
2

γ, β = MLP(t_embed + c_embed)
output = γ * normalize(x) + β

• 优点：计算高效，性能较好（FID ~45）
• 缺点：所有token共享���同的调制参数

4. adaLN-Zero（最优方案）

在adaLN基础上增加门控参数并零初始化：

1
2
3
4
5
6
7
8

γ, β, α = MLP(t_embed + c_embed)

# Transformer Block
x = x + α₁ * Attention(γ₁ * normalize(x) + β₁)
x = x + α₂ * FFN(γ₂ * normalize(x) + β₂)

# 初始化：MLP输出零向量 → α=0, γ=0, β=0
# 因此初始时：x = x + 0 = x（恒等函数）

为什么需要两组参数？

Transformer block有两个子层（Attention + FFN），每个子层需要独立的条件控制：

• 第一组 $(γ_1, β_1, α_1)$：用于Attention子层
• 第二组 $(γ_2, β_2, α_2)$：用于FFN子层

总共6个参数：shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp

零初始化的优势：

• 训练稳定性：初始时网络是恒等映射，梯度流动顺畅
• 更好的性能：FID显著优于其他方法（FID ~23）
• 渐进式学习：从恒等函数开始，逐步学习有用的变换

模型配置

设计四种规模的模型配置：

Point-wise FFN

Transformer中的标准组件，对序列中每个位置独立应用前馈网络：

1
2
3
4
5
6
7
8
9
10
11
12

class PointwiseFeedForward(nn.Module):
    def forward(self, x):
        # x: [batch, seq_len, d_model]
        x = Linear1(x)      # d_model → d_ff (通常4倍扩展)
        x = GELU(x)
        x = Linear2(x)      # d_ff → d_model
        return x

# 关键特点：
# 1. 每个token独立处理（point-wise）
# 2. 所有位置共享相同的权重
# 3. 可以完全并行计算

与Self-Attention的分工：

• Self-Attention：全局信息聚合（不同token间交互）
• Point-wise FFN：局部特征变换（每个token独立处理）

这是所有Transformer变体（LLM、ViT、DiT）的统一设计。

主要发现

1. Gflops与生成质量强相关

模型前向传播的计算量(Gflops)与FID呈强负相关（相关系数-0.93）：

• 增加模型深度/宽度 → 提升Gflops → 降低FID
• 减小patch大小 → 增加token数量 → 提升Gflops → 降低FID

关键洞察：参数量不是���一决定因素，计算量才是提升性能的关键。

2. adaLN-Zero显著优于其他条件注入方式

在400K训练步数时的FID对比：

3. 优秀的可扩展性

DiT展现出与ViT类似的可扩展性：

• 增加模型规模持续提升性能
• 训练高度稳定，无需学习率预热或特殊正则化
• 未观察到常见的loss spike现象

4. 计算效率优势

DiT-XL/2 (118.6 Gflops) 比传统方法更高效：

• 像素空间U-Net (ADM)：1120 Gflops（~10倍）
• 潜在空间U-Net (LDM-4)：103.6 Gflops（相近但性能更优）

5. 采样计算无法弥补模型计算不足

增加采样步数（增加测试时计算量）无法弥补模型规模不足：

• DiT-L/2 使用1000步采样：80.7 Tflops，FID=25.9
• DiT-XL/2 使用128步采样：15.2 Tflops，FID=23.7

结论：模型计算量比采样计算量更重要。

实验设计

数据集

• ImageNet：256×256和512×512分辨率
• 任务：类条件图像生成

训练配置

• 优化器：AdamW
• 学习率：$1 \times 10^{-4}$（常数，无warmup）
• 批大小：256
• 数据增强：仅水平翻转
• EMA：decay=0.9999
• 硬件：TPU v3-256 pod

评估指标

• 主要指标：FID-50K（使用250步DDPM采样）
• 次要指标：Inception Score、sFID、Precision/Recall

结果分析

ImageNet 256×256基准测试：

ImageNet 512×512基准测试：

DiT-XL/2在两个分辨率上都达到了SOTA性能。

讨论

优势

• 架构统一性：证明Transformer可以成功替代U-Net，推动生成模型架构统一化
• 优秀的可扩展性：计算量与性能呈强相关，为大规模模型发展指明方向
• 训练稳定性：无需特殊技巧即可稳定训练
• 计算效率：在相近或更少的计算量下达到更好的性能
• 更高的Recall：相比LDM，DiT在所有guidance scale下都有更高的recall值

局限性

• 依赖预训练VAE：使用Stable Diffusion的VAE，是混合架构而非纯Transformer
• 仅探索类条件生成：未涉及文生图等更复杂的条件生成任务
• 计算资源需求：大规模模型训练需要TPU集群
• patch大小的权衡：更小的patch提升性能但增加计算量

与传统U-Net的对比

相关工作

Transformer在生成模型中的应用

• 自回归模型：GPT系列、ImageGPT、DALL·E
• 掩码生成模型：MaskGIT、MAGE
• DALL·E 2：使用Transformer生成CLIP embedding

扩散模型架构

• DDPM (Ho et al., 2020)：首次引入U-Net作为扩散模型主干
• ADM (Dhariwal & Nichol, 2021)：改进U-Net设计，达到SOTA
• LDM (Rombach et al., 2022)：潜在空间扩散模型
• Concurrent work：U-ViT探索了类似的Transformer架构

架构复杂度分析

• ViT (Dosovitskiy et al., 2021)：证明Transformer在视觉任务中的可扩展性
• Scaling Laws：语言模型中计算量与性能的幂律关系

未来方向

1. 扩展到文生图任务：将DiT应用于DALL·E 2、Stable Diffusion等文生图模型
2. ���大规模的模型：继续扩展模型规模，探索scaling law
3. 纯Transformer架构：在像素空间训练DiT，摆脱VAE依赖
4. 多模态条件生成：探索更复杂的条件注入机制
5. 高效采样方法：结合DiT开发更快的采样算法
6. 架构搜索：自动化探索最优的DiT配置

技术细节补充

VAE vs VQ-VAE

Transformer Block完整结构

1
2
3
4
5
6
7
8
9
10
11
12
13
14

def dit_block(x, c):
    # 生成6个调制参数
    γ₁, β₁, α₁, γ₂, β₂, α₂ = adaLN_modulation(c)

    # 子层1: Multi-Head Self-Attention
    x = x + α₁ * Attention(adaLN(x, γ₁, β₁))

    # 子层2: Point-wise Feed-Forward
    x = x + α₂ * FFN(adaLN(x, γ₂, β₂))

    return x

def adaLN(x, γ, β):
    return γ * normalize(x) + β

计算复杂度分析

对于DiT-XL/2（$N=256$ tokens，$d=1152$）：

• Self-Attention：$O(N^2 \cdot d) \approx 75M$ ops
• Point-wise FFN：$O(N \cdot d^2 \cdot 2) \approx 680M$ ops

FFN占据了大部分计算量！

参考文献

• Peebles, W., & Xie, S. (2023). Scalable Diffusion Models with Transformers. ICCV 2023.
• Ho, J., et al. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020.
• Rombach, R., et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022.
• Dosovitskiy, A., et al. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021.
• Dhariwal, P., & Nichol, A. (2021). Diffusion Models Beat GANs on Image Synthesis. NeurIPS 2021.

关键要点总结

1. 架构创新：首次系统性地将纯Transformer应用于扩散模型
2. adaLN-Zero：零初始化的自适应归一化是性能关键
3. Gflops定律：计算量与生成质量强相关（-0.93）
4. SOTA性能：ImageNet 256×256达到FID 2.27
5. 统一架构：推动生成模型向Transformer统一的趋势

论文链接 | NVIDIA, 2025

研究背景

人形机器人作为通用机器人的理想硬件平台，需要强大的基础模型来实现智能自主操作。受大语言模型和视觉模型成功的启发，研究者希望通过在大规模异构数据上训练机器人基础模型，使其能够理解新场景、处理真实世界的变化并快速学习新任务。然而，与文本和图像领域不同，机器人领域缺乏互联网规模的训练数据，不同机器人的传感器、自由度、控制模式差异巨大，形成”数据孤岛”问题。

研究目标

本论文要解决的核心问题：

1. 数据稀缺问题：人形机器人数据收集成本高、耗时长，如何突破真实数据瓶颈
2. 跨具身泛化：如何统一不同机器人的状态和动作空间，实现跨具身学习
3. 数据效率：如何在有限数据下快速适应新任务并在真实环境中鲁棒执行
4. 端到端优化：如何将高层推理与低层控制统一到单一模型中

核心概念

Vision-Language-Action (VLA) 模型

视觉-语言-动作模型，接收图像观察和语言指令作为输入，直接输出机器人动作。与传统的分层方法（VLM规划 + 低层策略执行）不同，VLA模型实现端到端优化。

双系统架构 (Dual-System Architecture)

受人类认知理论启发（Kahneman, 2011），将模型分为：

• System 2（推理系统）：慢速、深思熟虑的高层推理
• System 1（反应系统）：快速、自动化的低层控制

数据金字塔 (Data Pyramid)

将异构训练数据按规模和具身特异性组织成三层结构：

• 底层：大规模网络数据和人类视频（通用先验）
• 中层：合成数据（仿真+神经生成，可扩展）
• 顶层：真实机器人数据（具身特定，高质量）

潜在动作 (Latent Actions)

通过VQ-VAE([[VQ-VAE-and-Latent-Action-for-Robotics]])学习的通用动作表示，能够统一不同具身体（包括人类）的动作空间，使无动作标签的视频数据可用于训练。

研究方法

模型架构

GR00T N1采用双系统组合架构，总参数量22亿（GR00T-N1-2B）：

System 2: Vision-Language Module

1
2
3
4
5
6
7
8

输入处理:
├─ 图像: SigLIP-2编码器 → 64个token (224×224)
└─ 文本: SmolLM2 tokenizer → 文本token

特征提取:
└─ Eagle-2 VLM (1.34B参数)
   ├─ 处理vision-language tokens
   └─ 输出: 中间层embeddings φ_t (第12层)

关键设计：

• 使用中间层而非最终层特征（更快推理+更高成功率）
• 语言组件冻结（保留预训练知识）
• 视觉编码器可训练（适应机器人任务）
• 运行频率：10Hz

System 1: Diffusion Transformer Module

1
2
3
4
5
6
7

DiT Block结构（重复N次）:
├─ Self-Attention
│  └─ 输入: noised action tokens + state embeddings
│
└─ Cross-Attention
   ├─ Query: action/state tokens
   └─ Key & Value: VLM输出的φ_t

动作生成流程：

1. 输入加噪动作 $A_t^{\tau} = \tau A_t + (1-\tau)\epsilon$，其中 $\tau \in [0,1]$
2. 通过DiT迭代去噪（K=4步）
3. 输出16步动作序列（action chunking）
4. 运行频率：120Hz

Flow-Matching损失：

$$
\mathcal{L}{fm}(\theta) = \mathbb{E}{\tau} |V_{\theta}(\varphi_t, A_t^{\tau}, q_t) - (\epsilon - A_t)|^2
$$

其中 $V_{\theta}$ 是[[Diffusion-Transformers-DiT]]模型，预测去噪向量场。

模块交互机制

1
2
3
4
5
6
7
8
9
10
11
12

信息流:
图像 + 语言指令
    ↓
[System 2: Eagle-2 VLM]
    ↓ (输出 φ_t)
[Cross-Attention Bridge]
    ↓
[System 1: DiT]
├─ Self-Attention (action + state)
└─ Cross-Attention (attend to φ_t)
    ↓
16步动作序列

端到端联合训练：

• 两个模块通过cross-attention紧密耦合
• 使用统一的flow-matching loss优化
• 辅助目标检测loss增强空间理解：

$$
\mathcal{L} = \mathcal{L}{fm} + \mathcal{L}{det}
$$

异构数据训练策略

1. 数据金字塔组织

总计：8,376小时训练数据

2. 潜在动作学习

VQ-VAE训练：

1
2
3
4
5
6
7
8
9
10
11

# 编码器
输入: (当前帧 x_t, 未来帧 x_{t+H})
     ↓
Encoder → 连续embedding → 量化到codebook
     ↓
潜在动作 z_t

# 解码器
输入: x_t + z_t
     ↓
Decoder → 重建 x_{t+H}

跨具身一致性：

• 同一潜在动作在不同具身体中语义一致
• 例如：潜在动作1 = “右臂向左移动”（对所有机器人和人类）

训练使用：

• 提取预量化连续embedding作为”LAPA具身体”的动作
• 使用flow-matching loss训练

3. 神经轨迹生成

目标：从88小时真实数据扩增到827小时（~10倍）

技术流程：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

步骤1: 微调视频生成模型
├─ 基础模型: WAN2.1-I2V-14B
├─ 方法: LoRA微调
├─ 数据: 3,000条轨迹，81帧@480P
└─ 训练: 100 epochs

步骤2: 生成反事实轨迹
├─ 输入: 初始帧 + 新语言指令
├─ 语言生成: 多模态LLM检测物体
│   生成"pick {object} from {A} to {B}"
└─ 输出: 高质量视频

步骤3: 质量过滤
├─ 采样8帧 → LLM判断是否遵循指令
└─ 不合格 → 重新标注

步骤4: 动作标注
├─ 潜在动作编码器 → LAPA
└─ 逆动力学模型 → 伪动作标签

生成能力：

• 改变操作手（左手↔右手）
• 改变目标位置和物体
• 处理仿真难题（液体、铰接物体）
• 多视角生成（4宫格视频）

4. 仿真数据自动生成

DexMimicGen系统：

1
2
3
4
5
6
7
8
9
10
11

输入: 少量人类演示（几十条）
     ↓
分割 → 物体中心的子任务片段
     ↓
变换 → 根据新物体位置调整
     ↓
组合 → 插值并组合片段
     ↓
验证 → 仿真执行，保留成功轨迹
     ↓
输出: 每任务10,000条演示

规模：

• 54个源-目标容器组合
• 540,000条预训练轨迹
• 11小时生成 = 6,500小时等效人类演示

5. 具身特定编码器/解码器

处理不同维度的状态和动作：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

embodiments = {
    "GR-1": {
        "state": [joint_pos, joint_vel, base_pos, ...],
        "action": [joint_targets, ...],
        "encoder": MLP_GR1,
        "decoder": MLP_GR1
    },
    "Franka": {
        "state": [ee_pos, ee_rot, gripper],
        "action": [ee_delta, gripper_cmd],
        "encoder": MLP_Franka,
        "decoder": MLP_Franka
    },
    "LAPA": {  # 潜在动作
        "action": [latent_embedding],
        "encoder": MLP_LAPA,
        "decoder": MLP_LAPA
    }
}

6. 统一训练框架

预训练阶段：

• 全局batch size: 16,384
• 训练步数: 200,000
• 数据混合采样：真实机器人(40%) + 仿真(30%) + 神经(20%) + 人类视频(10%)
• 计算资源: 最多1024个H100 GPU，约50,000 GPU小时

后训练阶段：

• Batch size: 128-1024
• 训练步数: 20,000-60,000
• 可选神经轨迹协同训练（1:1采样比例）
• 可在单个A6000 GPU上微调

主要发现

预训练泛化能力

在GR-1人形机器人上的零样本评估：

仿真基准测试

100条演示/任务的性能对比：

关键观察：

• GR00T N1在所有基准上均优于基线
• 在GR-1任务上优势最明显（+17.3%）

真实世界部署

GR-1人形机器人任务成功率：

数据效率：

• GR00T N1用10%数据（42.6%）≈ Diffusion Policy用全量数据（46.4%）
• 展现出色的样本效率

神经轨迹增强效果

RoboCasa基准（协同训练3K神经轨迹/任务）：

真实世界（协同训练100神经轨迹/任务）：

• 平均提升：+5.8%

观察：

• 低数据场景：LAPA略优（更通用的先验）
• 高数据场景：IDM更优（更接近真实动作）

定性分析

运动质量：

• GR00T N1运动更流畅，抓取精度更高
• Diffusion Policy常出现初始帧不动、抓取不准确

泛化能力：

• 预训练模型能执行未见过的双手交接任务
• 后训练模型在特定任务上更精确，但失去部分泛化能力

实验设计

仿真基准

RoboCasa Kitchen（24任务）：

• 机器人：Franka Emika Panda
• 任务：抓取放置、开关门、按按钮、转水龙头等
• 观察：3个RGB相机（左、右、腕部）
• 动作：末端执行器相对位姿 + 夹爪状态
• 数据：每任务3,000条MimicGen生成的演示

DexMimicGen Cross-Embodiment Suite（9任务）：

• 具身体：
  • 双臂Panda + 平行夹爪（穿线、组装、运输）
  • 双臂Panda + 灵巧手（清理、抬托盘）
  • GR-1人形 + 灵巧手（倒水、咖啡、分类）
• 数据：每任务1,000条演示

GR-1 Tabletop Tasks（24任务）：

• 机器人：GR-1人形 + Fourier灵巧手
• 任务：18个重排任务 + 6个铰接物体任务
• 观察：头部自我中心相机
• 动作：关节位置/旋转 + 腰部/颈部
• 数据：每任务1,000条DexMimicGen生成

真实世界基准

任务类别：

1. 抓取放置（5任务）：

   • 托盘→盘子、砧板→篮子、餐垫→碗等
   • 评估：见过和未见过物体

2. 铰接物体（3任务）：

   • 白色抽屉、深色柜子、木箱
   • 要求：放入物体并关闭

3. 工业操作（3任务）：

   • 机械零件打包
   • 网格杯倾倒
   • 圆柱体交接

4. 多机协作（2任务）：

   • 第1部分：抓取→放入网格杯→交给另一机器人
   • 第2部分：接收→放入黄色箱→倾倒剩余物

数据收集：

• 遥操作时长：15分钟-3小时/任务
• 过滤低质量轨迹

评估协议

仿真：

• 每任务100次试验
• 取最后5个checkpoint的最大值
• Checkpoint间隔：500步

真实机器人：

• 每任务10次试验（机械打包任务5次）
• 部分评分系统（捕捉不同执行阶段）
• 低数据场景：10%数据子采样

训练配置

预训练：

• 学习率：1e-4
• 优化器：AdamW (β1=0.95, β2=0.999)
• 学习率调度：cosine，warmup比例0.05
• Batch size：16,384
• 步数：200,000

后训练：

• Batch size：128-1024
• 步数：20,000-60,000
• 其他超参数同预训练

讨论

优势

1. 统一的跨具身学习：

   • 单一模型支持从桌面机械臂到双臂人形机器人
   • 潜在动作空间统一不同具身体

2. 卓越的数据效率：

   • 10%数据达到基线全量数据性能
   • 预训练提供强大的先验知识

3. 可扩展的数据生成：

   • 神经轨迹生成：10倍数据扩增
   • 仿真自动生成：11小时生成6,500小时等效数据

4. 端到端优化：

   • VLM推理与DiT控制联合训练
   • 避免分层方法的接口问题

5. 开源生态：

   • 公开22亿参数模型
   • 提供训练数据和仿真基准

局限性

1. 任务范围限制：

   • 当前主要关注短时域桌面操作
   • 未涉及长时域移动操作（loco-manipulation）

2. 合成数据质量：

   • 视频生成模型仍面临多样性和物理一致性挑战
   • 需要质量过滤和重新标注

3. 硬件依赖：

   • 需要高端GPU进行训练（H100集群）
   • 推理需要L40 GPU（63.9ms/16动作）

4. 泛化-专精权衡：

   • 后训练提升特定任务性能但损失部分泛化能力
   • 预训练模型能执行双手交接，后训练模型失去此能力

5. 视觉-语言骨干限制：

   • 当前VLM的空间推理和语言理解能力仍有提升空间
   • 更强的VLM可能进一步提升性能

相关工作

机器人基础模型

VLA模型：

• RT-1/RT-2 (Brohan et al., 2022, 2023)：早期VLA模型，使用Transformer架构
• π0 (Black et al., 2024)：使用mixture-of-experts连接VLM和动作生成
• Octo (Octo Model Team et al., 2024)：跨具身模型，但不微调VLM
• GR-2 (Cheang et al., 2024)：视频-语言-动作模型

GR00T N1的区别：

• 使用简单的cross-attention而非MoE
• 端到端微调VLM视觉编码器
• 支持潜在动作和IDM伪动作

机器人数据集

真实机器人数据：

• Open X-Embodiment (2024)：跨具身数据集联盟
• AgiBot-Alpha (2025)：100个机器人的大规模数据集
• 遥操作系统：VIVE、Apple Vision Pro、Leap Motion

人类视频数据：

• Ego4D (Grauman et al., 2022)：大规模自我中心视频
• EPIC-KITCHENS (Damen et al., 2018)：厨房活动
• Assembly-101 (Sener et al., 2022)：组装任务

GR00T N1的创新：

• 数据金字塔组织而非简单混合
• 潜在动作统一有/无标签数据

合成数据生成

仿真数据：

• MimicGen (Mandlekar et al., 2023)：演示变换和重放
• DexMimicGen (Jiang et al., 2024)：灵巧操作数据生成
• RoboCasa (Nasiriany et al., 2024)：厨房环境仿真

神经生成：

• 视频生成模型：Sora (Brooks et al., 2024)、WAN (Wan Team, 2025)
• 数据增强：GenAug (Chen et al., 2023)使用扩散模型增强

GR00T N1的规模：

• 827小时神经轨迹（前所未有）
• 540K仿真轨迹（11小时生成）

未来方向

1. 长时域移动操作：

   • 扩展到全身运动和导航
   • 需要改进硬件、模型架构和训练数据

2. 更强的视觉-语言骨干：

   • 提升空间推理能力
   • 增强语言理解和任务规划

3. 改进合成数据生成：

   • 提高视频生成的多样性和反事实能力
   • 增强物理一致性和真实感
   • 探索自动化初始帧生成（img2img扩散）

4. 新型模型架构：

   • 探索更高效的推理-控制耦合方式
   • 研究分层时间建模

5. 鲁棒性和泛化：

   • 提升对环境变化的适应能力
   • 增强零样本和少样本学习能力

6. 多模态感知：

   • 整合触觉、力觉等其他传感器
   • 探索多模态融合策略

7. 长时域视频生成：

   • 多轮视频生成实现长任务序列
   • 原子任务组合

参考文献

• NVIDIA (2025). GR00T N1: An Open Foundation Model for Generalist Humanoid Robots. arXiv:2503.14734v2.
• Black et al. (2024). π0: A vision-language-action flow model for general robot control. arXiv:2410.24164.
• Brohan et al. (2022). RT-1: Robotics transformer for real-world control at scale. arXiv:2212.06817.
• Brohan et al. (2023). RT-2: Vision-language-action models transfer web knowledge to robotic control. arXiv:2307.15818.
• Chi et al. (2024). Diffusion Policy: Visuomotor policy learning via action diffusion. IJRR.
• Jiang et al. (2024). DexMimicGen: Automated data generation for bimanual dexterous manipulation via imitation learning. CoRL.
• Mandlekar et al. (2023). MimicGen: A data generation system for scalable robot learning using human demonstrations. CoRL.
• Nasiriany et al. (2024). RoboCasa: Large-scale simulation of everyday tasks for generalist robots. RSS.
• Open X-Embodiment Collaboration et al. (2024). Open X-Embodiment: Robotic learning datasets and RT-X models.
• Ye et al. (2025). Latent action pretraining from videos. ICLR.
• Kahneman (2011). Thinking, Fast and Slow. Farrar, Straus and Giroux.

关键代码和资源

• 模型权重：HuggingFace
• 训练数据：HuggingFace Datasets
• 仿真基准：GitHub
• 数据格式：基于LeRobot格式扩展
• 训练基础设施：NVIDIA OSMO编排平台

技术细节补充

动作空间标准化

统一不同具身体的表示：

• 末端执行器旋转状态：6D旋转表示
• 末端执行器旋转动作：轴角表示
• 位置和关节：Min-max归一化
• 顺序：左臂→右臂，旋转→位置→夹爪

辅助目标检测损失

使用OWL-v2检测器标注目标物体边界框：

$$
\mathcal{L}{det} = |\mathbf{x}{pred} - \mathbf{x}_{gt}|^2
$$

其中 $\mathbf{x}$ 是归一化的边界框中心坐标。

推理性能

• GR00T-N1-2B：63.9ms采样16步动作（L40 GPU，bf16）
• VLM频率：10Hz
• 动作输出频率：120Hz
• 去噪步数：K=4

计算资源

• 预训练：最多1024个H100 GPU，约50,000 GPU小时
• 神经轨迹生成：3,600个L40 GPU，约105K GPU小时（1.5天）
• 后训练：单个A6000 GPU可微调（仅adapter层时batch size可达200）

Neural Discrete Representation Learning | VQ-BeT: Behavior Generation with Latent Actions

研究背景

在无监督学习和机器人学习领域，表示学习是核心问题之一。传统的变分自编码器（VAE, Variational AutoEncoder）使用连续潜在变量，但存在后验崩塌（posterior collapse）问题，即解码器过强导致忽略潜在编码。

在机器人学习中，直接学习连续高维动作空间面临以下挑战：

• 动作分布通常是多模态的（如抓取物体可以有多种方式）
• 行为克隆（Behavior Cloning）容易产生平均化的次优动作
• 连续动作空间的策略学习不稳定

VQ-VAE（Vector Quantised-Variational AutoEncoder）通过引入离散潜在变量，为这些问题提供了有效的解决方案。

研究目标

VQ-VAE的核心目标

1. 解决VAE中的后验崩塌问题
2. 学习有效的离散表示，适用于本质上离散的数据（语言、语音等）
3. 实现端到端的离散表示学习

机器人Latent Action的目标

1. 将连续高维动作空间压缩为离散的动作原语（action primitives）
2. 在离散空间中学习更稳定的策略
3. 实现时序动作抽象，降低决策频率
4. 提升多模态动作分布的建模能力

核心概念

VQ-VAE (Vector Quantised-Variational AutoEncoder)

VQ-VAE是一种使用离散潜在变量的生成模型，通过向量量化（Vector Quantization）技术将编码器输出映射到离散的码本空间。

关键组件：

• 编码器（Encoder）：将输入映射到连续潜在空间
• 码本（Codebook）：包含 $K$ 个 $d$ 维向量 $\mathbf{e} \in \mathbb{R}^{K \times d}$
• 量化层（Quantization）：将连续表示映射到最近的码本向量
• 解码器（Decoder）：从离散表示重建输入

Latent Action（潜在动作）

Latent Action是将连续动作序列编码为离散token的表示方法。每个离散token代表一个”动作原语”或”技能”，可以解码为一段连续的动作序列。

核心思想：

• 将动作序列 $\mathbf{a}_{t:t+H}$ 编码为单个离散索引 $z \in {1, …, K}$
• 策略网络在离散空间中选择动作：$\pi(\mathbf{o}_t) \rightarrow z$
• 解码器将离散索引恢复为连续动作：$z \rightarrow \mathbf{a}_{t:t+H}$

VQ-VAE方法详解

架构设计

1
2
3
4
5
6
7
8
9
10
11
12
13

输入 x
  ↓
[编码器] Encoder
  ↓
z_e(x) ∈ R^(H×W×D)  (连续潜在表示)
  ↓
[向量量化] Vector Quantization
  ↓
z_q(x) ∈ R^(H×W×D)  (离散潜在表示)
  ↓
[解码器] Decoder
  ↓
重建输出 x̂

向量量化过程

对于编码器输出的每个空间位置，找到最近的码本向量：

$$
z_q(\mathbf{x}) = \mathbf{e}_k, \quad \text{where} \quad k = \arg\min_j |\mathbf{z}_e(\mathbf{x}) - \mathbf{e}_j|_2
$$

损失函数

VQ-VAE使用三部分损失函数：

$$
L = \log p(\mathbf{x}|\mathbf{z}_q(\mathbf{x})) + |\text{sg}[\mathbf{z}_e(\mathbf{x})] - \mathbf{e}|_2^2 + \beta |\mathbf{z}_e(\mathbf{x}) - \text{sg}[\mathbf{e}]|_2^2
$$

其中：

• 重建损失（Reconstruction Loss）：$\log p(\mathbf{x}|\mathbf{z}_q(\mathbf{x}))$，确保重建质量
• 码本损失（Codebook Loss）：$|\text{sg}[\mathbf{z}_e(\mathbf{x})] - \mathbf{e}|_2^2$，更新码本向量使其靠近编码器输出
• 承诺损失（Commitment Loss）：$\beta |\mathbf{z}_e(\mathbf{x}) - \text{sg}[\mathbf{e}]|_2^2$，鼓励编码器输出靠近码本向量（$\beta = 0.25$）

其中 $\text{sg}[\cdot]$ 表示stop gradient操作，阻止梯度传播。

Straight-Through Estimator

问题：量化操作 $\mathbf{z}q = \arg\min{\mathbf{e}} |\mathbf{z}_e - \mathbf{e}|$ 不可微分

解决方案：在反向传播时，将解码器的梯度直接复制给编码器：

$$
\nabla_{\mathbf{z}e} L = \nabla{\mathbf{z}_q} L
$$

即在前向传播使用离散的 $\mathbf{z}_q$，在反向传播时假装量化操作是恒等映射。

数据尺度变化示例

以CIFAR-10图像重建为例：

信息压缩率：$(32 \times 32 \times 3) / (8 \times 8 \times \log_2 512) \approx 42$ 倍压缩（假设码本大小 $K=512$）

VQ-VAE在机器人中的应用

整体架构

1
2
3
4
5
6
7
8
9
10
11

观察 o_t (图像/状态)
    ↓
[策略网络 π]
    ↓
离散latent action z ∈ {1,...,K}
    ↓
[VQ-VAE解码器]
    ↓
连续动作序列 a_{t:t+H}
    ↓
执行到机器人

动作序列编码

输入：动作序列 $\mathbf{a}_{t:t+H} \in \mathbb{R}^{H \times d_a}$，其中 $H$ 是序列长度，$d_a$ 是动作维度

编码过程：

1. 通过1D卷积或Transformer编码时序信息
2. 输出单个向量 $\mathbf{z}_e \in \mathbb{R}^D$
3. 量化为离散索引 $k \in {1, …, K}$

解码过程：

1. 从码本中查找向量 $\mathbf{e}_k$
2. 通过解码器生成动作序列 $\hat{\mathbf{a}}_{t:t+H}$

训练流程

阶段1：训练VQ-VAE

使用专家演示数据训练VQ-VAE：

1
2
3
4
5
6
7
8
9
10
11
12
13
14

for batch in expert_demonstrations:
    action_seq = batch['actions']  # [B, H, action_dim]

    # 编码-量化-解码
    z_e = encoder(action_seq)
    z_q = quantize(z_e, codebook)
    action_recon = decoder(z_q)

    # 三部分损失
    loss_recon = MSE(action_seq, action_recon)
    loss_vq = MSE(sg(z_e), z_q)
    loss_commit = MSE(z_e, sg(z_q))

    loss = loss_recon + loss_vq + 0.25 * loss_commit

阶段2：训练策略网络

固定VQ-VAE，训练策略在离散空间中选择动作：

1
2
3
4
5
6
7
8
9
10
11

for batch in demonstrations:
    obs = batch['observations']  # [B, T, obs_dim]
    actions = batch['actions']   # [B, T, action_dim]

    # 将动作编码为离散token
    with torch.no_grad():
        z_indices = vqvae.encode(actions)  # [B, T//H]

    # 训练策略预测离散token
    z_pred = policy(obs)  # [B, T//H, K]
    loss = CrossEntropy(z_pred, z_indices)

主要应用案例

VQ-BeT (VQ-Behavior Transformer)

论文：Behavior Generation with Latent Actions (CoRL 2023)

核心思想：

1. 使用VQ-VAE将动作序列压缩为离散token
2. 使用Transformer建模观察到latent action的映射：$p(z_t | \mathbf{o}_{1:t})$
3. 执行时解码latent action为连续动作序列

优势：

• 有效处理多模态动作分布
• 避免行为克隆中的动作平均化问题
• 支持长时序动作规划（一次预测多步）

LISA (Latent Imagination with Skill Abstraction)

核心思想：结合世界模型和latent action

1
2
3
4
5
6
7
8
9

当前状态 s_t
    ↓
[世界模型] 在latent space中想象
    ↓
预测未来状态序列 ŝ_{t+1:t+H}
    ↓
[规划器] 选择最优latent action z*
    ↓
[VQ解码器] z* → 连续动作

SPiRL (Skill-based Model-based RL)

将VQ-VAE学习的离散表示视为”技能”，在强化学习中进行技能级别的规划。

实验设计与结果

VQ-VAE实验（原始论文）

数据集

• CIFAR-10：32×32彩色图像
• ImageNet：128×128和256×256图像
• VCTK语音数据集：英语语音数据
• DeepMind Lab：强化学习环境视频

关键参数

• 码本大小 $K$：512
• 编码维度 $D$：64
• 承诺损失系数 $\beta$：0.25

主要结果

机器人Latent Action实验

超参数选择

性能对比

VQ-BeT在多个机器人操作任务上的表现：

讨论

优势

VQ-VAE本身：

• 避免后验崩塌问题，潜在编码被充分利用
• 离散表示更适合某些模态（语言、符号）
• 可以学习到有意义的离散结构

在机器人中的优势：

• 多模态建模：离散分类比连续回归更容易处理多模态动作分布
• 时序抽象：一个latent action代表一段动作序列，降低决策频率
• 训练稳定性：离散空间避免连续动作的梯度不稳定
• 可解释性：码本向量可视为”技能原语”，便于分析和调试
• 泛化能力：学到的动作原语可以组合应用到新场景

局限性

VQ-VAE的挑战：

• 码本利用率问题（codebook collapse）：部分码本向量可能不被使用
• 重建误差：离散化导致信息损失
• 超参数敏感：$K$、$D$、$\beta$ 需要仔细调优

机器人应用的挑战：

• 重建精度：VQ-VAE无法完美重建动作，影响执行精度
• 序列长度选择：$H$ 的选择需要在抽象能力和精确控制之间权衡
• 计算开销：需要额外训练VQ-VAE模型
• 在线适应：预训练的码本可能不适合新任务

相关工作

离散表示学习

• VQ-VAE-2 (Razavi et al., 2019)：层次化VQ-VAE，提升生成质量
• DALL-E (Ramesh et al., 2021)：使用VQ-VAE的离散表示进行文本到图像生成
• Gumbel-Softmax VAE：另一种离散VAE方法，使用Gumbel-Softmax技巧

机器人技能学习

• Skill Discovery：无监督发现技能的方法（DIAYN, DADS等）
• Hierarchical RL：层次化强化学习，在不同抽象层次上决策
• Option Framework：时序抽象的经典框架

行为克隆与模仿学习

• Diffusion Policy (Chi et al., 2023)：使用扩散模型生成动作
• Action Chunking Transformer (Zhao et al., 2023)：直接预测动作序列
• BeT (Shafiullah et al., 2022)：使用离散化动作的行为Transformer

未来方向

方法改进

1. 层次化VQ-VAE：

   • 高层策略选择宏观latent action
   • 低层策略选择微观latent action
   • 实现多层次的时序抽象

2. 与扩散模型结合：

   • 使用VQ-VAE的离散表示作为扩散模型的条件
   • 在离散空间规划，在连续空间精细化
   • 结合两者优势：稳定性+精确性

3. 在线学习与适应：

   • 预训练VQ-VAE在大规模数据上
   • 在新任务上微调策略网络
   • 探索码本的在线更新机制

4. 解决码本崩塌：

   • 使用EMA（指数移动平均）更新码本
   • 引入正则化鼓励码本多样性
   • 动态调整码本大小

应用拓展

1. 多模态机器人学习：

   • 结合视觉、触觉、本体感觉
   • 学习跨模态的统一表示

2. 长时序任务规划：

   • 在latent action空间进行任务规划
   • 结合符号推理和连续控制

3. 迁移学习：

   • 在源任务上学习通用动作原语
   • 在目标任务上组合和微调

4. 人机协作：

   • 可解释的动作原语便于人类理解
   • 支持人类通过选择latent action进行干预

参考文献

核心论文

• van den Oord, A., Vinyals, O., & Kavukcuoglu, K. (2017). Neural Discrete Representation Learning. NIPS 2017. arXiv:1711.00937
• Shafiullah, N. M. M., et al. (2023). Behavior Generation with Latent Actions (VQ-BeT). CoRL 2023. arXiv:2403.03181

相关工作

• Razavi, A., et al. (2019). Generating Diverse High-Fidelity Images with VQ-VAE-2. NeurIPS 2019.
• Pertsch, K., et al. (2020). Accelerating Reinforcement Learning with Learned Skill Priors (SPiRL). CoRL 2020.
• Lynch, C., & Sermanet, P. (2020). Learning Latent Plans from Play (LISA). CoRL 2020.
• Chi, C., et al. (2023). Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. RSS 2023.

关键代码示例

VQ-VAE量化层实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

import torch
import torch.nn as nn

class VectorQuantizer(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, commitment_cost=0.25):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.num_embeddings = num_embeddings
        self.commitment_cost = commitment_cost

        # 初始化码本
        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.embedding.weight.data.uniform_(-1/num_embeddings, 1/num_embeddings)

    def forward(self, z_e):
        # z_e: [B, D] 编码器输出

        # 计算距离
        distances = torch.sum(z_e**2, dim=1, keepdim=True) + \
                    torch.sum(self.embedding.weight**2, dim=1) - \
                    2 * torch.matmul(z_e, self.embedding.weight.t())

        # 找到最近的码本向量
        encoding_indices = torch.argmin(distances, dim=1)
        z_q = self.embedding(encoding_indices)

        # 计算损失
        e_latent_loss = torch.mean((z_q.detach() - z_e)**2)  # 码本损失
        q_latent_loss = torch.mean((z_q - z_e.detach())**2)  # 承诺损失
        loss = e_latent_loss + self.commitment_cost * q_latent_loss

        # Straight-through estimator
        z_q = z_e + (z_q - z_e).detach()

        return z_q, loss, encoding_indices

动作序列编码器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

class ActionEncoder(nn.Module):
    def __init__(self, action_dim, hidden_dim, latent_dim, seq_len):
        super().__init__()
        self.seq_len = seq_len

        # 1D卷积编码时序信息
        self.conv1 = nn.Conv1d(action_dim, hidden_dim, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.Conv1d(hidden_dim, hidden_dim*2, kernel_size=4, stride=2, padding=1)
        self.fc = nn.Linear(hidden_dim*2 * (seq_len//4), latent_dim)

    def forward(self, action_seq):
        # action_seq: [B, seq_len, action_dim]
        x = action_seq.transpose(1, 2)  # [B, action_dim, seq_len]
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = x.flatten(1)
        z_e = self.fc(x)
        return z_e