Skip to content

Using LeRobot Datasets

Advanced usage patterns and best practices for LeRobot datasets.

Loading Data

Basic Loading

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset

# From Hugging Face Hub
dataset = LeRobotDataset("lerobot/pusht")

# From local directory
dataset = LeRobotDataset("username/custom_dataset", root="./data")

# Get single frame
frame = dataset[0]

# Get episode
episode_frames = dataset.get_episode(episode_index=0)

Filtering and Slicing

# Filter by task
dataset_task0 = dataset.filter(lambda x: x['task_index'] == 0)

# Get successful episodes only
successful = dataset.filter(lambda x: x['next.success'] == True)

# Time-based slicing
import pandas as pd
df = dataset.to_pandas()
recent_data = df[df['timestamp'] > '2024-01-01']

Custom Transforms

Image Transforms

from torchvision import transforms

class RoboticsTransform:
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.RandomCrop((200, 200)),
            transforms.ColorJitter(brightness=0.2, contrast=0.2),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

    def __call__(self, episode):
        episode['observation.image'] = self.transform(episode['observation.image'])
        return episode

dataset = LeRobotDataset(
    "lerobot/pusht",
    transforms=RoboticsTransform()
)

Action Transforms

class ActionNormalizer:
    def __init__(self, stats):
        self.mean = stats['mean']
        self.std = stats['std']

    def __call__(self, episode):
        action = episode['action']
        episode['action'] = (action - self.mean) / self.std
        return episode

dataset = LeRobotDataset(
    "lerobot/pusht",
    transforms=ActionNormalizer(dataset.stats['action'])
)

Creating Custom Datasets

From Scratch

import numpy as np
from pathlib import Path

class CustomDatasetCreator:
    def __init__(self, output_dir):
        self.output_dir = Path(output_dir)
        self.episodes = []

    def add_episode(self, observations, actions, success=False):
        episode_idx = len(self.episodes)
        episode_data = []

        for frame_idx, (obs, action) in enumerate(zip(observations, actions)):
            frame_data = {
                'episode_index': episode_idx,
                'frame_index': frame_idx,
                'timestamp': frame_idx / 30.0,  # Assuming 30 FPS
                'observation.state': obs['state'],
                'observation.image': f'videos/observation.image/episode_{episode_idx:06d}.mp4#frame={frame_idx}',
                'action': action,
                'next.done': frame_idx == len(observations) - 1,
                'next.success': success if frame_idx == len(observations) - 1 else False
            }
            episode_data.append(frame_data)

        self.episodes.append(episode_data)

    def save(self):
        # Create directories
        (self.output_dir / 'episodes').mkdir(parents=True, exist_ok=True)
        (self.output_dir / 'videos/observation.image').mkdir(parents=True, exist_ok=True)

        # Save episodes as parquet
        for episode_idx, episode in enumerate(self.episodes):
            df = pd.DataFrame(episode)
            df.to_parquet(self.output_dir / f'episodes/episode_{episode_idx:06d}.parquet')

        # Save metadata
        metadata = {
            'fps': 30,
            'total_episodes': len(self.episodes),
            'total_frames': sum(len(ep) for ep in self.episodes),
            'robot_type': 'custom'
        }

        import json
        (self.output_dir / 'meta').mkdir(exist_ok=True)
        with open(self.output_dir / 'meta/info.json', 'w') as f:
            json.dump(metadata, f, indent=2)

From Existing Data

def convert_to_lerobot(source_dir, output_dir, format='pkl'):
    """Convert existing dataset to LeRobot format"""

    creator = CustomDatasetCreator(output_dir)

    # Load episodes
    episode_files = sorted(Path(source_dir).glob(f'*.{format}'))

    for episode_file in episode_files:
        # Load episode (format-specific)
        if format == 'pkl':
            import pickle
            with open(episode_file, 'rb') as f:
                episode = pickle.load(f)

        # Extract observations and actions
        observations = episode['observations']
        actions = episode['actions']
        success = episode.get('success', False)

        # Add to dataset
        creator.add_episode(observations, actions, success)

    # Save
    creator.save()
    print(f"Converted {len(episode_files)} episodes to LeRobot format")

Advanced Features

Multi-Modal Observations

# Dataset with multiple camera views
dataset = LeRobotDataset("username/multi_view_dataset")

# Access different views
frame = dataset[0]
front_camera = frame['observation.image']
wrist_camera = frame['observation.wrist_image']
top_camera = frame['observation.top_image']

Language Instructions

# Dataset with language annotations
dataset = LeRobotDataset("username/language_dataset")

frame = dataset[0]
instruction = frame['language_instruction']  # "pick up the red cube"

Temporal Context

class TemporalDataset(Dataset):
    def __init__(self, base_dataset, history_length=5):
        self.base_dataset = base_dataset
        self.history_length = history_length

    def __getitem__(self, idx):
        # Get current and past frames
        frames = []
        for i in range(self.history_length):
            frame_idx = max(0, idx - i)
            frames.append(self.base_dataset[frame_idx])

        # Stack temporal context
        obs_history = torch.stack([f['observation.image'] for f in frames])
        current_action = frames[0]['action']

        return {
            'observation_history': obs_history,
            'action': current_action
        }

Performance Optimization

Caching

# Cache decoded images in memory
dataset = LeRobotDataset(
    "lerobot/pusht",
    cache_images=True  # Faster loading, more memory
)

# Partial caching
dataset.cache_episodes(range(100))  # Cache first 100 episodes

Parallel Loading

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,        # Parallel data loading
    prefetch_factor=4,    # Prefetch batches
    pin_memory=True       # Faster GPU transfer
)

Video Backend

# Use PyAV for faster video decoding
dataset = LeRobotDataset(
    "lerobot/pusht",
    video_backend="pyav"  # or "torchvision" (default)
)

Sharing Datasets

Push to Hugging Face Hub

from huggingface_hub import HfApi

# Login
from huggingface_hub import login
login(token="your_token")

# Load dataset
dataset = LeRobotDataset("username/my_dataset", root="./data")

# Push to hub
dataset.push_to_hub(
    repo_id="username/my_dataset",
    private=False,  # Public dataset
    commit_message="Initial dataset release"
)

Download from Hub

# Download dataset
dataset = LeRobotDataset("username/my_dataset")

# Automatically downloads to ~/.cache/huggingface/lerobot/

Next Steps