lerobot中libero数据集格式

数据集导出

import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset

repo_id = "/workspace/libero-dataset"
dataset = LeRobotDataset(repo_id)
sample = dataset[100]
print(sample)

{'observation.images.image': tensor([[[0.0510, 0.0510, 0.0510,  ..., 0.3490, 0.3098, 0.3059],
         [0.0510, 0.0510, 0.0510,  ..., 0.1765, 0.1961, 0.1843],
         [0.0510, 0.0510, 0.0510,  ..., 0.1882, 0.1686, 0.1961],
         ...,
         [0.3255, 0.3216, 0.3176,  ..., 0.2980, 0.2941, 0.2902],
         [0.3059, 0.2980, 0.2941,  ..., 0.3098, 0.3098, 0.3098],
         [0.2941, 0.2863, 0.2863,  ..., 0.3098, 0.3098, 0.3176]],

        [[0.0510, 0.0510, 0.0510,  ..., 0.3569, 0.3176, 0.3137],
         [0.0510, 0.0510, 0.0510,  ..., 0.1843, 0.2039, 0.1922],
         [0.0510, 0.0510, 0.0510,  ..., 0.1961, 0.1765, 0.2039],
         ...,
         [0.2510, 0.2471, 0.2431,  ..., 0.2118, 0.2078, 0.2039],
         [0.2275, 0.2196, 0.2196,  ..., 0.2235, 0.2235, 0.2235],
         [0.2157, 0.2078, 0.2118,  ..., 0.2235, 0.2235, 0.2314]],

        [[0.0510, 0.0510, 0.0510,  ..., 0.3529, 0.3137, 0.3098],
         [0.0510, 0.0510, 0.0510,  ..., 0.1804, 0.2000, 0.1882],
         [0.0510, 0.0510, 0.0510,  ..., 0.1922, 0.1725, 0.2000],
         ...,
         [0.1961, 0.1922, 0.1882,  ..., 0.1608, 0.1569, 0.1529],
         [0.1843, 0.1765, 0.1647,  ..., 0.1686, 0.1686, 0.1686],
         [0.1725, 0.1647, 0.1569,  ..., 0.1686, 0.1686, 0.1765]]]), 'observation.images.image2': tensor([[[0.3882, 0.3882, 0.3882,  ..., 0.3529, 0.3490, 0.3451],
         [0.3804, 0.3765, 0.3725,  ..., 0.3412, 0.3373, 0.3333],
         [0.3725, 0.3647, 0.3569,  ..., 0.3176, 0.3176, 0.3137],
         ...,
         [0.1412, 0.1412, 0.1451,  ..., 0.1490, 0.1490, 0.1490],
         [0.1451, 0.1451, 0.1451,  ..., 0.1490, 0.1490, 0.1490],
         [0.1451, 0.1451, 0.1451,  ..., 0.1490, 0.1490, 0.1490]],

        [[0.2824, 0.2824, 0.2824,  ..., 0.2510, 0.2471, 0.2431],
         [0.2745, 0.2706, 0.2667,  ..., 0.2392, 0.2353, 0.2314],
         [0.2667, 0.2588, 0.2510,  ..., 0.2235, 0.2235, 0.2196],
         ...,
         [0.1412, 0.1412, 0.1451,  ..., 0.1490, 0.1490, 0.1490],
         [0.1451, 0.1451, 0.1451,  ..., 0.1490, 0.1490, 0.1490],
         [0.1451, 0.1451, 0.1451,  ..., 0.1490, 0.1490, 0.1490]],

        [[0.2157, 0.2157, 0.2157,  ..., 0.2000, 0.1961, 0.1922],
         [0.2078, 0.2039, 0.2000,  ..., 0.1882, 0.1843, 0.1804],
         [0.2000, 0.1922, 0.1843,  ..., 0.1765, 0.1765, 0.1725],
         ...,
         [0.1412, 0.1412, 0.1451,  ..., 0.1490, 0.1490, 0.1490],
         [0.1451, 0.1451, 0.1451,  ..., 0.1490, 0.1490, 0.1490],
         [0.1451, 0.1451, 0.1451,  ..., 0.1490, 0.1490, 0.1490]]]), 'observation.state': tensor([-0.0064, -0.2608,  0.5791,  3.0175,  0.1379, -0.3110,  0.0195, -0.0198]), 'action': tensor([ 0.1554,  0.1714,  0.5009, -0.1200,  0.0000, -0.0504, -1.0000]), 'timestamp': tensor(10.), 'frame_index': tensor(100), 'episode_index': tensor(0), 'index': tensor(100), 'task_index': tensor(0), 'task': 'put the white mug on the left plate and put the yellow and white mug on the right plate'}

sample["observation.images.image"].shape

torch.Size([3, 256, 256])

多帧时序采样 (Multi-Frame Temporal Sampling)

delta_timestamps = {
    "observation.images.image": [-0.3,-0.2, -0.1, 0.0]  # 0.2s and 0.1s before current frame
}
dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)

sample = dataset[100]
print(sample["observation.images.image"].shape)  # [T, C, H, W], where T=3

torch.Size([4, 3, 256, 256])

DataLoader 批量数据加载 (Batch Data Loading with DataLoader)

batch_size = 16
data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
device = "cuda" if torch.cuda.is_available() else "cpu"
for batch in data_loader:
    observations = batch["observation.state"].shape
    actions = batch["action"].shape
    images = batch["observation.images.image"].shape
    print(f"observations:{observations}")
    print(f"actions:{actions}")
    print(f"images:{images}")
    break

