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 编码器-解码器生成最终的动作序列。推理时仅需单次前向传播，无需迭代去噪。
graph TB
subgraph Input["输入"]
IMG["相机图像<br/>(多视角)<br/>[B, C, H, W]"]
STATE["机器人状态<br/>(关节角度)<br/>[B, state_dim]"]
ENV["环境状态 (可选)<br/>[B, env_dim]"]
ACT_IN["动作序列 (仅训练)<br/>[B, chunk_size, action_dim]"]
end
subgraph VisionBackbone["ResNet18 视觉编码器"]
RES["ResNet18<br/>(预训练 ImageNet 权重)<br/>FrozenBatchNorm2d"]
PROJ_IMG["1x1 卷积投影<br/>512 → 512"]
POS2D["2D 正弦位置编码"]
end
subgraph VAE["VAE 编码器 (可选, 仅训练)"]
CLS["[CLS] Token"]
VAE_ENC["VAE Transformer 编码器<br/>(4层)"]
LATENT["隐变量采样<br/>z ~ N(mu, sigma)"]
end
subgraph TransformerCore["Transformer 编码器-解码器"]
ENC["Transformer 编码器<br/>(4层, 8头, dim=512)"]
DEC["Transformer 解码器<br/>(1层, 交叉注意力)"]
end
subgraph Output["输出"]
HEAD["动作回归头<br/>(Linear: 512 → action_dim)"]
ACT_OUT["预测动作序列<br/>[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<br/>[B, latent_dim]"| ENC
ENC -->|"编码器输出<br/>[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 正弦位置编码。
graph LR
subgraph Backbone["ResNet18 视觉编码器"]
IMG["输入图像<br/>[B, 3, H, W]"] --> CONV1["conv1 + bn1 + relu<br/>+ maxpool"]
CONV1 --> L1["layer1<br/>(2 BasicBlock)"]
L1 --> L2["layer2<br/>(2 BasicBlock)"]
L2 --> L3["layer3<br/>(2 BasicBlock)"]
L3 --> L4["layer4<br/>(2 BasicBlock)<br/>可选: 膨胀卷积替代步幅"]
end
subgraph Projection["特征投影"]
L4 -->|"feature_map<br/>[B, 512, h, w]"| CONV1x1["Conv2d 1x1<br/>512 → 512"]
CONV1x1 -->|"[B, 512, h, w]"| RESHAPE["展平为序列<br/>[h*w, B, 512]"]
end
subgraph PosEmbed["2D 正弦位置编码"]
RESHAPE --> ADD["+ pos_embed"]
PE["ACTSinusoidalPosition<br/>Embedding2d<br/>(dim=256)"] -->|"[1, 512, h, w]<br/>→ [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 编码器，直接使用零向量作为隐变量。
graph TB
subgraph VAEInput["VAE 编码器输入"]
CLS["[CLS] 嵌入<br/>(可学习, nn.Embedding)<br/>[B, 1, 512]"]
RS["机器人状态嵌入<br/>(Linear: state_dim → 512)<br/>[B, 1, 512]"]
ACT_SEQ["动作序列嵌入<br/>(Linear: action_dim → 512)<br/>[B, chunk_size, 512]"]
end
subgraph VAEEncoder["VAE Transformer 编码器 (4层)"]
CONCAT["拼接: [CLS, state, actions]<br/>[B, chunk_size+2, 512]"]
SINPE["固定正弦位置编码"]
PAD_MASK["填充掩码<br/>(action_is_pad)"]
TENC["ACTEncoder<br/>(4层 ACTEncoderLayer)<br/>+ LayerNorm"]
end
subgraph LatentSample["隐变量采样"]
CLS_OUT["[CLS] Token 输出<br/>[B, 512]"]
LINEAR_P["线性投影<br/>512 → 64<br/>(latent_dim * 2)"]
MU["mu<br/>[B, 32]"]
LOG_S["log(sigma^2)<br/>[B, 32]"]
REPARAM["重参数化技巧<br/>z = mu + sigma * eps<br/>eps ~ N(0, 1)"]
Z["隐变量 z<br/>[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 解码器通过交叉注意力从编码器输出中生成动作序列。
graph TB
subgraph EncoderInput["编码器输入 Token 序列"]
direction LR
T_Z["隐变量 Token<br/>(Linear: 32 → 512)<br/>[1, B, 512]"]
T_S["机器人状态 Token<br/>(Linear: state_dim → 512)<br/>[1, B, 512]"]
T_E["环境状态 Token (可选)<br/>(Linear: env_dim → 512)<br/>[1, B, 512]"]
T_IMG["图像特征 Token<br/>(每相机 h*w 个)<br/>[h*w*N_cam, B, 512]"]
end
subgraph Encoder["ACTEncoder (4层)"]
ENC_PE["位置编码<br/>1D token: 可学习 Embedding<br/>图像 token: 2D 正弦编码"]
ENC_LAYERS["ACTEncoderLayer x 4<br/>自注意力 + FFN<br/>(Post-Norm)"]
ENC_NORM["LayerNorm (若 Pre-Norm)"]
end
subgraph Decoder["ACTDecoder (1层)"]
DEC_QUERY["解码器查询<br/>全零初始化<br/>[chunk_size, B, 512]"]
DEC_PE["可学习位置嵌入<br/>(nn.Embedding)<br/>[chunk_size, 1, 512]"]
DEC_LAYER["ACTDecoderLayer x 1<br/>自注意力 → 交叉注意力 → 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 -->|"编码器输出<br/>[S, B, 512]"| DEC_LAYER
DEC_QUERY --> DEC_LAYER
DEC_PE --> DEC_LAYER
DEC_LAYER --> DEC_NORM
DEC_NORM -->|"解码器输出<br/>[chunk_size, B, 512]"| OUT["→ 动作头"]
style EncoderInput fill:#e8f4fd,stroke:#2196F3
style Encoder fill:#e8f5e9,stroke:#4CAF50
style Decoder fill:#fff3e0,stroke:#FF9800
ACTEncoderLayer 内部结构#
graph TB
IN["输入 x<br/>[S, B, 512]"] --> NORM1{"Pre-Norm?"}
NORM1 -->|"是"| LN1_PRE["LayerNorm"]
NORM1 -->|"否"| SA_POST["直接进入"]
LN1_PRE --> SA["多头自注意力<br/>(8头, dim=512)<br/>Q = K = x + pos_embed<br/>V = x"]
SA_POST --> SA
SA --> DROP1["Dropout"] --> ADD1["+ 残差连接"]
ADD1 --> NORM2{"Pre-Norm?"}
NORM2 -->|"是"| LN2_PRE["LayerNorm"]
NORM2 -->|"否"| LN1_POST["LayerNorm → 残差"]
LN2_PRE --> FFN["前馈网络<br/>Linear(512→3200) → ReLU → Dropout<br/>→ Linear(3200→512)"]
LN1_POST --> FFN
FFN --> DROP2["Dropout"] --> ADD2["+ 残差连接"]
ADD2 --> OUT["输出<br/>[S, B, 512]"]
style IN fill:#e3f2fd,stroke:#2196F3
style OUT fill:#e8f5e9,stroke:#4CAF50
ACTDecoderLayer 内部结构#
graph TB
IN["输入 x<br/>[chunk_size, B, 512]"] --> SA["多头自注意力<br/>Q = K = x + decoder_pos_embed<br/>V = x"]
SA --> DROP1["Dropout + 残差"]
DROP1 --> CA["多头交叉注意力<br/>Q = x + decoder_pos_embed<br/>K = encoder_out + encoder_pos_embed<br/>V = encoder_out"]
CA --> DROP2["Dropout + 残差"]
DROP2 --> FFN["前馈网络<br/>Linear(512→3200) → ReLU → Dropout<br/>→ Linear(3200→512)"]
FFN --> DROP3["Dropout + 残差"]
DROP3 --> OUT["输出<br/>[chunk_size, B, 512]"]
ENCODER["编码器输出<br/>[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 解码器的输出映射到动作空间。
graph LR
subgraph ActionHead["动作回归头"]
DEC_OUT["解码器输出<br/>[B, chunk_size, 512]"] --> LINEAR["nn.Linear<br/>512 → action_dim"]
LINEAR --> ACTIONS["预测动作<br/>[B, chunk_size, action_dim]"]
end
style ActionHead fill:#fce4ec,stroke:#E91E63
- 没有额外的非线性激活或归一化层，直接线性投影
- 输出维度等于动作空间维度 (如双臂 ALOHA 为 14)
- 一次性输出
chunk_size步动作 (默认 100 步)
3. 训练流水线#
训练时存在两条路径：使用 VAE 和 不使用 VAE。核心区别在于隐变量 z 的来源。
graph TB
subgraph DataInput["训练数据"]
IMG["相机图像<br/>[B, N_cam, C, H, W]"]
STATE["机器人状态<br/>[B, state_dim]"]
ENV["环境状态 (可选)<br/>[B, env_dim]"]
GT_ACT["真实动作序列<br/>[B, chunk_size, action_dim]"]
end
subgraph VisionEnc["视觉特征提取"]
IMG -->|"逐相机处理"| RESNET["ResNet18<br/>→ layer4 特征图"]
RESNET --> CONV["1x1 Conv 投影"]
CONV --> FLAT["展平 + 2D 位置编码<br/>[h*w*N_cam, B, 512]"]
end
subgraph VAEPath["VAE 路径 (use_vae=True)"]
STATE -->|"状态嵌入"| VAE_IN["VAE 编码器输入<br/>[CLS, state, actions]"]
GT_ACT -->|"动作嵌入"| VAE_IN
VAE_IN --> VAE_ENC["VAE Transformer<br/>编码器 (4层)"]
VAE_ENC -->|"[CLS] → Linear"| PARAMS["mu, log(sigma^2)"]
PARAMS --> SAMPLE["重参数化采样<br/>z ~ N(mu, sigma)"]
end
subgraph NoVAEPath["非 VAE 路径 (use_vae=False)"]
ZERO["z = 零向量<br/>[B, 32]"]
end
subgraph MainTransformer["主 Transformer"]
Z_PROJ["隐变量投影<br/>Linear: 32 → 512"]
STATE2["状态投影<br/>Linear: state_dim → 512"]
ENV2["环境投影 (可选)<br/>Linear: env_dim → 512"]
ENC["Transformer 编码器 (4层)<br/>输入: [z, state, (env), img_features]"]
DEC["Transformer 解码器 (1层)<br/>查询: 全零 + 可学习位置嵌入"]
HEAD["动作头<br/>Linear: 512 → action_dim"]
end
subgraph LossCalc["损失计算"]
L1["L1 损失<br/>|actions_pred - actions_gt|<br/>(带 action_is_pad 掩码)"]
KL["KL 散度损失 (仅 VAE)<br/>-0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)"]
TOTAL["总损失<br/>L1 + kl_weight * KL<br/>(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 设为零向量。模型通过单次前向传播输出整个动作块，无需迭代去噪。
graph TB
subgraph InferInput["推理输入"]
IMG["相机图像<br/>[B, N_cam, C, H, W]"]
STATE["机器人状态<br/>[B, state_dim]"]
ENV["环境状态 (可选)"]
end
subgraph SinglePass["单次前向传播"]
RESNET["ResNet18 特征提取<br/>+ 1x1 Conv 投影"]
ZERO_Z["零向量隐变量<br/>z = 0, [B, 32]"]
ENC["Transformer 编码器<br/>[z, state, (env), img] → 编码输出"]
DEC["Transformer 解码器<br/>全零查询 + 位置嵌入 → 交叉注意力"]
HEAD["动作头 → [B, chunk_size, action_dim]"]
end
subgraph ActionSelect["动作选择策略"]
CHUNK["完整动作块<br/>[B, chunk_size, action_dim]"]
OPTION_A["策略 A: 动作队列<br/>取前 n_action_steps 步<br/>逐步弹出执行"]
OPTION_B["策略 B: 时间集成<br/>exponential weighting<br/>多次预测加权平均"]
end
subgraph Output["输出"]
SINGLE_ACT["单步动作<br/>[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
动作选择策略对比#
graph LR
subgraph QueueMode["动作队列模式 (默认)"]
direction TB
P1["模型预测 chunk_size=100 步"]
P2["截取前 n_action_steps 步"]
P3["放入队列, 逐步弹出执行"]
P4["队列空后重新预测"]
P1 --> P2 --> P3 --> P4
end
subgraph EnsembleMode["时间集成模式 (temporal_ensemble)"]
direction TB
E1["每步都调用模型"]
E2["多次预测的重叠部分"]
E3["指数加权平均<br/>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 |