ACT (Action Chunking Transformer) 模型架构
1. 整体架构概览
ACT (Action Chunking Transformer) 是一种基于 Transformer 编码器-解码器结构的机器人操作策略模型,出自论文 Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware (Zhao et al., 2023)。其核心思想是 动作分块 (Action Chunking):模型一次性预测未来 chunk_size 步的动作序列,而非逐步预测,从而有效缓解复合误差问题。ACT 使用 ResNet18 作为视觉编码器提取图像特征,可选的 VAE 编码器学习动作分布的隐变量,以及 Transformer 编码器-解码器生成最终的动作序列。推理时仅需单次前向传播,无需迭代去噪。
(多视角)
[B, C, H, W]"] STATE["机器人状态
(关节角度)
[B, state_dim]"] ENV["环境状态 (可选)
[B, env_dim]"] ACT_IN["动作序列 (仅训练)
[B, chunk_size, action_dim]"] end subgraph VisionBackbone["ResNet18 视觉编码器"] RES["ResNet18
(预训练 ImageNet 权重)
FrozenBatchNorm2d"] PROJ_IMG["1x1 卷积投影
512 → 512"] POS2D["2D 正弦位置编码"] end subgraph VAE["VAE 编码器 (可选, 仅训练)"] CLS["[CLS] Token"] VAE_ENC["VAE Transformer 编码器
(4层)"] LATENT["隐变量采样
z ~ N(mu, sigma)"] end subgraph TransformerCore["Transformer 编码器-解码器"] ENC["Transformer 编码器
(4层, 8头, dim=512)"] DEC["Transformer 解码器
(1层, 交叉注意力)"] end subgraph Output["输出"] HEAD["动作回归头
(Linear: 512 → action_dim)"] ACT_OUT["预测动作序列
[B, chunk_size, action_dim]"] end IMG --> RES --> PROJ_IMG PROJ_IMG -->|"图像特征 + 2D 位置编码"| ENC POS2D --> ENC STATE -->|"状态嵌入"| ENC ENV -->|"环境嵌入 (可选)"| ENC ACT_IN --> VAE_ENC STATE --> VAE_ENC CLS --> VAE_ENC VAE_ENC --> LATENT -->|"隐变量 z
[B, latent_dim]"| ENC ENC -->|"编码器输出
[S, B, 512]"| DEC DEC --> HEAD --> ACT_OUT style Input fill:#e8f4fd,stroke:#2196F3 style VisionBackbone fill:#fff3e0,stroke:#FF9800 style VAE fill:#f3e5f5,stroke:#9C27B0 style TransformerCore fill:#e8f5e9,stroke:#4CAF50 style Output fill:#fce4ec,stroke:#E91E63
2. 核心组件详解
2.1 ResNet18 视觉编码器
ACT 使用 torchvision 提供的 ResNet18 作为图像特征提取器,默认加载 ImageNet 预训练权重。特征图取自 layer4 的输出,经过 1x1 卷积投影到 Transformer 的隐藏维度 (512),并附加 2D 正弦位置编码。
[B, 3, H, W]"] --> CONV1["conv1 + bn1 + relu
+ maxpool"] CONV1 --> L1["layer1
(2 BasicBlock)"] L1 --> L2["layer2
(2 BasicBlock)"] L2 --> L3["layer3
(2 BasicBlock)"] L3 --> L4["layer4
(2 BasicBlock)
可选: 膨胀卷积替代步幅"] end subgraph Projection["特征投影"] L4 -->|"feature_map
[B, 512, h, w]"| CONV1x1["Conv2d 1x1
512 → 512"] CONV1x1 -->|"[B, 512, h, w]"| RESHAPE["展平为序列
[h*w, B, 512]"] end subgraph PosEmbed["2D 正弦位置编码"] RESHAPE --> ADD["+ pos_embed"] PE["ACTSinusoidalPosition
Embedding2d
(dim=256)"] -->|"[1, 512, h, w]
→ [h*w, 1, 512]"| ADD end ADD --> OUT["送入 Transformer 编码器"] style Backbone fill:#fff3e0,stroke:#FF9800 style Projection fill:#e3f2fd,stroke:#2196F3 style PosEmbed fill:#e8f5e9,stroke:#4CAF50
关键设计细节:
- 使用
IntermediateLayerGetter提取layer4的特征图,而非 ResNet 的最终分类输出 - Batch Normalization 使用
FrozenBatchNorm2d,训练时不更新统计量 replace_final_stride_with_dilation选项可将layer4的 2x2 步幅替换为膨胀卷积,保留更高分辨率的特征图- 多相机视角的图像逐一通过同一个 backbone,特征拼接到编码器的序列维度上
2.2 VAE 编码器 (可选)
VAE 编码器是 ACT 的核心创新之一。训练时,它将真实动作序列编码为隐变量 z,捕获动作分布中的多模态性;推理时跳过 VAE 编码器,直接使用零向量作为隐变量。
(可学习, nn.Embedding)
[B, 1, 512]"] RS["机器人状态嵌入
(Linear: state_dim → 512)
[B, 1, 512]"] ACT_SEQ["动作序列嵌入
(Linear: action_dim → 512)
[B, chunk_size, 512]"] end subgraph VAEEncoder["VAE Transformer 编码器 (4层)"] CONCAT["拼接: [CLS, state, actions]
[B, chunk_size+2, 512]"] SINPE["固定正弦位置编码"] PAD_MASK["填充掩码
(action_is_pad)"] TENC["ACTEncoder
(4层 ACTEncoderLayer)
+ LayerNorm"] end subgraph LatentSample["隐变量采样"] CLS_OUT["[CLS] Token 输出
[B, 512]"] LINEAR_P["线性投影
512 → 64
(latent_dim * 2)"] MU["mu
[B, 32]"] LOG_S["log(sigma^2)
[B, 32]"] REPARAM["重参数化技巧
z = mu + sigma * eps
eps ~ N(0, 1)"] Z["隐变量 z
[B, 32]"] end CLS --> CONCAT RS --> CONCAT ACT_SEQ --> CONCAT SINPE --> TENC PAD_MASK --> TENC CONCAT --> TENC TENC -->|"取第0个token"| CLS_OUT CLS_OUT --> LINEAR_P LINEAR_P --> MU LINEAR_P --> LOG_S MU --> REPARAM LOG_S --> REPARAM REPARAM --> Z style VAEInput fill:#f3e5f5,stroke:#9C27B0 style VAEEncoder fill:#ede7f6,stroke:#673AB7 style LatentSample fill:#fce4ec,stroke:#E91E63
VAE 的设计意图:
- 人类演示数据中,同一状态可能对应多种合理的动作 (多模态性)。VAE 通过隐变量 z 对这种不确定性进行建模
- 训练时,VAE 编码器从真实动作序列中提取隐变量;推理时使用零向量 (先验均值),对应最可能的动作模式
- KL 散度损失约束隐变量分布接近标准正态分布 N(0, I),确保推理时零向量采样合理
2.3 Transformer 编码器-解码器
这是 ACT 的核心推理模块。Transformer 编码器融合隐变量、状态嵌入和图像特征;Transformer 解码器通过交叉注意力从编码器输出中生成动作序列。
(Linear: 32 → 512)
[1, B, 512]"] T_S["机器人状态 Token
(Linear: state_dim → 512)
[1, B, 512]"] T_E["环境状态 Token (可选)
(Linear: env_dim → 512)
[1, B, 512]"] T_IMG["图像特征 Token
(每相机 h*w 个)
[h*w*N_cam, B, 512]"] end subgraph Encoder["ACTEncoder (4层)"] ENC_PE["位置编码
1D token: 可学习 Embedding
图像 token: 2D 正弦编码"] ENC_LAYERS["ACTEncoderLayer x 4
自注意力 + FFN
(Post-Norm)"] ENC_NORM["LayerNorm (若 Pre-Norm)"] end subgraph Decoder["ACTDecoder (1层)"] DEC_QUERY["解码器查询
全零初始化
[chunk_size, B, 512]"] DEC_PE["可学习位置嵌入
(nn.Embedding)
[chunk_size, 1, 512]"] DEC_LAYER["ACTDecoderLayer x 1
自注意力 → 交叉注意力 → FFN"] DEC_NORM["LayerNorm"] end T_Z --> ENC_LAYERS T_S --> ENC_LAYERS T_E --> ENC_LAYERS T_IMG --> ENC_LAYERS ENC_PE --> ENC_LAYERS ENC_LAYERS --> ENC_NORM ENC_NORM -->|"编码器输出
[S, B, 512]"| DEC_LAYER DEC_QUERY --> DEC_LAYER DEC_PE --> DEC_LAYER DEC_LAYER --> DEC_NORM DEC_NORM -->|"解码器输出
[chunk_size, B, 512]"| OUT["→ 动作头"] style EncoderInput fill:#e8f4fd,stroke:#2196F3 style Encoder fill:#e8f5e9,stroke:#4CAF50 style Decoder fill:#fff3e0,stroke:#FF9800
ACTEncoderLayer 内部结构
[S, B, 512]"] --> NORM1{"Pre-Norm?"} NORM1 -->|"是"| LN1_PRE["LayerNorm"] NORM1 -->|"否"| SA_POST["直接进入"] LN1_PRE --> SA["多头自注意力
(8头, dim=512)
Q = K = x + pos_embed
V = x"] SA_POST --> SA SA --> DROP1["Dropout"] --> ADD1["+ 残差连接"] ADD1 --> NORM2{"Pre-Norm?"} NORM2 -->|"是"| LN2_PRE["LayerNorm"] NORM2 -->|"否"| LN1_POST["LayerNorm → 残差"] LN2_PRE --> FFN["前馈网络
Linear(512→3200) → ReLU → Dropout
→ Linear(3200→512)"] LN1_POST --> FFN FFN --> DROP2["Dropout"] --> ADD2["+ 残差连接"] ADD2 --> OUT["输出
[S, B, 512]"] style IN fill:#e3f2fd,stroke:#2196F3 style OUT fill:#e8f5e9,stroke:#4CAF50
ACTDecoderLayer 内部结构
[chunk_size, B, 512]"] --> SA["多头自注意力
Q = K = x + decoder_pos_embed
V = x"] SA --> DROP1["Dropout + 残差"] DROP1 --> CA["多头交叉注意力
Q = x + decoder_pos_embed
K = encoder_out + encoder_pos_embed
V = encoder_out"] CA --> DROP2["Dropout + 残差"] DROP2 --> FFN["前馈网络
Linear(512→3200) → ReLU → Dropout
→ Linear(3200→512)"] FFN --> DROP3["Dropout + 残差"] DROP3 --> OUT["输出
[chunk_size, B, 512]"] ENCODER["编码器输出
[S, B, 512]"] -->|"Key, Value"| CA style IN fill:#e3f2fd,stroke:#2196F3 style OUT fill:#e8f5e9,stroke:#4CAF50 style ENCODER fill:#fff3e0,stroke:#FF9800
注意: 原始 ACT 代码中
n_decoder_layers设为 7,但由于 代码 bug,实际只有第一层被使用。LeRobot 实现中默认设为 1 以匹配原始行为。
2.4 动作头
动作头是一个简单的线性层,将 Transformer 解码器的输出映射到动作空间。
[B, chunk_size, 512]"] --> LINEAR["nn.Linear
512 → action_dim"] LINEAR --> ACTIONS["预测动作
[B, chunk_size, action_dim]"] end style ActionHead fill:#fce4ec,stroke:#E91E63
- 没有额外的非线性激活或归一化层,直接线性投影
- 输出维度等于动作空间维度 (如双臂 ALOHA 为 14)
- 一次性输出
chunk_size步动作 (默认 100 步)
3. 训练流水线
训练时存在两条路径:使用 VAE 和 不使用 VAE。核心区别在于隐变量 z 的来源。
[B, N_cam, C, H, W]"] STATE["机器人状态
[B, state_dim]"] ENV["环境状态 (可选)
[B, env_dim]"] GT_ACT["真实动作序列
[B, chunk_size, action_dim]"] end subgraph VisionEnc["视觉特征提取"] IMG -->|"逐相机处理"| RESNET["ResNet18
→ layer4 特征图"] RESNET --> CONV["1x1 Conv 投影"] CONV --> FLAT["展平 + 2D 位置编码
[h*w*N_cam, B, 512]"] end subgraph VAEPath["VAE 路径 (use_vae=True)"] STATE -->|"状态嵌入"| VAE_IN["VAE 编码器输入
[CLS, state, actions]"] GT_ACT -->|"动作嵌入"| VAE_IN VAE_IN --> VAE_ENC["VAE Transformer
编码器 (4层)"] VAE_ENC -->|"[CLS] → Linear"| PARAMS["mu, log(sigma^2)"] PARAMS --> SAMPLE["重参数化采样
z ~ N(mu, sigma)"] end subgraph NoVAEPath["非 VAE 路径 (use_vae=False)"] ZERO["z = 零向量
[B, 32]"] end subgraph MainTransformer["主 Transformer"] Z_PROJ["隐变量投影
Linear: 32 → 512"] STATE2["状态投影
Linear: state_dim → 512"] ENV2["环境投影 (可选)
Linear: env_dim → 512"] ENC["Transformer 编码器 (4层)
输入: [z, state, (env), img_features]"] DEC["Transformer 解码器 (1层)
查询: 全零 + 可学习位置嵌入"] HEAD["动作头
Linear: 512 → action_dim"] end subgraph LossCalc["损失计算"] L1["L1 损失
|actions_pred - actions_gt|
(带 action_is_pad 掩码)"] KL["KL 散度损失 (仅 VAE)
-0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)"] TOTAL["总损失
L1 + kl_weight * KL
(kl_weight 默认 = 10.0)"] end SAMPLE -->|"z"| Z_PROJ ZERO -->|"z"| Z_PROJ STATE --> STATE2 ENV --> ENV2 Z_PROJ --> ENC STATE2 --> ENC ENV2 --> ENC FLAT --> ENC ENC --> DEC --> HEAD HEAD -->|"预测动作"| L1 GT_ACT -->|"真实动作"| L1 PARAMS -->|"mu, log_sigma_x2"| KL L1 --> TOTAL KL --> TOTAL style DataInput fill:#e8f4fd,stroke:#2196F3 style VisionEnc fill:#fff3e0,stroke:#FF9800 style VAEPath fill:#f3e5f5,stroke:#9C27B0 style NoVAEPath fill:#e0f2f1,stroke:#009688 style MainTransformer fill:#e8f5e9,stroke:#4CAF50 style LossCalc fill:#ffebee,stroke:#f44336
损失函数细节:
| 损失项 | 公式 | 说明 |
|---|---|---|
| L1 重建损失 | mean(\|actions_pred - actions_gt\| * ~action_is_pad) |
主损失,填充位置被掩码排除 |
| KL 散度 (仅 VAE) | -0.5 * mean(sum(1 + log(sigma^2) - mu^2 - sigma^2)) |
约束隐变量接近 N(0, I) |
| 总损失 | L1 + kl_weight * KL |
kl_weight 默认为 10.0 |
4. 推理流水线
推理时 不使用 VAE 编码器,隐变量 z 设为零向量。模型通过单次前向传播输出整个动作块,无需迭代去噪。
[B, N_cam, C, H, W]"] STATE["机器人状态
[B, state_dim]"] ENV["环境状态 (可选)"] end subgraph SinglePass["单次前向传播"] RESNET["ResNet18 特征提取
+ 1x1 Conv 投影"] ZERO_Z["零向量隐变量
z = 0, [B, 32]"] ENC["Transformer 编码器
[z, state, (env), img] → 编码输出"] DEC["Transformer 解码器
全零查询 + 位置嵌入 → 交叉注意力"] HEAD["动作头 → [B, chunk_size, action_dim]"] end subgraph ActionSelect["动作选择策略"] CHUNK["完整动作块
[B, chunk_size, action_dim]"] OPTION_A["策略 A: 动作队列
取前 n_action_steps 步
逐步弹出执行"] OPTION_B["策略 B: 时间集成
exponential weighting
多次预测加权平均"] end subgraph Output["输出"] SINGLE_ACT["单步动作
[B, action_dim]"] end IMG --> RESNET --> ENC STATE --> ENC ENV --> ENC ZERO_Z --> ENC ENC --> DEC --> HEAD --> CHUNK CHUNK --> OPTION_A --> SINGLE_ACT CHUNK --> OPTION_B --> SINGLE_ACT style InferInput fill:#e8f4fd,stroke:#2196F3 style SinglePass fill:#e8f5e9,stroke:#4CAF50 style ActionSelect fill:#fff3e0,stroke:#FF9800 style Output fill:#fce4ec,stroke:#E91E63
动作选择策略对比
w_i = exp(-coeff * i)"] E4["输出加权后的动作"] E1 --> E2 --> E3 --> E4 end style QueueMode fill:#e3f2fd,stroke:#2196F3 style EnsembleMode fill:#fff9c4,stroke:#FFC107
5. 关键超参数表
| 参数 | 默认值 | 说明 |
|---|---|---|
chunk_size |
100 | 动作分块大小,一次预测的动作步数 |
n_action_steps |
100 | 每次推理实际执行的动作步数 (<=chunk_size) |
n_obs_steps |
1 | 观测帧数 (当前仅支持 1) |
dim_model |
512 | Transformer 隐藏维度 |
n_heads |
8 | 多头注意力头数 |
dim_feedforward |
3200 | FFN 中间层维度 |
feedforward_activation |
relu |
FFN 激活函数 |
n_encoder_layers |
4 | Transformer 编码器层数 |
n_decoder_layers |
1 | Transformer 解码器层数 (原代码 bug 导致实际为 1) |
use_vae |
True | 是否启用 VAE 路径 |
latent_dim |
32 | VAE 隐变量维度 |
n_vae_encoder_layers |
4 | VAE 编码器 Transformer 层数 |
kl_weight |
10.0 | KL 散度损失权重 |
dropout |
0.1 | Transformer 各层 Dropout 率 |
vision_backbone |
resnet18 |
视觉编码器骨干网络 |
pretrained_backbone_weights |
ResNet18_Weights.IMAGENET1K_V1 |
预训练权重 |
replace_final_stride_with_dilation |
False | 是否用膨胀卷积替代 layer4 步幅 |
pre_norm |
False | 是否使用 Pre-Norm (默认 Post-Norm) |
temporal_ensemble_coeff |
None | 时间集成系数 (None 表示不启用) |
optimizer_lr |
1e-5 | 学习率 |
optimizer_lr_backbone |
1e-5 | 视觉骨干网络学习率 |
optimizer_weight_decay |
1e-4 | 权重衰减 |
optimizer_grad_clip_norm |
1.0 | 梯度裁剪范数 |
6. 关键源文件表
| 组件 | 类名 | 文件 |
|---|---|---|
| 策略封装 | ACTPolicy |
lerobot/policies/act/modeling_act.py:41 |
| 核心模型 | ACT |
lerobot/policies/act/modeling_act.py:252 |
| Transformer 编码器 | ACTEncoder |
lerobot/policies/act/modeling_act.py:509 |
| 编码器层 | ACTEncoderLayer |
lerobot/policies/act/modeling_act.py:528 |
| Transformer 解码器 | ACTDecoder |
lerobot/policies/act/modeling_act.py:567 |
| 解码器层 | ACTDecoderLayer |
lerobot/policies/act/modeling_act.py:590 |
| 时间集成 | ACTTemporalEnsembler |
lerobot/policies/act/modeling_act.py:161 |
| 2D 位置编码 | ACTSinusoidalPositionEmbedding2d |
lerobot/policies/act/modeling_act.py:680 |
| 1D 位置编码 | create_sinusoidal_pos_embedding |
lerobot/policies/act/modeling_act.py:662 |
| 配置 | ACTConfig |
lerobot/policies/act/configuration_act.py:25 |