robomimic基础教程（四）——开源数据集

robomimic开源了大量数据集及仿真环境，数据集标准格式为HDF5

目录

一、基础要求

二、使用步骤

1. 下载数据集

2. 后处理

3. 训练

4. 查看训练结果

三、HDF5数据集结构与可视化

1. 数据集结构

（1）根级别（data 组 group）

（2）轨迹（组 group ）

（2.1）属性

（2.2）数据集

（2.3）观测（组 group）

（3）掩码（组 group）

2. 可视化

（1）查看HDF5数据结构

（2）查看数据集轨迹

一、基础要求

编译器：使用pycharm或者Colab notebook均可

pycharm安装在UMI复现基础环境安装配置全流程（二）——实用软件安装及卸载写过

二、使用步骤

robomimic数据集通过记录的环境数据，并用作给定离线RL或IL算法的输入。之后，可以通过以下方式使用robomimic数据集：

1. 下载所需的数据集

2. 对数据集进行后处理，保证与robomimic兼容

3. 用数据集训练机器人

1. 下载数据集

robomimic目前支持以下开箱即用的数据集。点击相应的(1)下载链接下载数据集，点击相应的(2)后处理链接对数据集进行后处理。

也可以在robomimic原程序文件夹中使用download_datasets.py函数下载，例如：

python robomimic/scripts/download_datasets.py --tasks lift --dataset_types ph

数据集存储为datasets/lift/ph/low_dim_v141.hdf5

2. 后处理

如果下载了low_dim或image数据集，那么数据集可以开箱即用!不需要后处理

如果下载了原始数据集，则必须对数据集进行后处理，因为没有存储观测值。必须运行dataset_states_to_obs.py

3. 训练

在下载和后处理之后，使用train.py对数据集进行训练

python train.py --dataset <PATH_TO_POSTPROCESSED_DATASET> --config <PATH_TO_CONFIG>

比如针对1中下载的low_dim_v141.hdf5，可以选择运行behavior cloning (BC)算法训练

python robomimic/scripts/train.py --config robomimic/exps/templates/bc.json --dataset datasets/lift/ph/low_dim_v141.hdf5 --debug

4. 查看训练结果

在tests文件夹中的tmp_model_dir文件夹中包括这几个文件夹，

更多详细内容查看robomimic应用教程（一）——模型训练

三、HDF5数据集结构与可视化

1. 数据集结构

所有后处理的 robomic 兼容数据集被存储为具有高度组织和层次结构的 HDF5 文件

单个数据集是具有以下结构的单个HDF5文件（链接）

HDF5 结构详细说明：

（1）根级别（data 组 group）

• total（属性 attribute）：表示数据集中状态-动作样本的总数，提供整体大小信息

• env_args（属性 attribute）：一个包含环境元数据的 JSON 字符串，记录了数据收集时的环境信息，该元数据包括：

  • env_name：环境或任务的名称。
  • env_type：环境的类型（例如 robosuite）
  • env_kwargs：传递给环境的其他关键字参数，用于配置

（2）轨迹（组 group ）

每个轨迹都被存储为一个组（如 demo_0、demo_1 等），在每个轨迹组内，包含以下内容：

（2.1）属性

• num_samples（属性 attribute）：该轨迹中的状态-动作样本数量

• model_file（属性 attribute）：MJCF MuJoCo 模型的 XML 字符串，对于 robosuite 数据集是特定的，对于非 robosuite 数据集则省略

（2.2）数据集

• states（数据集 dataset）：包含按时间顺序排列的 MuJoCo 状态，形状为 （N, D），其中：

  • N：轨迹中的样本数量。
  • D：状态向量的维度。对于非 robosuite 数据集，可能为空或填充虚拟值

• actions（数据集 dataset）：包含环境中执行的动作，按时间顺序排列。其形状为（N, A），其中：

  • N：轨迹中的样本数量
  • A：动作空间的维度

