概述与QA

本系列为我在vla上的第一个实践：基于lerobot做π0.5的推理、微调 、仿真。本文主要介绍libero数据集的可视化。

下面列出几个QA来引入：

Q：什么是libero？

A：大家都在训练各种 AI 机器人模型，怎么证明你的模型很牛呢？不能光靠嘴说，得拉到同一个考场里，用同一套卷子考一考，看谁分数高。LIBERO 就是这个考场。

具体来说，它包含了四项东西：

1.任务库：定义了机器人在这个考场上需要完成的任务，具体类别有libero_spatial（空间关系）、libero_object（物体类别）等等。

2.仿真环境：其基于mujoco和robosuite搭建了仿真环境，模型通过简单的接口即可实现可视化。

3.标准数据集：其包含了各种任务的标准演示数据集，这也是我们本章需要可视化的东西。

4.评测指标：libero封装了工具用于模型在各种任务上的评测。

Q：为什么要做这个数据可视化？

A：我们之后将会在libero数据集上做模型的微调，因此必须要先验证这个数据集的格式满足pi0.5模型的输入要求。

---

1. 数据集基本信息

在可视化具体数据之前，先宏观的了解一下libero数据集的组成：

from time import perf_counter
# LeRobotDatasetMetadata不包含真实的数据，只包含数据信息
from lerobot.datasets import LeRobotDatasetMetadata

repo_id = "lerobot/libero"
start = perf_counter()
meta = LeRobotDatasetMetadata(repo_id)
elapsed = perf_counter() - start

print("========== A3 数据集基本信息 ==========")
print(f"数据集ID:{repo_id}")
print(f"格式版本:{meta.info.codebase_version}")
print(f"回合(episode)数:{meta.total_episodes}")
print(f"总帧数:{meta.total_frames}")
print(f"fps:{meta.fps}")
print(f"任务数:{meta.total_tasks}")
print(f"元数据读取耗时:{elapsed}")
print(f"本地元数据目录:{meta.root}")
print("A3_METADATA_OK")

输出：

# ========== A3 数据集基本信息 ==========
# 数据集ID:lerobot/libero
# 格式版本:v3.0
# 回合(episode)数:1693
# 总帧数:273465
# fps:10.0
# 任务数:40
# 元数据读取耗时:3.976761218160391
# 本地元数据目录:/root/autodl-tmp/huggingface/lerobot/hub/datasets--lerobot--libero/snapshots/a1aaacb7f6cd6ee5fb43120f673cebb0cfea7dd4
# A3_METADATA_OK

episode意思是回合，一条完整的演示便可以认为是一个episode，任务数为40，也就是说这个数据集一共有40种任务，这40种任务拥有1693条演示，总帧数为273465。

---

2. 数据集完整字段

回到之前QA中的问题，我们需要检查该数据集中的字段能不能跟pi0.5需要的输入字段对上，运行如下代码：

# LeRobotDatasetMetadata不包含真实的数据，只包含数据信息
from lerobot.datasets import LeRobotDatasetMetadata

repo_id = "lerobot/libero"
meta = LeRobotDatasetMetadata(repo_id)

print("========== A3 数据集全部字段 ==========")
for key, feature in meta.features.items():
    print(
        f"{key}:"
        f"dtype={feature.get('dtype')},"
        f"shape={feature.get('shape')}"
    )

print("========== libero数据集比对π0.5字段 ==========")
print(f"相机键名:{meta.camera_keys}")

for key in [
    "observation.images.image",
    "observation.images.image2",
    "observation.state",
    "action",
]:
    if key in meta.features:
        feature = meta.features[key]
        print(
            f"{key}:"
            f"dtype={feature.get('dtype')},"
            f"shape={feature.get('shape')}"
        )
    else:
        print(f"{key}:缺失")

print("A3_FEATURES_OK")

输出：

========== A3 libero数据集全部字段 ==========
observation.images.image:dtype=video,shape=(256, 256, 3)
observation.images.image2:dtype=video,shape=(256, 256, 3)
observation.state:dtype=float32,shape=(8,)
action:dtype=float32,shape=(7,)
timestamp:dtype=float32,shape=(1,)
frame_index:dtype=int64,shape=(1,)
episode_index:dtype=int64,shape=(1,)
index:dtype=int64,shape=(1,)
task_index:dtype=int64,shape=(1,)
========== libero数据集比对π0.5字段 ==========
相机键名:['observation.images.image', 'observation.images.image2']
observation.images.image:dtype=video,shape=(256, 256, 3)
observation.images.image2:dtype=video,shape=(256, 256, 3)
observation.state:dtype=float32,shape=(8,)
action:dtype=float32,shape=(7,)
A3_FEATURES_OK

可以发现，libero的数据集字段中也包含pi0.5所需要的两路image+state，换句话说，libero数据集完全能适配pi0.5的输入，我们后续不用做数据处理就能让模型在这个数据集上微调。

---

3.任务描述字段

通过下面的代码简单看看实际每一条数据对应的任务描述是什么：

# LeRobotDatasetMetadata不包含真实的数据，只包含数据信息
from lerobot.datasets import LeRobotDatasetMetadata
from lerobot.datasets import LeRobotDataset

repo_id = "lerobot/libero"
meta = LeRobotDatasetMetadata(repo_id)

print("========== 任务描述表 ==========")
print(f"任务文本字段名:{meta.tasks.index.name}")
print(f"任务表中的列:{meta.tasks.columns.tolist()}")
print(f"任务总数:{len(meta.tasks)}")

print("\n前五个任务")
for task_text, row in meta.tasks.head(5).iterrows():
    print(f"task_index={row['task_index']}:{task_text}")

print("\n========== episode 0真实数据 ==========")
dataset = LeRobotDataset(
    repo_id=repo_id,
    episodes=[0],
    download_videos=False,
)

first_frame = dataset.get_raw_item(0)
task_index = int(first_frame["task_index"])
task_text = meta.tasks.iloc[task_index].name

print(f"第一帧的task_index:{task_index}")
print(f"对应的任务文本：{task_text}")
print("A3_TASKS_OK")

输出：

========== 任务描述表 ==========
任务文本字段名: task
任务表中的列: ['task_index']
任务总数: 40

前五个任务
task_index=0: put the white mug on the left plate and put the yellow and white mug on the right plate
task_index=1: put the white mug on the plate and put the chocolate pudding to the right of the plate
task_index=2: put the yellow and white mug in the microwave and close it
task_index=3: turn on the stove and put the moka pot on it
task_index=4: put both the alphabet soup and the cream cheese box in the basket

========== episode 0 ==========
第一帧的task_index: 0
对应的任务文本：put the white mug on the left plate and put the yellow and white mug on the right plate
A3_TASKS_OK

最后的输出表明第一个episode对应的任务是“：put the white mug on the left plate and put the yellow and white mug on the right plate”。

---

4.rerun可视化

通过下面的指令可视化第一个episode的数据：

lerobot-dataset-viz \
  --repo-id lerobot/libero \
  --episode-index 0 \
  --save 1 \
  --output-dir /root/autodl-tmp/workspace/my_test/a3_viz \
  --num-workers 0

最终会得到一个rrd文件，用rerun打开：

---

总结

本文介绍了libero数据集的基本信息、字段、任务描述，并使用rerun进行了单条数据的可视化。

文章仅做学习分享，如有谬误，欢迎指正。