🤖 LeRobot Tutorial with MuJoCo

🦾 这是一个基于 LeRobot 和 MuJoCo 的抓取放置任务复现仓库，包含数据采集、数据回放、ACT 训练与部署，以及 pi_0 和 SmolVLA 的训练/推理示例。

“基于开源项目 https://github.com/jeongeun980906/lerobot-mujoco-tutorial 的复现”

📦 项目内容

• 机械臂 teleoperation 采集示范数据
• 回放与可视化采集到的数据
• ACT 训练与仿真部署
• 语言条件环境的数据采集与可视化
• pi_0 和 SmolVLA 的训练与部署配置

🗂 仓库结构

• 1.collect_data.ipynb: 数据采集
• 2.visualize_data.ipynb: 数据回放与可视化
• 3.train.ipynb: ACT 训练
• 4.deploy.ipynb: ACT 部署
• 5.language_env.ipynb: 语言条件环境采集
• 6.visualize_data.ipynb: 语言条件数据可视化
• 7.pi0.ipynb: pi_0 部署
• 8.smolvla.ipynb: SmolVLA 部署
• train_model.py: 统一训练入口
• pi0_omy.yaml: pi_0 训练配置
• smolvla_omy.yaml: SmolVLA 训练配置

🛠 安装

测试环境：Python 3.10，MuJoCo 3.1.6。

建议直接安装依赖，不要用 pip install lerobot 单独安装官方包，否则容易和当前仓库中的代码版本不一致。

pip install -r requirements.txt

解压示例资产：

cd asset/objaverse
unzip plate_11.zip

🚀 快速开始

1. 采集演示数据

运行 1.collect_data.ipynb。任务是把 mug 放到 plate 上。

控制方式：

• WASD: XY 平面移动
• RF: Z 轴移动
• QE: 俯仰/倾斜
• 方向键: 其余旋转
• Space: 切换夹爪开合
• Z: 重置环境并丢弃当前 episode

数据中包含多视角图像与状态、动作信息，示例格式如下：

fps = 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", "channels"],
  },
  "observation.state": {
    "dtype": "float32",
    "shape": (6,),
    "names": ["state"],
  },
  "action": {
    "dtype": "float32",
    "shape": (7,),
    "names": ["action"],
  },
  "obj_init": {
    "dtype": "float32",
    "shape": (6,),
    "names": ["obj_init"],
  },
}

采集后的数据会保存到 ./demo_data，目录结构类似：

.
├── data
│   └── chunk-000
│       ├── episode_000000.parquet
│       └── ...
└── meta
  ├── episodes.jsonl
  ├── info.json
  ├── stats.json
  └── tasks.jsonl

仓库中还提供了 示例数据 便于直接跑通流程。

2. 回放数据

运行 2.visualize_data.ipynb。

这个步骤会在重建的仿真场景里回放动作，并显示数据集中对应的叠加图像。

3. 训练 ACT

运行 3.train.ipynb。通常需要约 30 到 60 分钟，具体取决于数据量和显卡性能。

训练完成后，checkpoint 会保存在 ./ckpt/act_y。

如果你遇到如下报错：

PicklingError: Can't pickle <function <lambda> at ...>: attribute lookup <lambda> on __main__ failed

把 DataLoader 的 num_workers 改成 0 即可。

4. 部署 ACT

运行 4.deploy.ipynb。如果你不打算本地训练，也可以直接下载已有 checkpoint 进行部署。

5. 语言条件环境

• 5.language_env.ipynb: 采集语言条件数据
• 6.visualize_data.ipynb: 可视化采集数据

🤗 pi_0 与 SmolVLA

pi_0

• 训练入口：train_model.py
• 配置文件：pi0_omy.yaml
• 部署 notebook：7.pi0.ipynb

训练命令：

python train_model.py --config_path pi0_omy.yaml

训练日志与 rollout 示例：

SmolVLA

• 训练入口：train_model.py
• 配置文件：smolvla_omy.yaml
• 部署 notebook：8.smolvla.ipynb

训练命令：

python train_model.py --config_path smolvla_omy.yaml

训练日志与 rollout 示例：

🔌 离线使用说明

如果你在没有网络的环境里运行 pi_0 或 SmolVLA，请确保相关模型文件已经放在本地，并在对应 notebook 里使用本地路径加载。这样可以避免默认从 Hugging Face Hub 拉取模型或 tokenizer。

🔎 参考数据与模型

• pi_0 finetuned: Jeongeun/omy_pnp_pi0
• SmolVLA finetuned: Jeongeun/omy_pnp_smolvla
• Dataset: Jeongeun/omy_pnp_language

🙏 致谢

• 机器人本体资产来自 robotis_mujoco_menagerie
• MuJoCo parser 改自 yet-another-mujoco-tutorial-v3
• 训练与数据处理流程参考了 LeRobot examples
• plate 与 mug 资产来自 Objaverse

📝 说明

本仓库为个人复现与学习项目，不代表原始项目官方仓库。