世界模型 | WAM 环境搭建：AHA-WAM、Wan 权重与 RoboTwin 数据格式一次理清

转载自公众号：敢敢AUTOHUB

0. 简介

AHA-WAM 是一种用于机器人操作的异步世界-动作模型。它把较慢的视频驱动世界规划器和快速的动作驱动执行器拆开，再通过观察引导的上下文路由机制，把复用的世界上下文和最新机器人状态重新对齐。这里的工作和 FAST-WAM 很接近，但 AHA-WAM 的代码、配置和数据入口都有一些更新，尤其是 RoboTwin 数据部分，如果只按“下载数据然后改路径”的直觉去跑，很容易卡在 LeRobot 元数据、Parquet 文件和 dataset_stats.json 这些细节上。

这篇博客主要记录三件事：第一，AHA-WAM 环境和 Wan Base 组件如何放置；第二，RoboTwin 预处理数据到底应该长什么样；第三，如果手里只有上游 RoboTwin 原始数据，应该怎么转换成当前训练加载器能读的 LeRobot 格式。后半部分会重点解释那段容易看懵的配置说明，也会把原始 RoboTwin 格式、LeRobot 目标格式、字段映射和转换优化建议放在一起看。

1. 先看仓库结构，知道配置和入口在哪里

1.1 配置入口

AHA-WAM 仓库的结构比较清晰，训练、模型、数据、部署各自放在固定目录下。这里的关键是，后面所有路径覆盖都围绕 configs/ 展开：数据入口在 configs/data/robotwin.yaml，任务入口在 configs/task/robotwin_ahawam*.yaml，训练总入口在 configs/train.yaml。如果只改某一个 YAML，而任务配置里又覆盖了同名字段，最终生效路径可能不是你以为的那个，所以更推荐第一次跑时用 Hydra 命令行覆盖关键路径。

1.2 目录视图

AHA-WAM/
├── assets/                   # README figures
├── configs/
│   ├── data/                 # Dataset configs
│   ├── model/                # AHA-WAM model configs
│   ├── task/                 # Training task configs
│   ├── deploy.yml            # Hydra config for real-robot deployment
│   ├── sim_robotwin.yaml     # RoboTwin evaluation config
│   └── train.yaml            # Base training config
├── scripts/
│   ├── train.py
│   ├── train_zero1.sh        # DeepSpeed ZeRO-1 training entrypoint
│   ├── train_zero2.sh        # DeepSpeed ZeRO-2 training entrypoint
│   ├── preprocess_action_dit_backbone.py
│   └── precompute_text_embeds.py
├── experiments/
│   └── robotwin/             # RoboTwin evaluation entrypoints and policy adapter
├── deploy/                   # Real-robot server/client deployment examples
├── src/ahawam/               # Python package
├── data/                     # Local datasets
├── runs/                     # Training outputs
└── evaluate_results/         # Evaluation outputs

直觉理解：可以把 configs/data/robotwin.yaml 理解成“数据长什么样”，把 configs/task/robotwin_ahawam.yaml 理解成“这次训练用哪份数据、跑多少步、输出到哪里”。如果两边都写了路径，任务配置里的覆盖项更接近真正启动训练时用到的值。

2. Python 环境：先把训练依赖装齐

2.1 安装顺序

基础环境使用 Python 3.10。本仓库 pyproject.toml 已经写了 accelerate、deepspeed、datasets、pyarrow、av、torchcodec、imageio、modelscope 等依赖，所以安装顺序可以先装 PyTorch，再 pip install -e . 安装仓库本身。这里的 ffmpeg 很重要，因为 LeRobot 视频数据会涉及 mp4 编码和解码，缺了它以后，数据目录看起来完整，真正取视频帧时仍然会失败。更实际一点说，很多数据问题表面上像字段错误，最后查下来只是运行环境缺了视频后端。

2.2 环境命令

conda create -n ahawam python=3.10 -y
conda activate ahawam
pip install -U pip setuptools wheel
pip install torch==2.7.1+cu128 torchvision==0.22.1+cu128 
  --extra-index-url <https://download.pytorch.org/whl/cu128>
