OpenDM05-Memory 训练开源

OpenDM05-Memory 开源：让机器人利用历史视觉记忆完成长时序操作任务

在机器人操作任务中，仅依赖当前时刻的相机画面往往是不够的。

例如，在 cover_blocks 任务中，机器人需要先观察不同颜色积木的位置；当积木被完全遮挡后，再依靠此前的视觉信息，按照指定顺序揭开盖子。此时，当前画面已经无法提供完整答案，策略必须真正利用历史观测。

为了解决这类问题，我开源了 OpenDM05-Memory：

GitHub：https://github.com/garlic-byte/OpenDM05-Memory

该项目基于 Dexmal 官方 OpenDM / DM0.5，补充了一套可复现的视觉 Memory 数据转换、训练和验证流程，支持在 DM0.5 中引入稀疏历史图像，并针对 RoboDojo Memory 类型任务进行训练。

如果这个项目对你有帮助，欢迎在 GitHub 点一个 Star。

一、项目解决了什么问题？

普通 VLA 策略通常以当前状态、当前相机图像和语言指令作为输入：

复制代码

当前图像 + 当前机器人状态 + 语言指令 → 动作序列

但 Memory 任务还需要模型访问过去的重要视觉信息：

复制代码

当前图像
+ 当前机器人状态
+ 语言指令
+ 历史视觉观测
→ 动作序列

OpenDM05-Memory 会从同一个 episode 中，按照指定间隔采样过去的图像：

复制代码

t - memory_frames × memory_stride
...
t - 2 × memory_stride
t - memory_stride

所有历史帧按照从旧到新的顺序排列：

不包含当前帧；
不跨越 episode；
episode 开头历史不足时，可以使用 Masked Slot 左填充；
Memory 图像经过视觉编码与池化后，写入语言模型前缀中的专用 Memory Token；
无效的填充位置不会贡献图像特征，也不会参与注意力计算。

例如：

复制代码

MEMORY_FRAMES=20
MEMORY_STRIDE=25

表示每个训练样本最多使用 20 个历史观测，相邻历史观测之间间隔 25 个数据步。

二、项目特点

OpenDM05-Memory 主要增加了以下能力：

支持 DM0.5 稀疏视觉 Memory 训练；
支持 LeRobot 2.1 数据转换；
支持 LeRobot 3.0 转换为 2.1，再转换为 DM05 Memory JSONL；
支持图片和 MP4 视频形式的相机数据；
支持 PyAV 和 TorchCodec 视频读取；
支持多进程 DataLoader 下的 Decoder 缓存；
支持 episode 开头固定长度 Memory 左填充；
提供完整的 8 卡复现训练参数；
默认不启用 W&B，不在代码中保存任何 API Key；
不提交数据集、模型权重、训练输出和本地环境。

项目地址：

复制代码

https://github.com/garlic-byte/OpenDM05-Memory

上游 OpenDM：

复制代码

https://github.com/dexmal/opendm

三、环境搭建

推荐使用官方 Docker 环境，避免 CUDA、PyTorch 和 FlashAttention 版本不一致。

1. Docker 方式

复制代码

git clone https://github.com/garlic-byte/OpenDM05-Memory.git
cd OpenDM05-Memory

docker run -it --rm --gpus all --network host \
  --name opendm-memory \
  --shm-size=16g \
  -v "$PWD":/app/opendm \
  -w /app/opendm \
  dexmal/opendm:latest /bin/bash

conda activate opendm
pip install -e .

2. 本地 Conda 方式

复制代码

conda create -n opendm python=3.10 -y
conda activate opendm

pip install torch torchvision \
  --index-url https://download.pytorch.org/whl/cu128

pip install ninja packaging
MAX_JOBS=2 pip install flash-attn --no-build-isolation

pip install -e .

如果训练时使用 TorchCodec 视频后端，还需要安装与当前 PyTorch、CUDA 和 FFmpeg 兼容的 TorchCodec。

四、下载 DM0.5 基础模型

基础权重默认放在：

复制代码

./checkpoints/DM05

下载命令：

复制代码

huggingface-cli download Dexmal/DM05 \
  --local-dir ./checkpoints/DM05

也可以把模型放在其他位置，训练时通过环境变量指定：

复制代码

export DM05_MODEL_PATH=/path/to/DM05

模型权重已经加入 .gitignore，不应该提交到 GitHub。

五、支持 LeRobot 2.1 和 3.0

项目使用 LeRobot 2.1 作为统一中间格式：

复制代码

LeRobot 2.1 ───────────────────────→ DM05 Memory JSONL

LeRobot 3.0 → LeRobot 2.1 ────────→ DM05 Memory JSONL

