简单讲一下Mobile Aloha模拟环境代码

Mobile Aloha模拟环境有两部分代码:

1. https://github.com/tonyzhaozh/act   --- 里面有ACT论文的链接

2. https://github.com/MarkFzp/act-plus-plus   --- 这个是我们要介绍的代码

ACT算法全称是 Action Chunking with Transformers, 大体意思是: 用机械臂的n个帧的视频数据和机械臂状态数据来训练和推理, 使用Transformer神经网络;

简单介绍一下代码结构:

1. assets下面是mujoco的配置文件, 里面包含了机械臂的信息;

2. detr下面是神经网络的部分, 参考论文看代码, 大体是VAE+Transformer, 然后还参考了BERT神经网络的一些东西;

3. constants.py下面是一些配置信息, 列如训练数据的目录, 每个episode有多少帧, 摄像头配置等;

4. ee_sim_env.py是environment, 加载mujoco模拟环境, 还写了reward函数, 训练的时候reward生成success信息, 可以方便的检查训练的结果;

5. imitate_episodes.py里面是创建神经网络, 训练和评估; 每500步会执行一次评估, 评估会把success数据放在Weights and Biases中, 我们可以在W and B上看到success信息, 来判断训练的结果;

6. record_sim_episodes.py是生成训练用的演示数据集; 生成的是hdf5文件, 文件包含了每一帧的视频, 每一帧视频对应的机械臂的状态;

7. scripted_policy.py 是一个编写好的机械臂运行轨迹, 参考self.left_trajectory和self.left_trajectory两个变量, 里面定义了机械臂关键帧的状态,总共400帧, 中间帧的机械臂状态插值生成;

8. visualize_episodes.py可以把hdf5的文件读出来, 生成视频和机械臂400帧的状态;

我稍微改了一下代码, 可以在没有显卡的普通PC上跑起来, 顺便把Weights and Biases改为Aim(Weights and Biases收费, Aim免费);

对于Transfer Cube这个任务, 大约两万多次训练之后, success能达到最高;