209 lines
7.1 KiB
Markdown
209 lines
7.1 KiB
Markdown
# RoboIMI
|
||
|
||
基于 MuJoCo 的机器人仿真与模仿学习框架，实现了使用扩散策略的视觉-语言-动作（VLA）模型，用于机器人操作任务。
|
||
|
||
## 主要特性
|
||
|
||
- **多机器人平台支持**：支持 Diana 和 vx300s 机械臂，可扩展至其他机器人
|
||
- **扩散策略**：采用最先进的扩散模型（DDPM/DDIM）进行动作序列预测
|
||
- **视觉-语言-动作模型**：使用 ResNet-18 视觉骨干网络和空间 softmax 进行视觉特征提取
|
||
- **灵活的控制模式**：支持关节空间和末端执行器（笛卡尔）控制
|
||
- **Hydra 配置系统**：模块化配置系统，便于实验
|
||
- **HDF5 数据集格式**：高效存储和加载演示数据
|
||
- **单臂和双臂任务**：支持单臂和双臂操作任务
|
||
|
||
## 安装
|
||
|
||
### 环境要求
|
||
|
||
- Python 3.8+
|
||
- 支持 CUDA 的 GPU（训练时推荐）
|
||
- Conda 或 Miniconda
|
||
|
||
### 安装步骤
|
||
|
||
```bash
|
||
# 克隆仓库
|
||
git clone <repository-url>
|
||
cd robo-imi-act
|
||
|
||
# 创建并激活 conda 环境
|
||
conda env create -f environment.yml
|
||
conda activate roboimi
|
||
|
||
# 以开发模式安装包
|
||
pip install -e .
|
||
```
|
||
|
||
## 快速开始
|
||
|
||
### 1. 数据采集
|
||
|
||
在仿真环境中记录演示轨迹：
|
||
|
||
```bash
|
||
# 为 vx300s 机器人记录轨迹
|
||
python roboimi/demos/record_sim_episodes.py
|
||
|
||
# 为 Diana 机器人记录轨迹
|
||
python roboimi/demos/diana_record_sim_episodes.py
|
||
```
|
||
|
||
轨迹数据以 HDF5 文件格式保存，包含机器人状态、动作和相机观测。
|
||
|
||
### 2. 计算数据集统计信息
|
||
|
||
训练前需要计算归一化统计数据：
|
||
|
||
```bash
|
||
python roboimi/vla/scripts/calculate_stats.py
|
||
```
|
||
|
||
该命令会生成 `data_stats.pkl` 文件，包含动作和观测的均值/标准差或最小值/最大值。
|
||
|
||
### 3. 训练 VLA 模型
|
||
|
||
使用采集的数据训练视觉-语言-动作模型：
|
||
|
||
```bash
|
||
# 使用默认配置训练
|
||
python roboimi/demos/vla_scripts/train_vla.py
|
||
|
||
# 覆盖特定参数
|
||
python roboimi/demos/vla_scripts/train_vla.py train.batch_size=32 train.lr=5e-5 train.max_steps=50000
|
||
|
||
# 使用不同的模型架构
|
||
python roboimi/demos/vla_scripts/train_vla.py agent=resnet_diffusion data=resnet_dataset
|
||
```
|
||
|
||
训练输出保存至 `outputs/<日期>/<时间>/`，模型检查点保存至 `checkpoints/`。
|
||
|
||
### 4. 评估模型
|
||
|
||
在仿真环境中评估训练好的模型：
|
||
|
||
```bash
|
||
# 使用默认配置评估（使用最佳检查点）
|
||
python roboimi/demos/vla_scripts/eval_vla.py
|
||
|
||
# 指定检查点和评估轮数
|
||
python roboimi/demos/vla_scripts/eval_vla.py eval.ckpt_path=checkpoints/vla_model_step_8000.pt eval.num_episodes=5
|
||
|
||
# 启用动作平滑以获得更流畅的执行
|
||
python roboimi/demos/vla_scripts/eval_vla.py eval.use_smoothing=true eval.smooth_alpha=0.5
|
||
```
|
||
|
||
## 项目结构
|
||
|
||
```
|
||
robo-imi-act/
|
||
├── roboimi/
|
||
│ ├── assets/ # 机器人模型和资源
|
||
│ │ ├── models/manipulators/ # URDF 和 MuJoCo XML 文件
|
||
│ │ └── robots/ # 机器人抽象类
|
||
│ ├── envs/ # 仿真环境
|
||
│ │ ├── mujoco_base.py # MuJoCo 环境基类
|
||
│ │ ├── single_base.py # 单臂任务基类
|
||
│ │ └── double_base.py # 双臂任务基类
|
||
│ ├── vla/ # 视觉-语言-动作模型
|
||
│ │ ├── agent.py # VLAAgent（训练与推理）
|
||
│ │ ├── models/
|
||
│ │ │ ├── backbones/ # 视觉编码器（ResNet 等）
|
||
│ │ │ └── heads/ # 策略头（扩散 UNet1D）
|
||
│ │ ├── conf/ # Hydra 配置文件
|
||
│ │ └── scripts/ # 训练和工具脚本
|
||
│ └── demos/ # 演示脚本和示例
|
||
├── checkpoints/ # 保存的模型检查点
|
||
├── outputs/ # 训练输出（Hydra）
|
||
├── environment.yml # Conda 环境配置
|
||
└── CLAUDE.md # Claude Code 开发指南
|
||
```
|
||
|
||
## 架构设计
|
||
|
||
### VLA 训练流程
|
||
|
||
```
|
||
HDF5 轨迹数据 → Dataset → DataLoader → VLAAgent → 模型检查点
|
||
```
|
||
|
||
**模型组件**：
|
||
- **视觉骨干网络**：ResNet-18 + 空间 softmax，用于从相机图像中提取视觉特征
|
||
- **扩散头**：条件 UNet1D，使用 DDPM/DDIM 预测动作序列
|
||
- **VLAAgent**：组合视觉编码器和扩散策略，处理训练和推理
|
||
|
||
### 配置系统
|
||
|
||
基于 Hydra 的配置文件位于 `roboimi/vla/conf/`：
|
||
- `config.yaml`：主要训练配置（批次大小、学习率、设备）
|
||
- `agent/resnet_diffusion.yaml`：模型架构（动作维度、观测维度、时间窗口）
|
||
- `data/resnet_dataset.yaml`：数据集路径、相机名称、归一化类型
|
||
- `eval/eval.yaml`：评估设置（检查点路径、轮数、平滑参数）
|
||
|
||
使用配置插值保持一致性：`${agent.obs_horizon}`
|
||
|
||
### 数据集格式
|
||
|
||
HDF5 轨迹文件（`episode_*.hdf5`）包含：
|
||
- `action`：机器人动作 `[T, action_dim]`
|
||
- `observations/qpos`：关节位置 `[T, obs_dim]`
|
||
- `observations/images/<cam_name>`：相机图像 `[T, H, W, C]`
|
||
|
||
统计文件（`data_stats.pkl`）存储归一化参数（最小值/最大值/均值/标准差）。
|
||
|
||
## 开发指南
|
||
|
||
### 添加新机器人
|
||
|
||
1. 在 `roboimi/assets/models/manipulators/<robot_name>/` 创建 URDF/XML 文件
|
||
2. 在 `roboimi/assets/robots/<robot_name>.py` 定义机器人类（继承自 `arm_base.py`）
|
||
3. 在 `roboimi/envs/<robot_name>_*.py` 创建环境类
|
||
4. 如需要，在常量中注册机器人
|
||
|
||
### 修改 VLA 架构
|
||
|
||
1. **自定义骨干网络**：在 `roboimi/vla/models/backbones/` 创建新类，继承 `VLABackbone`
|
||
2. **自定义头部**：在 `roboimi/vla/models/heads/` 创建新类，继承 `VLAHead`
|
||
3. **更新配置**：在 `roboimi/vla/conf/agent/` 添加新的 YAML 文件
|
||
4. **接口定义**：参考 `roboimi/vla/core/interfaces.py` 的抽象基类
|
||
|
||
### 训练最佳实践
|
||
|
||
- 采集新数据后务必运行 `calculate_stats.py`
|
||
- 训练时会归一化输入/输出；推理时使用检查点中保存的统计信息进行反归一化
|
||
- 模型预测 `pred_horizon` 步，但只执行前 `action_horizon` 步
|
||
- 推理使用 DDIM（10 步）快速采样；训练使用 DDPM（100 步）
|
||
- 监控验证损失以防止过拟合
|
||
|
||
## 技术细节
|
||
|
||
- **坐标系**：关节空间（qpos）或末端执行器空间（xyz + rpy + 夹爪）
|
||
- **动作时间窗口**：`obs_horizon` 为观测窗口，`pred_horizon` 为预测窗口，`action_horizon` 为执行窗口
|
||
- **归一化**：对稳定训练至关重要 - 训练前务必计算统计信息
|
||
- **推理加速**：使用 DDIM 调度器，比训练时的 DDPM 快 10 倍
|
||
- **设备配置**：通过 `train.device` 配置（cuda/cpu）
|
||
|
||
## 许可证
|
||
|
||
[在此添加许可证信息]
|
||
|
||
## 引用
|
||
|
||
如果您在研究中使用了本代码库，请引用：
|
||
|
||
```bibtex
|
||
[在此添加引用信息]
|
||
```
|
||
|
||
## 贡献
|
||
|
||
欢迎贡献！请随时提交 Pull Request 或开启 Issue。
|
||
|
||
## 致谢
|
||
|
||
本项目基于以下开源项目构建：
|
||
- [MuJoCo](https://mujoco.org/) - 物理仿真引擎
|
||
- [PyTorch](https://pytorch.org/) - 深度学习框架
|
||
- [Hydra](https://hydra.cc/) - 配置管理系统
|
||
- [Diffusers](https://github.com/huggingface/diffusers) - 扩散模型库
|