Pi0 & Pi0.5 On Robocasa

Infrastructure

Robocasa Dataset to Lerobot Dataset Converter

HDF5格式的Robocasa数据集需要转换为Lerobot数据集格式以供Pi0和Pi0.5模型使用。

Pi0 & Pi0.5 Inference on Robocasa

在Openpi仓库的基础上，我们加入了src/policies/robocasa_policy.py用于适配robocasa环境，构建了/examples/robocasa的推理环境。

Robocasa结果

针对robocasa数据集，在dp、pi0、pi0.5三个模型上进行微调。

基础微调结果比对

在模型微调实验中，分别设置了30个、100个以及300个demos作为训练样本以探索数据规模对模型性能的影响。各模型所使用的训练轮数设置如下：dp模型训练80k轮，pi0模型在100个demos条件下训练300k轮，在300个demos条件下训练150k轮，pi0.5模型在300个demos条件下训练60k轮。在模型训练过程中，以监测到训练损失稳定为结束训练的依据。

评估过程中，我们发现对于OpenDoubleDoors这个任务，模型的准确率为0%，其主要原因并非模型无法完成任务，而是测试的时间范围不足以使模拟器判定任务完成。因此，将评估过程中的action horizon从默认值调整为1000以扩大任务完成判定时间范围。

为了确保公平性，对100demos和300demos的dp在评估时进行了action horizon的调整，结果显示，dp模型在使用100和300个demos的微调实验中的成功率与之前未调整action horizon的实验相比并未显著提升。

数据集调整训练

方法

基于以上微调实验结果，针对模型性能不足的问题采用了数据集调整的方法以进行优化训练。在调整数据集的过程中，根据RoboCasa数据集中的 原子任务（Atomic Tasks） 特性，将任务动作归类为七大类混合任务，包括：Pick and Place Tasks、Door Tasks、Drawer Tasks、Turning Lever Tasks、Twisting Knob Tasks、Insert Tasks以及Pressing Button Tasks。

在构造上述混合任务的数据集时，采用加权的方式。例如在pi05模型的300demos训练中，Pressing Button Tasks的任务数据集来源于：Coffee Press Button、Turn On Microwave、Turn Off Microwave各300demos，以及其他数据集各50demos，以尽量避免训练过程中动作遗忘的问题。

调整后的混合任务数据集采用同样的norm_state.json，以保证正确归一化。

在训练时，依次对每个混合任务训10k epochs，旨在期望能够强化模型学习到的任务相关的动作（例如插入、扭动杠杆等），随后对全任务集合进行整合训练。

结果

针对300demos的数据集，pi05在基础训练和加权训练下的结果。