pip install -e .
conda install -c conda-forge ffmpeg -y

进一步看，如果你是在服务器上跑，最好启动训练前确认当前 shell 真的进了 ahawam 环境。因为只用系统 /bin/python 的话，可能连 datasets 都找不到；这种情况下，看 meta/info.json 没问题，不代表训练入口能正常实例化数据集。排查时可以先看 which python 和 python -c "import datasets"，把环境问题和数据问题分开。

3. Wan Base 组件：文件名比下载源更重要

3.1 加载器看什么

Wan Base 的加载方式类似 DiffSynth：如果文件不存在，首次运行时会尝试自动下载。默认下载源是 ModelScope，文件会缓存到 ./checkpoints/。在 configs/model/ahawam.yaml 中，当 redirect_common_files: true 时，文本编码器和 VAE 会被重定向到 DiffSynth-Studio/Wan-Series-Converted-Safetensors。这里的关键是，加载器不是只看“你下载过某个仓库”，而是看本地路径下是否存在它期待的文件模式。

3.2 期望文件

解读：默认 AHA-WAM 配置会寻找下面这些内容：

Wan-AI/Wan2.2-TI2V-5B/diffusion_pytorch_model*.safetensors
DiffSynth-Studio/Wan-Series-Converted-Safetensors/models_t5_umt5-xxl-enc-bf16.safetensors
DiffSynth-Studio/Wan-Series-Converted-Safetensors/Wan2.2_VAE.safetensors
Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl/

3.3 手动下载

解读：如果用 ModelScope 手动对齐，可以这样下载：

pip install -U modelscope

MODEL_ROOT=checkpoints

python -m modelscope.cli.cli download 
  --model Wan-AI/Wan2.2-TI2V-5B 
  --local_dir "$MODEL_ROOT/Wan-AI/Wan2.2-TI2V-5B" 
  --include 
    'diffusion_pytorch_model*.safetensors' 
    'diffusion_pytorch_model.safetensors.index.json' 
    'config.json' 
    'configuration.json'

python -m modelscope.cli.cli download 
  --model DiffSynth-Studio/Wan-Series-Converted-Safetensors 
  --local_dir "$MODEL_ROOT/DiffSynth-Studio/Wan-Series-Converted-Safetensors" 
  --include 
    'models_t5_umt5-xxl-enc-bf16.safetensors' 
    'Wan2.2_VAE.safetensors'

python -m modelscope.cli.cli download 
  --model Wan-AI/Wan2.1-T2V-1.3B 
  --local_dir "$MODEL_ROOT/Wan-AI/Wan2.1-T2V-1.3B" 
  --include 'google/umt5-xxl/*'

Wan Base 组件下载结果

3.4 路径对齐

Hugging Face 上也可以使用 SereneC/wan-series-checkpoint 这类镜像完成下载。换句话说，ModelScope 和 Hugging Face 不是核心矛盾，核心矛盾在于本地目录要和加载器期望的 model_id/file_pattern 对齐。文件不在预期位置时，即使磁盘上已经有 safetensors，代码也会继续报找不到模型。这里建议把下载后的目录截图或 find checkpoints -maxdepth 3 输出保留下来，后面复现实验时会很省事。

Wan Base 路径对齐结果

4. ActionDiT：没有现成权重就从 Wan 初始化

4.1 为什么要生成

ActionDiT 这块如果没有现成的 ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt，需要先运行预处理脚本生成。这个步骤的作用不是重新训练动作模型，而是把 Wan 视频 DiT 的部分结构转换成 ActionDiT 可以使用的初始化权重。对于第一次复现实验的人来说，这一步很容易被当成“可选项”，但在从头训练 AHA-WAM 时，它实际上是训练配置里会引用的关键文件。

4.2 生成命令

python scripts/preprocess_action_dit_backbone.py 
  --model-config configs/model/ahawam.yaml 
  --output checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt 
  --device cuda 
  --dtype bfloat16

解读：然后在训练配置或命令行里把它传给：

model.action_dit_pretrained_path=checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt

工程价值：ActionDiT 初始化就像给执行器一个“会看视频 DiT 结构”的起点。直接随机初始化也许能跑，但训练成本和稳定性都会变差；用这个预处理脚本，相当于把已有 Wan 表征迁移到动作分支。

5. 数据准备：真正容易踩坑的是 RoboTwin

5.1 下载路线

LIBERO 和 RoboTwin 都可以走下载预处理数据的路线。LIBERO 解压后的目录是多个 LeRobot 子数据集；RoboTwin 解压后则应该得到一个内层 robotwin2.0/ 根目录。这里的根目录不是随便起的文件夹，而是 LeRobot 数据集根目录：里面必须同时包含 data/、meta/、videos/。如果你把 dataset_dirs 指到外层下载目录，或者指到某个 videos/chunk-*，加载器都会在读取 metadata 时失败。

# LIBERO
mkdir -p data/libero_mujoco3.3.2

huggingface-cli download yuanty/LIBERO-fastwam 
  --repo-type dataset 
  --local-dir data/libero_mujoco3.3.2 
  --include "*.tar.gz"

cd data/libero_mujoco3.3.2
for f in *.tar.gz; do
  tar -xzf "$f"
done
cd -

# RoboTwin
mkdir -p data/robotwin2.0

huggingface-cli download yuanty/robotwin2.0-fastwam 
  --repo-type dataset 
  --local-dir data/robotwin2.0 
  --include "robotwin2.0.tar.gz.part-*"

cd data/robotwin2.0
cat robotwin2.0.tar.gz.part-* | tar -xzf -
cd -

5.2 LIBERO 解压结果

解读：LIBERO 提取后的目录通常是：

data/libero_mujoco3.3.2/
├── libero_10_no_noops_lerobot/
├── libero_goal_no_noops_lerobot/
├── libero_object_no_noops_lerobot/
└── libero_spatial_no_noops_lerobot/

LIBERO 数据目录截图 1LIBERO 数据目录截图 2

5.3 RoboTwin 解压结果

解读：RoboTwin 提取后的有效目录应该类似：

data/robotwin2.0/
└── robotwin2.0/
    ├── data/
    ├── meta/
    └── videos/

RoboTwin 数据目录截图 1RoboTwin 数据目录截图 2

5.4 加载器真正读取的内容

这里要厘清的是：训练加载器不是直接读取 RoboTwin 上游的 HDF5，而是通过仓库内置 LeRobot 包装器读取这个目录协议。所以 videos/ 只是其中一部分，meta/info.json 和 data/chunk-*/episode_*.parquet 同样不可缺少。只有视频没有 parquet，代码不知道每一帧的动作、状态、任务索引；只有 parquet 没有视频，模型又拿不到三路视觉输入。

其中  对应 num_frames， 对应 action_video_freq_ratio。这条关系式解释了为什么 RoboTwin 配置里 num_frames=65 且 action_video_freq_ratio=8 时，模型不是读取 65 张视频帧，而是在动作高频序列里抽取较低频的视频上下文。视频分支和动作分支的时间尺度并不相同，这正是 WAM 类方法要显式处理的地方。

其中  是 global_sample_stride，fps 来自 meta/info.json。如果转换时随手写错 fps，采样时间点就会整体偏移；模型仍然能读到张量，但视频帧、动作和状态之间的物理时间关系已经变了。这个问题最麻烦的地方在于它不一定立刻报错，而是表现为训练变慢、动作预测发散或评估效果不稳定。

6. AHA说明到底在说什么

6.1 原始占位符

AHA 说明里让我们更新这些占位符。表面看它只是在说“把 /path/to/... 换成自己的路径”，但真正要理解的是路径背后的数据协议：一个路径指数据集根目录，另一个路径指归一化统计文件。把这两类路径混在一起，是第一次配置 RoboTwin 训练时最常见的困惑，也会导致报错信息看起来像数据损坏。

data.train.dataset_dirs: [/path/to/robotwin2.0]
data.val.dataset_dirs: [/path/to/robotwin2.0]
data.train.pretrained_norm_stats: /path/to/dataset_stats.json
data.val.pretrained_norm_stats: /path/to/dataset_stats.json