observations:torch.Size([16, 8])
actions:torch.Size([16, 7])
images:torch.Size([16, 4, 3, 256, 256])

4. 数据集元数据检查 (Dataset Metadata Inspection)

from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata

ds_meta = LeRobotDatasetMetadata(repo_id)

print(f"Total number of episodes: {ds_meta.total_episodes}")
print(f"Average number of frames per episode: {ds_meta.total_frames / ds_meta.total_episodes:.3f}")
print(f"Frames per second used during data collection: {ds_meta.fps}")
print(f"Robot type: {ds_meta.robot_type}")
print(f"keys to access images from cameras: {ds_meta.camera_keys=}\n")

print("Tasks:")
print(ds_meta.tasks)
print("Features:")
print(ds_meta.features)
print(ds_meta)

Total number of episodes: 1693
Average number of frames per episode: 161.527
Frames per second used during data collection: 10.0
Robot type: panda
keys to access images from cameras: ds_meta.camera_keys=['observation.images.image', 'observation.images.image2']

Tasks:
                                                    task_index
task                                                          
put the white mug on the left plate and put the...           0
put the white mug on the plate and put the choc...           1
put the yellow and white mug in the microwave a...           2
turn on the stove and put the moka pot on it                 3
put both the alphabet soup and the cream cheese...           4
put both the alphabet soup and the tomato sauce...           5
put both moka pots on the stove                              6
put both the cream cheese box and the butter in...           7
put the black bowl in the bottom drawer of the ...           8
pick up the book and place it in the back compa...           9
put the bowl on the plate                                   10
put the wine bottle on the rack                             11
open the top drawer and put the bowl inside                 12
put the cream cheese in the bowl                            13
put the wine bottle on top of the cabinet                   14
push the plate to the front of the stove                    15
turn on the stove                                           16
put the bowl on the stove                                   17
put the bowl on top of the cabinet                          18
open the middle drawer of the cabinet                       19
pick up the orange juice and place it in the ba...          20
pick up the ketchup and place it in the basket              21
pick up the cream cheese and place it in the ba...          22
pick up the bbq sauce and place it in the basket            23
pick up the alphabet soup and place it in the b...          24
pick up the milk and place it in the basket                 25
pick up the salad dressing and place it in the ...          26
pick up the butter and place it in the basket               27
pick up the tomato sauce and place it in the ba...          28
pick up the chocolate pudding and place it in t...          29
pick up the black bowl next to the cookie box a...          30
pick up the black bowl in the top drawer of the...          31
pick up the black bowl on the ramekin and place...          32
pick up the black bowl on the stove and place i...          33
pick up the black bowl between the plate and th...          34
pick up the black bowl on the cookie box and pl...          35
pick up the black bowl next to the plate and pl...          36
pick up the black bowl next to the ramekin and ...          37
pick up the black bowl from table center and pl...          38
pick up the black bowl on the wooden cabinet an...          39
Features:
{'observation.images.image': {'dtype': 'image', 'shape': (256, 256, 3), 'names': ['height', 'width', 'channel'], 'fps': 10.0}, 'observation.images.image2': {'dtype': 'image', 'shape': (256, 256, 3), 'names': ['height', 'width', 'channel'], 'fps': 10.0}, 'observation.state': {'dtype': 'float32', 'shape': (8,), 'names': ['state'], 'fps': 10.0}, 'action': {'dtype': 'float32', 'shape': (7,), 'names': ['actions'], 'fps': 10.0}, 'timestamp': {'dtype': 'float32', 'shape': (1,), 'names': None, 'fps': 10.0}, 'frame_index': {'dtype': 'int64', 'shape': (1,), 'names': None, 'fps': 10.0}, 'episode_index': {'dtype': 'int64', 'shape': (1,), 'names': None, 'fps': 10.0}, 'index': {'dtype': 'int64', 'shape': (1,), 'names': None, 'fps': 10.0}, 'task_index': {'dtype': 'int64', 'shape': (1,), 'names': None, 'fps': 10.0}}
LeRobotDatasetMetadata({
    Repository ID: '/workspace/libero-dataset',
    Total episodes: '1693',
    Total frames: '273465',
    Features: '['observation.images.image', 'observation.images.image2', 'observation.state', 'action', 'timestamp', 'frame_index', 'episode_index', 'index', 'task_index']',
})',

数据集子集选择

dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])

# And see how many frames you have:
print(f"Selected episodes: {dataset.episodes}")
print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}")

Selected episodes: [0, 10, 11, 23]
Number of episodes selected: 4
Number of frames selected: 1270

dataset = LeRobotDataset(repo_id)
print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}")

# The previous metadata class is contained in the 'meta' attribute of the dataset:
print(dataset.meta)

Number of episodes selected: 1693
Number of frames selected: 273465
LeRobotDatasetMetadata({
    Repository ID: '/workspace/libero-dataset',
    Total episodes: '1693',
    Total frames: '273465',
    Features: '['observation.images.image', 'observation.images.image2', 'observation.state', 'action', 'timestamp', 'frame_index', 'episode_index', 'index', 'task_index']',
})',

逐帧访问与 Episode 范围查询

dataset.meta.camera_keys

['observation.images.image', 'observation.images.image2']

episode_index = 0
dataset.meta.episodes["dataset_from_index"][episode_index]

0

dataset.meta.episodes["dataset_to_index"][episode_index]

214

from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]

# Then we grab all the image frames from the first camera:
camera_key = dataset.meta.camera_keys[0]
frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)]

# The objects returned by the dataset are all torch.Tensors
print(type(frames[0]))
print(frames[0].shape)

<class 'torch.Tensor'>
torch.Size([3, 256, 256])