斯坦福 Mobile Aloha 开源ACT代码复现教程

0. 项目地址：Mobile ALOHA

GitHub - MarkFzp/act-plus-plus: Imitation learning algorithms with Co-training for Mobile ALOHA: ACT, Diffusion Policy, VINNhttps://github.com/MarkFzp/act-plus-plus

1. 环境配置

新建Pycharm工程，从github上下载项目拷贝过去，在termina里进行环境配置：

1.1 安装部分依赖项

conda create -n aloha python=3.8.10
conda activate aloha
pip install pyquaternion
pip install pyyaml
pip install rospkg
pip install pexpect
pip install mujoco==2.3.7
pip install dm_control==1.0.14
pip install opencv-python
pip install matplotlib
pip install einops
pip install packaging
pip install h5py
pip install ipython
pip install diffusers
cd act/detr && pip install -e .

1.2 安装torch

pytorch需要单独安装，保证安装的是gpu版本，根据自己的cuda版本在torch官网选择安装命令：

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

不需要额外安装cudNN

1.3 安装egl-probe

egl-probe直接安装（pip install or pip build）会有一些问题（如下图所示：找不到egl-probe.glad package&&CMake指令报错），需要去github下载源码进行编译安装。

i.首先安装一下CMake进行后续编译

pip install CMake

ii.从egl-probe的github项目地址拉取源代码

iii.修改setup.py如下(注释掉的是原来的)

        # build using cmake
        #subprocess.check_call("cmake ..; make -j", cwd=build_dir, shell=True)
        subprocess.check_call("cmake ..", cwd=build_dir, shell=True)

        # build using cmake
        #subprocess.check_call("cmake ..; make -j", cwd=build_dir, shell=True)
        subprocess.check_call("cmake ..", cwd=build_dir, shell=True)

iv.编译并安装egl-probe

#先进入到下载的egl-probe项目根目录里
cd egl_probe-master
python setup.py build
python setup.py install

至此，egl-probe终于安装完成了（和gpt勾心斗角了好久）

1.4 Robomimic

Robomimic是斯坦福大学开发的一个用于机器人演示学习的框架。它提供了在机器人操作领域收集的广泛的演示数据集，以及从这些数据集学习的学习算法。该项目是通过模拟环境更深层次的推进机器人智能(ARISE)计划的一部分，旨在降低人工智能和机器人交叉领域前沿研究的进入门槛。

直接pip安装Robomimic框架后会报错：

ModuleNotFoundError: No module named 'robomimic.algo.diffusion_policy'
# 原因：robomimic的release版本还不支持diffusion_policy模块。

下载,离线安装diffusion-policy-mg分支即可：GitHub - ARISE-Initiative/robomimic at diffusion-policy-mg

git clone https://github.com/ARISE-Initiative/robomimic.git -b diffusion-policy-mg
cd robomimic && pip install -v -e .

1.5 util

配置环境的最后一步，安装detr所需要的util库，这个库已经在Mobile ALOHA的源码里了，路径为

act-plus-plus-main\detr\util

将util整个文件夹复制到刚才创建的conda虚拟环境下的这个路径：

D:\Anaconda3\envs\aloha\Lib\site-packages

2.SIM仿真环境训练

2.1数据集构建

模拟实验(LEGACY table-top ALOHA environments)数据集收集

模拟实验(LEGACY table-top ALOHA environments)中运行sim_transfer_cube_scripted任务，生成50集脚本数据，运行:

python record_sim_episodes.py --task_name sim_transfer_cube_scripted --dataset_dir <data save dir> --num_episodes 50

#运行示例
python record_sim_episodes.py --task_name sim_transfer_cube_scripted --dataset_dir data/sim_transfer_cube_scripted --num_episodes 50

若要实时查看仿真环境数据收集效果，可以在上述运行代码后加上：

 --onscreen_render

50个demo以hdf5格式保存至--dataset_dir 

python record_sim_episodes.py --dataset_dir <data save dir> --episode_idx <num>