6.2 路径含义

这四行其实说的是两类东西。dataset_dirs 指向 LeRobot 数据根目录，也就是包含 data/、meta/、videos/ 的那一层；pretrained_norm_stats 指向 AHA-WAM 自己的归一化统计文件，也就是训练动作和状态时要用的 dataset_stats.json。dataset_stats.json 不是 LeRobot 的 meta/stats.json，它是 AHA-WAM processor 针对动作和状态重新计算出来的统计文件。

6.3 首次统计生成逻辑

如果你还没有 dataset_stats.json ，那么在第一次运行时，将 pretrained_norm_stats 设置为 null 。训练过程会将数据集统计信息写入运行输出目录；你可以重复使用该文件来进行后续的训练。

进一步看代码，第一次训练时如果 pretrained_norm_stats 为空，训练集会遍历数据并写出 dataset_stats.json；验证集不会自己计算统计，而是优先复用训练侧产生的那份文件。这也是为什么第一次运行可以把两个位置都设为 null，后续再显式传入上一次输出目录里的 dataset_stats.json。

# src/ahawam/datasets/lerobot/robot_video_dataset.py
if not pretrained_norm_stats:
    if not is_training_set:
        raise ValueError(
            "pretrained_norm_stats must be provided for validation/test sets since we don't want to calculate stats on them."
        )
    dataset_stats = self.lerobot_dataset.get_dataset_stats(processor)
    save_dataset_stats_to_json(
        dataset_stats, os.path.join(work_dir, "dataset_stats.json")
    )
else:
    dataset_stats = load_dataset_stats_from_json(pretrained_norm_stats)

6.4 首次训练覆盖

解读：第一次训练可以这样覆盖路径：

bash scripts/train_zero1.sh 8 
  task=robotwin_ahawam 
  model=ahawam 
  data.train.dataset_dirs='[data/robotwin2.0/robotwin2.0]' 
  data.val.dataset_dirs='[data/robotwin2.0/robotwin2.0]' 
  data.train.pretrained_norm_stats=null 
  data.val.pretrained_norm_stats=null 
  data.train.text_embedding_cache_dir=data/robotwin2.0/robotwin2.0/text_embeds_cache 
  data.val.text_embedding_cache_dir=data/robotwin2.0/robotwin2.0/text_embeds_cache 
  model.action_dit_pretrained_path=checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt 
  output_dir=runs/robotwin_ahawam_first

6.5 复用统计文件

解读：第二次以后就复用统计：

bash scripts/train_zero1.sh 8 
  task=robotwin_ahawam 
  model=ahawam 
  data.train.dataset_dirs='[data/robotwin2.0/robotwin2.0]' 
  data.val.dataset_dirs='[data/robotwin2.0/robotwin2.0]' 
  data.train.pretrained_norm_stats=runs/robotwin_ahawam_first/dataset_stats.json 
  data.val.pretrained_norm_stats=runs/robotwin_ahawam_first/dataset_stats.json 
  data.train.text_embedding_cache_dir=data/robotwin2.0/robotwin2.0/text_embeds_cache 
  data.val.text_embedding_cache_dir=data/robotwin2.0/robotwin2.0/text_embeds_cache 
  model.action_dit_pretrained_path=checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt

难点提示（为什么验证集不自己算统计）：归一化统计相当于训练集的“尺子”。验证集如果重新拿自己的数据算尺子，就会把评估条件偷偷改掉；正确做法是训练、验证、评估、部署都用同一把尺子，也就是同一个 dataset_stats.json。

这里  就是写到 dataset_stats.json 的归一化统计，输入来自训练集的动作和状态。验证、评估和部署应该复用同一份统计，而不是各自重新计算。这样做可以保证模型在不同阶段看到的是同一个数值尺度，避免评估阶段因为统计口径变化而产生额外偏差。

