ACT（Action Chunking with Transformers）+ ALOHA：论文与开源代码一体化解读

一、这篇工作要解决什么问题

论文关注的是真实世界中的精细双臂操作（fine-grained bimanual manipulation）：例如打开半透明调料杯、把电池插入槽位、穿线/绑扎等。

这类任务难点集中在：

• 毫米级精度 + 双手协同：动作链很长，任何一步的小误差都会把系统推离可恢复状态。
• 强接触与非线性：撬开、捏合、撕开等接触过程很难用解析模型精确建模。
• 闭环视觉反馈：需要依赖多视角 RGB 输入来不断纠错，而不是一次性规划。

因此，作者提出一个组合式答案：

1. 低成本可复现的双臂遥操作硬件（ALOHA）：把“高质量示范数据采集”变成现实。
2. 面向长时序模仿学习的算法（ACT）：用“动作分块预测 + Transformer + CVAE”降低模仿学习的误差累积问题。

二、论文核心思路（抓住 ACT 的两个关键词）

1）Action Chunking：把“单步预测”改成“预测一段动作”

经典 BC（Behavior Cloning）通常做的是：给定观测 (o_t)，预测 (a_t)。
在精细操作里，一旦预测 (a_t) 有小偏差，就会导致状态漂移、越走越偏（compounding error）。

ACT 的关键做法是：在同一时刻，直接预测一个长度为 (K) 的动作序列：

• 输出：([a_t, a_{t+1}, …, a_{t+K-1}])
• 直觉：把有效控制地平线“缩短”，减少“每一步都可能犯错”的次数。

2）CVAE + Transformer：用“风格变量 z”建模示范的多模态

同一个观测下，人类示范可能有不同“风格”（速度、路径偏好、接触策略）。
ACT 用 Conditional VAE：

• Encoder（训练时用）：把 (qpos) + 动作序列编码成潜变量 (z)
• Decoder（推理时用）：丢弃 encoder，把 (z) 设为先验均值（代码里是 0），由 decoder 生成动作 chunk

实践含义：在数据不多的情况下，用生成式建模吸收示范多样性，提升泛化稳定性。

三、论文给出的开源仓库

论文项目主页给出了两套核心代码库：ALOHA（遥操作/采集）与 ACT+Sim（算法+仿真）$TRAE_REF

• ALOHA Codebase：https://github.com/tonyzhaozh/aloha
• ACT+Sim Codebase：https://github.com/tonyzhaozh/act

本文的代码解读主要围绕 act 仓库展开，并在需要时指向 aloha 里的数据采集/真机运行入口。

四、代码仓库整体结构（先看 3 个入口文件）

1）ACT+Sim：tonyzhaozh/act

README.md 已经把关键入口列得很清楚（建议按这个顺序读）：

• imitate_episodes.py：训练与评估主脚本（最核心）
• policy.py：ACTPolicy 的训练/推理接口封装（loss 在这里）
• detr/：ACT 的模型定义（基于 DETR 改）
• record_sim_episodes.py / scripted_policy.py：仿真环境生成数据
• utils.py：数据集（HDF5）读取、归一化、随机切片采样等
• sim_env.py / ee_sim_env.py：MuJoCo + dm_control 环境封装

2）ALOHA：tonyzhaozh/aloha

这个仓库偏“系统工程”：

• aloha_scripts/：遥操作、录制、回放等脚本
• launch/：ROS launch（4 相机 + 4 机械臂）
• config/：机器人端口/配置

它的作用很明确：解决“高质量示范数据怎么采”。

五、论文到代码的对应关系（最重要的一节）

下面按论文描述的组件，把“论文概念 → 代码位置”列出来（你后续做复现/二次开发，基本就靠这张映射表）：

论文组件	代码入口	你该看什么
ACT 训练与评估流程	imitate_episodes.py	参数如何组织、怎么切 train/val、怎么保存 ckpt、eval loop 如何 rollout
动作分块（chunk）	imitate_episodes.py + policy.py	chunk_size → num_queries；推理时输出 chunk 后如何取当前步动作
Temporal ensembling（时序集成）	imitate_episodes.py	--temporal_agg 分支：把多次预测的 chunk 在同一时刻做指数加权融合
CVAE（含 KL）训练目标	policy.py	loss = l1 + kl_weight * kl；训练时只取 actions[:, :num_queries]
Transformer + CVAE 结构	detr/models/detr_vae.py	encoder：[CLS] + qpos + action_seq → z；decoder：图像特征 + proprio + latent → action chunk
数据集格式与采样策略	utils.py	HDF5 路径、随机切片 start_ts、padding 与 is_pad、归一化统计

六、关键实现细节（把论文话术落到代码行级）