#运行示例
python visualize_episodes.py --dataset_dir data/sim_transfer_cube_scripted --episode_idx 0

（可选：直接下载项目自带的模拟数据集）

最终，将数据集目录更改成如下格式：

act-plus-plus/data
    ├── sim_insertion_human
    │   ├── episode_0.hdf5
        ├── ...
    ├── sim_insertion_scripted
    │   ├── episode_0.hdf5
        ├── ... 
    ├── sim_transfer_cube_human
    │   ├── episode_0.hdf5
    │   ├── ...
    └── sim_transfer_cube_scripted
        ├── episode_0.hdf5
        ├── ...

2.2训练

i. 修改数据集读取路径，找到constants.py，修改如下

### Task parameters
#DATA_DIR = '/home/zfu/interbotix_ws/src/act/data'
DATA_DIR = 'D:/aloha/data'

DATA_DIR是存放数据集的位置，每个子文件夹对应着训练的子任务，名字一一对应。

ii. 修改源码，找到/detr/model/detr_vae.py文件285行，修改如下

        #encoder = build_transformer(args)
        encoder = build_encoder(args)

iii. 运行训练命令：

python imitate_episodes.py --task_name sim_transfer_cube_scripted --ckpt_dir <ckpt dir> --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 512 --batch_size 8 --dim_feedforward 3200 --num_steps 2000 --lr 1e-5 --seed 0


#运行示例
python imitate_episodes.py --task_name sim_transfer_cube_scripted --ckpt_dir trainings --policy_class ACT --kl_weight 1 --chunk_size 10 --hidden_dim 512 --batch_size 1 --dim_feedforward 3200 --lr 1e-5 --seed 0 --num_steps 5

# 训练时会提示下面内容，由于没有W&B account，直接选择3
wandb: (1) Create a W&B account
wandb: (2) Use an existing W&B account
wandb: (3) Don't visualize my results
wandb: Enter your choice:

出现如下界面训练成功开始：

2.3 可选：使用wandb(官方使用手册)：

i. 注册wandb账号，获取API密钥

ii. 在imitate_episodes.py里插入密钥

os.environ["WANDB_API_KEY"] = '密钥'

iii. 设置proxy

set HTTP_PROXY=http://127.0.0.1:7890
set HTTPS_PROXY=http://127.0.0.1:7890

iv.简单修改imitate_episodes.py并使用

        #wandb.init(project="mobile-aloha2", reinit=True, entity="mobile-aloha2", name=expr_name)
        wandb.init(project="myaloha",name=expr_name)

2.4 测试

python imitate_episodes.py --eval --task_name sim_transfer_cube_scripted --ckpt_dir trainings\test1 --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 512 --batch_size 8 --dim_feedforward 3200 --lr 1e-5 --seed 0 --num_steps 100 --onscreen_render

要评估策略，运行相同的命令，但要添加--eval，这将加载最佳验证检查点。sim_transfer_cube的成功率约为90%，insertion成功率约为50%。可以添加--onscreen_render来查看评估期间的实时渲染。

对于难以建模的真实世界数据，在损失趋于平稳后，至少训练5000次或3-4倍的时间。有关更多信息，请参考调优提示。如果你的ACT策略不稳定，或者在中间停顿，那就训练更长时间吧!成功率和平滑度可以在损失平台期后提高。

2.5 导出eval视频

# 定义输出视频的参数
height, width = 480, 640
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('output_video.avi', fourcc, 30.0, (width, height))

# 在循环中保存每一帧图像
#以下注释为源代码，修改时源码勿动，在下补充输出视频的代码
                #if onscreen_render:
                    #image = env._physics.render(height=480, width=640, camera_id=onscreen_cam)
                    #plt_img.set_data(image)
                    #plt.pause(DT)
                    ##输出视频
                    frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # 将RGB格式转换为BGR格式
                    out.write(frame)

# 释放资源
out.release()

参考博客：

源码复现|10分钟带你复现Mobile ALOHA SIM，超详细教程（附范文+代码） - 知乎 (zhihu.com)