当 norm_default_mode: z-score 时，可以用这个关系理解动作归一化。实际代码还会处理不同字段、padding 和异常模式，但核心思想就是用训练集均值和方差把动作、状态拉到模型更稳定的数值范围。这里的  可以理解成数值稳定项，避免某些维度方差过小时出现过大的缩放。

7. Option B：从原始 RoboTwin 转成 LeRobot

7.1 原始数据长什么样

如果你不是下载预处理好的 yuanty/robotwin2.0-fastwam，而是从 RoboTwin 上游采集结果开始，一般会看到这样的原始结构。它通常服务于 RoboTwin 自己的采集和评估脚本，里面的 HDF5、mp4、scene_info.json 各司其职，但还没有被整理成 LeRobot 统一读取所需的 metadata 和 parquet 协议，因此需要额外转换后才能训练。

<raw_robotwin_task>/
├── data/
│   ├── episode0.hdf5
│   ├── episode1.hdf5
│   └── ...
├── video/
│   ├── episode0.mp4
│   ├── episode1.mp4
│   └── ...
├── scene_info.json
└── seed.txt

7.2 目标 LeRobot 结构

解读：这个结构能被 RoboTwin 自己的脚本读，但不能直接被 AHA-WAM 训练入口读。AHA-WAM 需要的是 LeRobot v2.1 风格目录：

<converted_robotwin_lerobot>/
├── data/
│   └── chunk-000/
│       ├── episode_000000.parquet
│       └── episode_000001.parquet
├── meta/
│   ├── info.json
│   ├── tasks.jsonl
│   ├── episodes.jsonl
│   └── episodes_stats.jsonl
└── videos/
    └── chunk-000/
        ├── observation.images.cam_high/
        ├── observation.images.cam_left_wrist/
        └── observation.images.cam_right_wrist/

7.3 转换不是改后缀

这里的关键是，转换不是“把 HDF5 改个后缀”，而是把一条 episode 拆成两部分：视频帧进入 mp4，动作、状态、时间戳、episode 索引、任务索引进入 parquet，整个数据集结构和字段定义进入 meta JSON/JSONL。LeRobot 加载器先读 meta/info.json 知道字段和路径模板，再读 parquet 拿非视频数据，最后按 timestamp 到 videos/ 里解码对应帧。

8. 源格式和目标格式逐项对比

8.1 字段语义翻译

原始 RoboTwin 的 HDF5 里通常会存 observation/head_camera/rgb、observation/left_camera/rgb、observation/right_camera/rgb、joint_action/vector 这类字段。AHA-WAM 这边不直接认这些名字，而是通过 configs/data/robotwin.yaml 的 shape_meta 去找 LeRobot 字段。换句话说，转换脚本的职责就是把 RoboTwin 的字段语义翻译成 AHA-WAM 期待的字段语义。只要字段名或维度偏一点，后面的 processor 就会在 shape assert 处报错。

8.2 对照表

8.3 default 字段的真实映射

这里要特别说清一个容易误导的点：有些说明会把 configs/data/robotwin.yaml 里的 key: default 写成 observation.state.default 和 action.default。但在当前仓库代码里，default 是特殊 key，会映射成不带 .default 后缀的字段，也就是 observation.state 和 action。这不是猜测，可以直接从加载器里看到。当前本机预处理数据的 meta/info.json 也使用这两个字段名。

# src/ahawam/datasets/lerobot/base_lerobot_dataset.py
for meta in self.state_meta:
    key = meta["key"]
    meta["lerobot_key"] = f"observation.state.{key}" if key != "default" else "observation.state"

for meta in self.action_meta:
    key = meta["key"]
    meta["lerobot_key"] = f"action.{key}" if key != "default" else "action"

8.4 当前仓库应使用的字段

解读：

这里  可以对应 observation.state，动作 action 也按同样的 14 维协议组织。状态和动作维度相同不代表它们语义完全相同，但它们必须使用同一套左右臂顺序，否则归一化统计和执行器输出都会错位。最直观的后果是模型预测出来的某一维动作，可能会被执行器解释成另一只手臂或夹爪的控制量。

