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["相机图像
(多视角)
[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 正弦位置编码。

graph LR subgraph Backbone["ResNet18 视觉编码器"] IMG["输入图像
[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

关键设计细节:

2.2 VAE 编码器 (可选)

VAE 编码器是 ACT 的核心创新之一。训练时,它将真实动作序列编码为隐变量 z,捕获动作分布中的多模态性;推理时跳过 VAE 编码器,直接使用零向量作为隐变量。

graph TB subgraph VAEInput["VAE 编码器输入"] CLS["[CLS] 嵌入
(可学习, 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 的设计意图:

2.3 Transformer 编码器-解码器

这是 ACT 的核心推理模块。Transformer 编码器融合隐变量、状态嵌入和图像特征;Transformer 解码器通过交叉注意力从编码器输出中生成动作序列。

graph TB subgraph EncoderInput["编码器输入 Token 序列"] direction LR T_Z["隐变量 Token
(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 内部结构

graph TB IN["输入 x
[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 内部结构

graph TB IN["输入 x
[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 解码器的输出映射到动作空间。

graph LR subgraph ActionHead["动作回归头"] DEC_OUT["解码器输出
[B, chunk_size, 512]"] --> LINEAR["nn.Linear
512 → action_dim"] LINEAR --> ACTIONS["预测动作
[B, chunk_size, action_dim]"] end style ActionHead fill:#fce4ec,stroke:#E91E63

3. 训练流水线

训练时存在两条路径:使用 VAE不使用 VAE。核心区别在于隐变量 z 的来源。

graph TB subgraph DataInput["训练数据"] IMG["相机图像
[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 设为零向量。模型通过单次前向传播输出整个动作块,无需迭代去噪。

graph TB subgraph InferInput["推理输入"] IMG["相机图像
[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

动作选择策略对比

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["指数加权平均
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