LeRobot MuJoCo 训练ACT、SmolVLA、Pi0教程

本仓库提供了一个最小可运行示例：用于采集示教数据，并在自定义数据集上训练（或微调）视觉-语言-动作（VLA）模型。

安装

我们在 Python 3.10 上测试通过。

不建议直接使用 pip install lerobot，可能会报错。

安装 MuJoCo 相关依赖和 lerobot：

bashconda create -n py310 python=3.10
pip install -r requirements.txt
conda install jupyterlab
pip install ipywidgets ipykernel
python -m ipykernel install --user --name py310 --display-name "py310"
jupyter lab .
# 在当前目录启动


# 如果torch的cuda有问题：
pip uninstall -y torch torchvision torchaudio
pip install --no-cache-dir --force-reinstall --index-url https://download.pytorch.org/whl/cu124 torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0
python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())"

请确认 MuJoCo 版本为 3.1.6。

解压资源文件：

bashcd asset/objaverse
unzip plate_11.zip

更新计划

[x] Viewer 更新
[x] 增加多种 mug、plate，对应不同语言指令
[x] 增加 pi_0 训练与推理
[x] 增加 SmolVLA

1. 采集示教数据

运行 1.collect_data.ipynb

默认使用 ROBOTIS OMY 六轴机械臂及夹爪的 MuJoCo 仿真模型，场景文件为 asset/example_scene_y.xml，无需实体机械臂。运行采集单元后，让 MuJoCo 窗口获得键盘焦点再操作。本章固定版本的 LeRobotDataset 按 LeRobot v2.1 格式保存数据。

在给定环境中采集示教数据。任务是抓起杯子并放到盘子上。当杯子在盘子上、夹爪打开且末端执行器位于杯子上方时，环境判定成功。

键位说明：

WASD：x-y 平面移动
R/F：z 轴移动
Q/E：倾斜
方向键：其余旋转
空格：切换夹爪状态
Z：重置环境，并丢弃当前回合缓存数据

叠加图像说明：

右上：Agent 视角
右下：腕部（第一人称）视角
左上：侧视图
左下：俯视图

数据集结构：

pythonfps = 20,
features={
    "observation.image": {
        "dtype": "image",
        "shape": (256, 256, 3),
        "names": ["height", "width", "channels"],
    },
    "observation.wrist_image": {
        "dtype": "image",
        "shape": (256, 256, 3),
        "names": ["height", "width", "channel"],
    },
    "observation.state": {
        "dtype": "float32",
        "shape": (6,),
        "names": ["state"], # x, y, z, roll, pitch, yaw
    },
    "action": {
        "dtype": "float32",
        "shape": (7,),
        "names": ["action"], # 6 个关节角 + 1 个夹爪
    },
    "obj_init": {
        "dtype": "float32",
        "shape": (6,),
        "names": ["obj_init"], # 仅物体初始位置，训练中不使用
    },
},

数据默认保存在 ./demo_data 目录。仓库中的历史示例 demo_data_example 可用于查看目录结构和回放；训练请优先使用本次修复后重新采集的数据。

开始操作后，代码以 20 Hz 缓存图像、状态和动作；成功判定成立后保存一个 episode。Z 会丢弃当前未保存回合，成功保存或重置后需重新开始操作才会记录新回合。默认 NUM_DEMO=1 只采集一条示教；用于训练时应增加数量并覆盖不同物体位置。

动作与观测的区别：四个键盘采集入口的 observation.state 是下发本帧命令前的末端位姿，action 是逆运动学求出的绝对目标关节角加归一化夹爪命令（0 为开、1 为闭）。OMY、Nova5、xArm6 为 7 维，xArm7 为 8 维。调用 PnPEnv.step(action) 准备控制命令后，用 PnPEnv.get_commanded_joint_action() 读取动作标签；step() 的返回值仍是观测，不能作为目标动作。夹爪有平滑或接触限位时，标签记录经过这些限制后实际下发的命令。