这条式子只是把配置约束写成显式关系：configs/data/robotwin.yaml 里的 raw_shape: 14、shape: 14 和 processor 里的 action_output_dim: 14、proprio_output_dim: 14 必须彼此一致。转换脚本不要为了迁就某个原始字段临时改成 12 维或 16 维，除非你同步修改模型、processor 和评估端的动作协议，否则训练和部署会脱节。

所以转换当前仓库可用的数据时，目标字段应该是下面这些名字。这里建议把它当成转换脚本的单元测试清单：写完 parquet 和 metadata 后，先逐项检查 features 里是否存在这些 key，再进入耗时的视频读取和训练流程。这样可以把字段错误拦在最早阶段，减少无效排查。

observation.images.cam_high
observation.images.cam_left_wrist
observation.images.cam_right_wrist
observation.state
action
task

一句话理解：key: default 不是让你在数据里新建 .default 字段，而是告诉加载器“这个模态只有一个默认分支，直接读 observation.state 或 action”。如果转换脚本照字面写成 observation.state.default，当前加载器反而可能找不到。

9. meta/info.json 应该声明什么

9.1 元数据入口

meta/info.json 是整个 LeRobot 数据集的入口文件。它描述了数据版本、fps、总 episode 数、总帧数、parquet 路径模板、视频路径模板和每个字段的 dtype/shape。AHA-WAM 这边最关心的是三路视频字段、14 维状态、14 维动作、时间戳、frame index、episode index 和 task index。当前本机预处理 RoboTwin2.0 的 fps 是 50，自己转换时必须按真实采集频率填写。

9.2 示例结构

{
  "codebase_version": "v2.1",
  "robot_type": "robotwin",
  "total_episodes": 27500,
  "total_frames": 6075103,
  "total_tasks": 921032,
  "total_videos": 82500,
  "total_chunks": 28,
  "chunks_size": 1000,
  "fps": 50,
  "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet",
  "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4",
  "features": {
    "observation.images.cam_high": {"dtype": "video", "shape": [480, 640, 3]},
    "observation.images.cam_left_wrist": {"dtype": "video", "shape": [480, 640, 3]},
    "observation.images.cam_right_wrist": {"dtype": "video", "shape": [480, 640, 3]},
    "observation.state": {"dtype": "float32", "shape": [14]},
    "action": {"dtype": "float32", "shape": [14]},
    "timestamp": {"dtype": "float32", "shape": [1]},
    "frame_index": {"dtype": "int64", "shape": [1]},
    "episode_index": {"dtype": "int64", "shape": [1]},
    "index": {"dtype": "int64", "shape": [1]},
    "task_index": {"dtype": "int64", "shape": [1]}
  }
}

9.3 路径模板

实际转换脚本里不建议手写所有统计字段。更稳的做法是复用仓库里 LeRobot 相关工具，或者至少严格复用它们的路径模板：data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet 和 videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4。路径模板错了，meta/info.json 就算字段写对，加载器也会去错误位置找文件。

LeRobot 默认按 chunks_size=1000 把 episode 分块，所以 episode index 决定了 parquet 和视频落在哪个 chunk 下。路径模板和 episode 编号必须同时正确，只改文件名不改 metadata，或者只改 metadata 不移动文件，都会让加载器去错误位置查找。批量转换时尤其要避免 episode 编号从 1 开始，而 metadata 又按从 0 开始写，这种错位很隐蔽。

其中  是 observation.images.cam_high、observation.images.cam_left_wrist 或 observation.images.cam_right_wrist。这也是为什么三路相机目录名必须和 info.json 的 video key 完全一致。目录名里多一个下划线、少一个前缀，都会让视频 key 和真实文件路径脱节，最后在解码阶段才暴露出来，排查时还会误以为是视频文件损坏。

进一步看，转换质量比“能不能跑起来”更重要。AHA-WAM 这种模型同时吃视频、动作、状态和文本，如果 fps、动作语义或任务文本不稳定，训练仍然能启动，但学到的对齐关系会偏。下面这张表可以作为转换后的自查清单。每一项都不是格式洁癖，而是会直接影响时间对齐、动作尺度或语言条件的一致性。