1）“chunk_size”在代码里对应什么？

在 imitate_episodes.py 里：

• policy_config['num_queries'] = args['chunk_size']
• 这直接进入 DETRVAE(num_queries=...)
• 推理时，ACTPolicy 会输出形状大致为：(batch, num_queries, action_dim) 的动作序列

你可以把它理解为：DETR 的 object queries 被改造成 action queries。

2）CVAE 的“z”是怎么做出来的？

在 detr/models/detr_vae.py 的 DETRVAE.forward()：

• 训练阶段（actions is not None）：
  1. encoder_action_proj(actions) 把动作序列投到 hidden_dim
  2. encoder_joint_proj(qpos) 把关节状态投到 hidden_dim
  3. 拼成 encoder_input = [CLS, qpos, action_seq]
  4. Transformer encoder 输出取 CLS token，过 latent_proj 得到 mu/logvar
  5. reparametrize() 采样 z，再投回 hidden_dim 作为 decoder 的 latent_input
• 推理阶段：
  • latent_sample = zeros(...)，也就是把 (z) 固定到先验均值

这和论文里的描述是严丝合缝的：训练时用 encoder 学到“风格”，推理时用固定 prior。

3）训练 loss 是什么？哪里实现？

在 policy.py 的 ACTPolicy.__call__()：

• l1 = L1(actions, a_hat)（并用 ~is_pad mask 掉 padding）
• kl = kl_divergence(mu, logvar)
• 总损失：loss = l1 + kl_weight * kl

这也是你后续做 ablation/替换 loss（例如换成 Huber、加速度正则、contact-aware loss）的关键切入点。

4）Temporal ensembling 到底怎么做？

在 imitate_episodes.py 的 eval loop：

• 先把每次预测到的 all_actions 写到 all_time_actions[t, t:t+num_queries]
• 然后取“所有对当前时刻 t 有预测的 chunk”，做指数衰减加权平均：
  • exp_weights = exp(-k * arange(N))，再归一化
  • 最终 raw_action = sum(actions_for_curr_step * exp_weights)

这个实现非常直观：越新的 chunk 权重越大，起到“平滑 + 抗抖动”的作用。

七、怎么跑通或怎么复现（以仿真为最低成本路径）

我这里不直接在你的机器上跑训练（依赖 MuJoCo/torch/conda 环境），但把“最短可复现路径”整理清楚，确保你接下来可以按步骤跑通。

1）生成仿真数据（脚本策略）

在 act 仓库里，按 README：

python3 record_sim_episodes.py \
  --task_name sim_transfer_cube_scripted \
  --dataset_dir <data save dir> \
  --num_episodes 50

生成的数据是 episode_*.hdf5，结构在 utils.py 能看到：

• /observations/qpos
• /observations/qvel
• /observations/images/<cam_name>
• /action
• root.attrs['sim']

2）训练 ACT

python3 imitate_episodes.py \
  --task_name sim_transfer_cube_scripted \
  --ckpt_dir <ckpt dir> \
  --policy_class ACT \
  --kl_weight 10 \
  --chunk_size 100 \
  --hidden_dim 512 \
  --batch_size 8 \
  --dim_feedforward 3200 \
  --num_epochs 2000 \
  --lr 1e-5 \
  --seed 0

3）评估 + 视频

在相同命令基础上加 --eval。
若要启用 temporal ensembling，加 --temporal_agg，对应的实现就在 imitate_episodes.py。

八、优点、局限与我的理解（偏工程视角）

优点（为什么 ACT 适合精细操作）

• 把“长时序误差累积”变成“短时序 chunk 回归”，这是最关键的工程优势。
• 模型结构高度可复用：DETR 改造 + CVAE，代码层面很好改（换 backbone、加传感器、加语言条件等）。
• 数据管线非常朴素：HDF5 + 随机切片 + padding mask，利于真实数据快速迭代。

局限（你复现时最容易踩的坑）

• 对高质量示范很敏感：这也是为什么论文必须配 ALOHA。
• chunk_size/kl_weight/训练轮数很关键：README 里也强调“loss plateau 之后继续训，动作会更顺滑”。
• 视觉分辨率与多相机拼接：当前实现把多相机特征在宽度维 concat（torch.cat(all_cam_features, axis=3)），对显存/速度有影响，后续扩展更多相机要注意。

九、适合继续深入的阅读入口（按优先级）

1. imitate_episodes.py：先把训练/评估主流程读通
2. policy.py：搞清楚训练时输入输出与 loss
3. detr/models/detr_vae.py：把 CVAE 的 encoder/decoder 对上论文图
4. utils.py：看数据怎么被随机切片、padding、归一化
5. aloha/aloha_scripts：当你要上真机采数据时再深入