旧数据处理（2026-10-02）：旧版本的 1.collect_data.ipynb、1.collect_data_nova5.ipynb、1.collect_data_xarm6.ipynb、1.collect_data_xarm7.ipynb 误将当前关节状态保存为动作标签。建议使用本次修复后的代码重新采集并训练；仅更新代码不能修复已有数据中的标签。请保留旧数据备份，改用新的 ROOT 目录，并同步修改训练数据路径，避免混入旧数据。更新环境代码后需重启 Notebook kernel，再从头执行。

第 5 节的语言条件采集入口也使用相同的目标动作接口，但其 observation.state 保留为当前的六个机械臂关节角，而非末端位姿；训练和部署时应使用各自对应的状态定义。

安装本章依赖后，可在本章目录运行 python -m unittest -v test_collection_actions 检查动作标签、夹爪命令、回合边界及五个采集入口。该检查无需打开图形窗口或运行 GPU 训练。

2. 回放数据

按采集时使用的机器人选择回放入口，数据路径和机器人模型必须匹配：

在重建后的仿真场景中可视化你的动作。主窗口会回放动作；右上和右下叠加图像来自数据集。

3. 训练 Action-Chunking-Transformer（ACT）

运行 3.train.ipynb

大约需要 30~60 分钟。

在自定义数据集上训练 ACT。示例中 chunk_size=10。

训练好的 checkpoint 会保存在 ./ckpt/act_y。

可通过与数据集真值动作对比，评估策略误差。

PicklingError: Can't pickle <function <lambda>...>

如遇 pickling 错误，请将 `num_workers` 设为 `0`，例如：

pythondataloader = torch.utils.data.DataLoader(
    dataset,
    num_workers=0, # 4
    batch_size=64,
    shuffle=True,
    pin_memory=device.type != "cpu",
    drop_last=True,
)

4. 部署 ACT 策略

运行 4.deploy.ipynb

如果没有可用于训练的 GPU，可从 Google Drive 下载 checkpoint：

https://drive.google.com/drive/folders/1UqxqUgGPKU04DkpQqSWNgfYMhlvaiZsp?usp=sharing

5-6. 语言条件环境中的采集与可视化

5.language_env.ipynb：键盘遥操作采集数据（键位与第一个环境一致）
1中只采集一条数据，5中采集20条数据，两个任务，红杯子和蓝杯子
6.visualize_data.ipynb：可视化已采集数据

数据示例

模型与数据集

获取语言条件数据集

Pi0 和 SmolVLA 使用同一份语言条件数据，选择以下一种方式即可。所有命令均在本章目录运行，数据统一放在 ./demo_data_language。

方式一：下载公开样例

Datawhale 公开样例是 OMY 场景的 LeRobot v2.1 数据集，包含 20 个 episode、2621 帧，以及红杯子和蓝杯子两个任务。

先安装 Git LFS，然后执行：

bashgit lfs install
git clone https://huggingface.co/datasets/Datawhale/datawhale_eai_pnp_language demo_data_language
git -C demo_data_language lfs pull

已有同名目录时不要重复 clone 或删除原数据。如果目录是该数据集的 Git 仓库，可直接执行最后一条命令补齐 LFS 文件。仅下载到 LFS 指针文件不代表数据已完整。

方式二：自行采集

运行 5.language_env.ipynb，保留 REPO_NAME = 'datawhale_eai_pnp_language' 和 ROOT = "./demo_data_language"，采集并保存红杯子、蓝杯子示教，然后用 6.visualize_data.ipynb 检查轨迹。该 Notebook 在目录已存在时包含删除旧数据的确认步骤；已有下载数据或示教数据时，应先备份或更换目录。

两种方式均可使用本章的 pi0_datawhale_eai.yaml 和 smolvla_datawhale_eai.yaml，无需改配置：

yamldataset:
  repo_id: datawhale_eai_pnp_language
  root: ./demo_data_language

