diff --git a/README.md b/README.md new file mode 100644 index 0000000..b5d0744 --- /dev/null +++ b/README.md @@ -0,0 +1,208 @@ +# RoboIMI + +基于 MuJoCo 的机器人仿真与模仿学习框架,实现了使用扩散策略的视觉-语言-动作(VLA)模型,用于机器人操作任务。 + +## 主要特性 + +- **多机器人平台支持**:支持 Diana 和 vx300s 机械臂,可扩展至其他机器人 +- **扩散策略**:采用最先进的扩散模型(DDPM/DDIM)进行动作序列预测 +- **视觉-语言-动作模型**:使用 ResNet-18 视觉骨干网络和空间 softmax 进行视觉特征提取 +- **灵活的控制模式**:支持关节空间和末端执行器(笛卡尔)控制 +- **Hydra 配置系统**:模块化配置系统,便于实验 +- **HDF5 数据集格式**:高效存储和加载演示数据 +- **单臂和双臂任务**:支持单臂和双臂操作任务 + +## 安装 + +### 环境要求 + +- Python 3.8+ +- 支持 CUDA 的 GPU(训练时推荐) +- Conda 或 Miniconda + +### 安装步骤 + +```bash +# 克隆仓库 +git clone +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/`:相机图像 `[T, H, W, C]` + +统计文件(`data_stats.pkl`)存储归一化参数(最小值/最大值/均值/标准差)。 + +## 开发指南 + +### 添加新机器人 + +1. 在 `roboimi/assets/models/manipulators//` 创建 URDF/XML 文件 +2. 在 `roboimi/assets/robots/.py` 定义机器人类(继承自 `arm_base.py`) +3. 在 `roboimi/envs/_*.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) - 扩散模型库