• rewards（数据集 dataset）：存储轨迹中从环境获得的奖励，形状为 （N, ），表示每个时间步的奖励

• dones（数据集 dataset）：指示在每个动作后，情节是否结束（1表示结束，0表示未结束），形状为 （N, ）

（2.3）观测（组 group）

• obs（组 group）：包含观测键的多个数据集

  • <obs_key_1>（数据集 dataset）：第一个观测键，数据集的名称和形状可能不同（例如agentview_image，形状为（N，84，84，3））
  • <obs_key_2>（数据集 dataset）：第二个观测键，依此类推

• next_obs（组 group）：包含下一时间步的观测键，结构与 obs 组类似

  • <obs_key_1>（数据集）：对应的下一步观测
  • <obs_key_2>（数据集）：依此类推

（3）掩码（组 group）

• mask（组 group）：如果数据集包含过滤键（filter keys），则存在此组，用于选择数据集的子集（例如，验证轨迹）
  • <filter_key_1>（数据集）：包含轨迹标识符的列表，例如 [“demo_0”, “demo_19”, “demo_35”]，表示用于验证的过滤轨迹
  • <filter_key_2>（数据集）：其他过滤键，依此类推

该结构在不同的轨迹（如 demo_0、demo_1 等）中是一致的，每个轨迹都包含类似的属性和数据集格式

这种设置方式可以访问特定轨迹的数据、元数据、观测和环境中的动作，使其适合使用 robomimic 框架进行策略的训练和测试

2. 可视化

（1）查看HDF5数据结构

存储库提供了一个简单的实用程序脚本（get_dataset_info.py）来查看 hdf5 数据集结构和 hdf5 数据集的一些统计信息，脚本显示如下信息：

• 关于轨迹的统计信息（数量、平均长度等）
• 数据集中的筛选键（filter keys）
• 数据集中的环境元数据（environment metadata），用于构建收集数据的相同模拟器环境
• 第一个演示的数据集结构

可以通过 --verbose 参数打印每个过滤键下的演示键列表，以及所有用于演示的数据集结构

例如，使用 tests/assets/test_v141.hdf5 中的存储库打包得到的小型 hdf5 数据集

python get_dataset_info.py --dataset ../../tests/assets/test_v141.hdf5

对于如何编写自定义代码以处理 robomimic 数据集，可查看 jupyter 代码

相关分析在 robomimic应用教程（三）——深入理解robomimic数据集

（2）查看数据集轨迹

对于上面打包的小型 hdf5 数据集（基于 robosuite v1.4.1），使用 playback_dataset.py 脚本查看数据集轨迹

# For the first 5 trajectories, load environment simulator states one-by-one, and render "agentview" and "robot0_eye_in_hand" cameras to video at /tmp/playback_dataset.mp4
$ python playback_dataset.py --dataset ../../tests/assets/test_v141.hdf5 --render_image_names agentview robot0_eye_in_hand --video_path /tmp/playback_dataset.mp4 --n 5

# Directly visualize the image observations in the dataset. This is especially useful for real robot datasets where there is no simulator to use for rendering.
$ python playback_dataset.py --dataset ../../tests/assets/test_v141.hdf5 --use-obs --render_image_names agentview_image --video_path /tmp/obs_trajectory.mp4

# Visualize depth observations as well.
$ python playback_dataset.py --dataset /path/to/dataset.hdf5 --use-obs --render_image_names agentview_image --render_depth_names agentview_depth --video_path /tmp/obs_trajectory.mp4

# Play the dataset actions in the environment to verify that the recorded actions are reasonable.
$ python playback_dataset.py --dataset ../../tests/assets/test_v141.hdf5 --use-actions --render_image_names agentview --video_path /tmp/playback_dataset_with_actions.mp4

# Visualize only the initial demonstration frames.
$ python playback_dataset.py --dataset ../../tests/assets/test_v141.hdf5 --first --render_image_names agentview --video_path /tmp/dataset_task_inits.mp4