root 是相对于运行目录的本地数据路径；完整数据已在该目录时，固定版本 LeRobot 使用本地文件。repo_id 是数据集标识，上述不带组织名的值沿用采集 Notebook 的命名，不能当作可搜索的 Hugging Face 完整地址。需要从 Hub 自动获取时应使用完整 ID Datawhale/datawhale_eai_pnp_language。更改本地目录时，需同步修改 YAML 和部署 Notebook 中的路径。

模型参考

Model 🤗	Dataset 🤗
pi_0 finetuned	dataset
smolvla finetuned	同上

7. 训练与部署 pi_0

以下命令均在本章目录执行，使用 requirements.txt 固定的 LeRobot 10b7b3532543b4adfb65760f02a49b4c537afde7。先执行 python check_training_configs.py 检查两个配置（不训练、不下载权重）。 YAML 命令中 --config_path 与文件名之间保留空格；该版本的等号写法会走 JSON checkpoint 加载分支。

train_model.py：训练脚本
pi0_datawhale_eai.yaml：训练配置
7.pi0.ipynb：部署示例

训练命令：

bashpython train_model.py --config_path pi0_datawhale_eai.yaml

部署效果：

训练日志：

配置以本节链接的 YAML 为准。默认读取第 5 节的 ./demo_data_language，dataset.repo_id 为 datawhale_eai_pnp_language；使用其他数据时需同时修改两项。默认关闭 W&B，启用前需设置自己的项目和账号。训练需要完整数据、模型权重和 CUDA 环境；配置解析通过不代表完整训练已验证。

8. 训练与部署 SmolVLA

train_model.py：训练脚本
smolvla_datawhale_eai.yaml：训练配置
8.smolvla.ipynb：部署示例

训练命令：

bashpython train_model.py --config_path smolvla_datawhale_eai.yaml

部署效果：

训练日志：

10. 用 π0.5 微调 OMY

运行 10.pi05_omy.ipynb。流程接续第 1 节的 OMY 单杯抓放数据：检查示教及动作标签、转换 LeRobot v2.1 为独立 v3 数据、计算训练回合分位数、五步训练检查、正式微调、留出动作对比和 MuJoCo 仿真评测。

使用独立的 Python 3.12 / LeRobot 0.6.1 环境及 requirements_pi05.txt，不要升级第 1～9 节的旧环境。具体安装、可编辑参数和启动步骤在笔记本中；辅助接口见 pi05_omy_utils.py。默认从 lerobot/pi05_base 微调，也可用 OMY_PI05_PRETRAINED 指定兼容的本地预训练目录。批量大小为 1、只训练动作专家，先做五步检查，完整训练和评测结果由实际运行生成。

pi05_omy_train.py 调用官方训练入口并要求全部预训练参数加载成功，权重读取失败时停止训练；推理加载也做相同检查，防止固定版本加载器的随机初始化回退被误报为微调成功。远程短步验证的设置、来源与结果见 pi05_omy_short_check.md。

输入是两路图像、六维末端位姿和语言指令；输出是六个绝对目标关节角及归一化夹爪命令，保持 20 Hz 采集频率，关闭相对动作转换。第 5 节的状态为关节角且场景接口不同，不能直接使用该流程。历史错误动作标签不能通过格式迁移恢复，请优先使用修复后的第 1 节重新采集并回放。

转换保留原数据，目标目录必须为空；三个以上回合按整回合留出约 20% 评测，只用训练回合计算归一化统计量。一两条示教适合接口检查，不能作为泛化成功率依据。训练需要合适的 CUDA 工作站；短步训练通过、损失下降均不代表抓放已经复现成功。

安装独立环境后，在本章目录运行 python -m unittest -v test_pi05_omy，检查新版数据接口、图像与动作保真、回合边界及控制时序。测试不下载模型权重、不启动完整训练。

后续，我们还会为大家补充展示多条数据训练ACT和diffusion训练抓取实验。敬请期待。

致谢

Robotis-OMY 机械臂资源来自 robotis_mujoco_menagerie
MuJoco Parser Class 改自 yet-another-mujoco-tutorial-v3
教程参考了 lerobot examples
plate 与 mug 资源来自 Objaverse