具身智能论文和开源系统之RLDS及辅助工具库介绍

本篇文章将向大家介绍RLDS数据集格式定义及其提供的辅助工具库的使用，和COCO(Common Objects in Context，上下文中的常见物体，微软开源的计算机视觉任务数据集)数据集的定义和辅助工具cocotools类似。RLDS(Reinforcement Learning Data Schema)为google团队提出的面向强化学习任务的数据集数据模式的定义，本篇文章将简要介绍该数据模式以及相关辅助数据集处理的工具库及关键函数。

具身智能开源数据集OpenX-Embodiment是通过RLDS的格式构建的，将20多种不同的机器人数据进行了综合，由于原先不同数据集的格式存在着较大差异，如Berkeley Cable Routing采用HDFS格式进行存储，Stanford Hydra Dataset采用多个层级的json格式存储actions和观测数据observations等等，RLDS则定义了统一的存储规则即数据模式，同时在实现上以TFRecord进行二进制序列化，并提供了tensorflow_dataset库将其解析为tf.data.Dataset类方便模型训练和验证。

RLDS以层级字典的形式进行标准化的数据格式定义，强化学习的一次完整采样Rollout一般称为trajectory，其数据形式类似于{“steps”:{“observation”:[obs_t1,obst2,…obstn],”action”:[act_t1,act_t2,…act_tn]}}。以下代码为对OpenX-Embodiment的数据集的一个名为fractal20220817_data的子集采用tensorflow_dataset库进行解析的代码示例。OpenX-Embodiment的数据集在google cloud storage（如gs://gresearch/robotics/fractal20220817_data/0.1.0)上采用TFRecord序列化格式进行存储，关于RLDS的数据集的解释请参考文档[5]。

import tensorflow_datasets as tfds
dataset = 'fractal20220817_data' 
display_key = 'image'

b = tfds.builder_from_directory(builder_dir=dataset2path(dataset))
if display_key not in b.info.features['steps']['observation']:
  raise ValueError(
      f"The key {display_key} was not found in this dataset.\n"
      + "Please choose a different image key to display for this dataset.\n"
      + "Here is the observation spec:\n"
      + str(b.info.features['steps']['observation']))

ds = b.as_dataset(split='train[:10]').shuffle(10)   # take only first 10 episodes
episode = next(iter(ds))
images = [step['observation'][display_key] for step in episode['steps']]
images = [Image.fromarray(image.numpy()) for image in images]
display.Image(as_gif(images))

map() 是python的一个内置函数，但在这里tf.data.Dataset类实现了自定义的map函数可以使用在dataset的对象上。tf.data.Dataset类也实现了数据集采样的函数sample_from_datasets，可以从多个数据集按照权重进行采样，具体更加详细的介绍请参见[5]。

在crossfromer代码实现中，采用了dlimp函数库，其对符合RLDS格式的数据集做了进一步的包装，如在处理时去掉了steps[4]层级，数据的表示更加扁平化，具体的更多介绍可以参考[3][6]。

References

• [1]、计算机视觉任务COCO数据集：COCO – Common Objects in Context
• [2]、OpenX-Embodiment数据集： google-deepmind/open_x_embodiment
• [3]、dlimp函数库： kvablack/dlimp: dataloading is my passion
• [4]、数据扁平化预处理：dlimp/dataset.py#L367
• [5]、关于RLDS的解释：RLDS.md
• [6]、crossformer/data.md