之所以采用两步转换，是因为两种 LeRobot 格式的存储方式不同：

LeRobot 2.1：每个 episode 对应独立的 Parquet 和 MP4；
LeRobot 3.0：多个 episode 被合并到共享 Parquet 和 MP4 分片中，需要通过 metadata 解析 episode offset。

相关转换代码已经放在项目的 script/ 目录中。

六、LeRobot 3.0 转 2.1

首先创建一个独立的转换环境：

复制代码

conda create -n lerobot-convert python=3.10 -y
conda activate lerobot-convert

pip install lerobot==0.4.0 \
  jsonlines \
  pyarrow \
  numpy \
  tqdm \
  huggingface-hub

确认系统已经安装 FFmpeg：

复制代码

ffmpeg -version

执行转换：

复制代码

python script/convert_lerobot_v3_to_v21.py \
  --input-root /path/to/dataset_v30 \
  --output-root /path/to/dataset_v21

该转换器采用非破坏式设计：

不修改原始 v3.0 数据；
不移动原始数据；
不删除原始数据；
输出目录已存在时拒绝覆盖；
将共享 Parquet 和 MP4 重建为 episode 独立的 v2.1 结构。

七、LeRobot 2.1 转 DM05 Memory JSONL

如果数据本身就是 LeRobot 2.1，可以直接执行这一步。

复制代码

conda activate opendm

python script/convert_lerobot_v21_to_memory_jsonl.py \
  --input-root /path/to/dataset_v21 \
  --output-dir /path/to/dataset_v21/dm05_jsonl \
  --state-key observation.state \
  --action-key action \
  --camera-keys observation.images.head,observation.images.left_wrist,observation.images.right_wrist \
  --output-image-keys images_1,images_2,images_3

转换完成后，每个 episode 会对应一个 JSONL 文件：

复制代码

dm05_jsonl/
├── episode_000000.jsonl
├── episode_000001.jsonl
├── episode_000002.jsonl
└── conversion_manifest.json

转换器不会重复复制 MP4，而是在 JSONL 中记录：

复制代码

{
  "type": "video",
  "url": "videos/chunk-000/observation.images.head/episode_000000.mp4",
  "frame_idx": 10
}

conversion_manifest.json 会记录：

JSONL 输出目录；
原始媒体根目录；
episode 数量；
frame 数量；
state/action 字段映射；
相机字段映射。

八、DM05 Memory JSONL 格式

每一行代表一个时间步，主要包含：

复制代码

{
  "state": [0.0, 0.1],
  "action": [0.0, 0.2],
  "prompt": "place the object",
  "images_1": {
    "type": "video",
    "url": "videos/head.mp4",
    "frame_idx": 1
  },
  "images_2": {
    "type": "video",
    "url": "videos/left_wrist.mp4",
    "frame_idx": 1
  },
  "images_3": {
    "type": "video",
    "url": "videos/right_wrist.mp4",
    "frame_idx": 1
  }
}

每一个 JSONL 文件必须只包含一个 episode。这样 Memory 采样时不会跨越不同轨迹。

九、启动 Memory 训练

完成数据转换后，设置数据路径：

复制代码

export JSONL_DIR=/path/to/dataset_v21/dm05_jsonl
export IMAGE_DIR=/path/to/dataset_v21

对于 14 维双臂 ARX X5，状态描述为：

复制代码

export STATE_DESC=joint,joint,joint,joint,joint,joint,gripper,joint,joint,joint,joint,joint,joint,gripper

启动训练：

复制代码

bash script/train_memory_sft.sh

如需先检查路径和最终命令，不启动训练：

复制代码

DRY_RUN=1 bash script/train_memory_sft.sh

十、RoboDojo Memory 复现参数

项目中的 train_memory_sft.sh 默认采用已复现的 RoboDojo ARX X5 Memory 训练配置。

参数	默认值	说明
`ACTION_MODE`	`ABSOLUTE`	绝对关节位置
`CHUNK_SIZE`	`50`	预测未来50步动作
`MEMORY_IMAGE_KEYS`	`images_1`	全局相机作为Memory
`MEMORY_FRAMES`	`20`	20个历史观测
`MEMORY_STRIDE`	`25`	历史观测间隔25步
`MAX_MEMORY_IMAGES`	`20`	Memory最大图像数
`LEFT_PAD_MEMORY`	`True`	episode开头左填充
`MODEL_MAX_LENGTH`	`1536`	多模态Token上限
`VIDEO_BACKEND`	`torchcodec`	视频读取后端
`AUGMENTATION_PROBABILITY`	`0.0`	关闭图像增强
`NPROC_PER_NODE`	`8`	8张GPU
`PER_DEVICE_BATCH_SIZE`	`4`	单卡Batch Size
`GRADIENT_ACCUMULATION_STEPS`	`8`	梯度累积
有效全局Batch	`256`	`8 × 4 × 8`
`NUM_TRAIN_STEPS`	`10000`	总训练步数
`LEARNING_RATE`	`4e-5`	MuonAdamW学习率
`WARMUP_STEPS`	`1000`	Warmup步数
`SAVE_STEPS`	`2000`	Checkpoint保存间隔
`USE_LORA`	`False`	全参数微调
Attention	`SDPA`	LLM、Vision、Action
`SEED`	`42`	随机种子

默认配置会检查：

复制代码

8张GPU
每张显存不少于79000 MiB

如果需要在其他硬件上训练，可以设置：

复制代码

VALIDATE_GPU_MEMORY=False

但应同步调整：

NPROC_PER_NODE
PER_DEVICE_BATCH_SIZE
GRADIENT_ACCUMULATION_STEPS

建议尽量保持有效全局 Batch Size 为 256：

复制代码

有效全局Batch
= GPU数量
× 单卡Batch
× 梯度累积

十一、为什么使用绝对位置？

RoboDojo Memory 的复现配置使用：

复制代码

ACTION_MODE=ABSOLUTE

也就是直接预测绝对关节目标。

这与使用相对关节增量的训练方式不同。训练、归一化统计、Checkpoint 和推理阶段必须使用相同的动作表示，否则模型输出的含义会不一致。

因此，复现该配置时不要把 ACTION_MODE 修改成 RELATIVE。

如果换成其他机器人或数据集，需要同时确认：

state 的物理含义；
action 的物理含义；
state/action 维度；
gripper 是否使用绝对值；
训练数据是否为绝对关节位置；
推理控制器需要绝对目标还是增量目标。

十二、RoboDojo Memory 任务结果

根据 RoboDojo 官方 OpenDM05 Simulation Rollout Leaderboard，DM0.5 在部分 Memory 任务上的结果如下：

|--------------------------------|----------|-----------|---------|
| Memory任务 | 机器人 | Avg Score | 成功率 |
| cover_blocks | arx_x5 | 100.00 | 100.00% |
| press_by_number | arx_x5 | 95.33 | 95.00% |
| match_and_pick_from_conveyor | arx_x5 | 70.67 | 71.00% |

排行榜地址：

复制代码

https://robodojo-benchmark.com/leaderboard/rollouts/OpenDM05?bench=sim

其中：

cover_blocks 要求机器人记住被遮挡积木的颜色位置；
press_by_number 要求机器人保留此前出现的数字或操作信息；
match_and_pick_from_conveyor 要求机器人记住目标，并在传送带场景中完成匹配与抓取。

这些指标来自模拟器 Rollout，不是训练 Loss 或离线验证集 Accuracy。完整复现还需要保持以下配置一致：

RoboDojo 模拟器版本；
机器人与相机配置；
DM0.5 基础权重；
训练后 Checkpoint；
Action Chunk；
Memory 帧数与间隔；
动作模式；
归一化统计；
评测随机种子。

十三、W&B 使用与隐私安全

项目默认不启用 W&B。

如需记录训练曲线，可以在仓库外登录：

复制代码

wandb login

然后设置项目名：

复制代码

WANDB_PROJECT=dm05-memory \
bash script/train_memory_sft.sh

不要把 WANDB_API_KEY 写入：

Python 文件；
Shell 脚本；
README；
Git 配置；
提交记录。

非交互训练任务应通过作业系统环境变量或 Secrets Manager 注入。

十四、项目验证

代码提供以下检查方式：

复制代码

pytest -q tests/test_memory_dataset.py

执行全部提交前检查：

复制代码

pre-commit run --all-files

只检查训练命令：

复制代码

DRY_RUN=1 bash script/train_memory_sft.sh

十五、总结

OpenDM05-Memory 的目标不是重新实现一套 DM0.5，而是在官方 OpenDM 基础上，补齐视觉 Memory 训练最容易缺失的工程环节：

LeRobot 2.1 / 3.0 数据转换；
episode 级 JSONL 构建；
历史视觉采样；
多视频稀疏读取；
Memory Token 注入；
固定槽位 Mask；
可复现训练参数；
RoboDojo Memory 任务结果说明；
权重、数据和隐私信息隔离。

如果你正在研究：

Vision-Language-Action；
机器人长时序任务；
历史视觉建模；
具身智能 Memory；
DM0.5 微调；
RoboDojo；
LeRobot 数据转换；

欢迎试用 OpenDM05-Memory。

GitHub：

复制代码

https://github.com/garlic-byte/OpenDM05-Memory

欢迎提交 Issue、Pull Request，也欢迎 Star 和 Fork。