Integration guide

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
import numpy as np

# Load dataset
dataset = LeRobotDataset(
    repo_id="motionledger_franka_v1",
    split="train"
)

# Basic info
print(f"Total episodes: {dataset.num_episodes}")
print(f"Total frames: {len(dataset)}")
print(f"FPS: {dataset.fps}")

# Check features
print(f"\nFeatures: {list(dataset.features.keys())}")

# Sample an episode
sample = dataset[0]
print(f"\nSample keys: {list(sample.keys())}")
print(f"State shape: {sample['observation.state'].shape}")
print(f"Action shape: {sample['action'].shape}")

# Verify normalization stats exist
print(f"\nNorm stats available: {list(dataset.stats.keys())}")
print(f"State mean: {dataset.stats['observation.state']['mean'][:4]}...")
print(f"Action std: {dataset.stats['action']['std'][:4]}...")

# Check image dimensions
for key in sample.keys():
    if 'images' in key:
        print(f"{key} shape: {sample[key].shape}")

import torch
from torch.utils.data import DataLoader
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset

# Load dataset
dataset = LeRobotDataset(
    repo_id="motionledger_franka_v1",
    split="train"
)

# Access normalization statistics
stats = dataset.stats
state_mean = torch.tensor(stats['observation.state']['mean'])
state_std = torch.tensor(stats['observation.state']['std'])
action_q01 = torch.tensor(stats['action']['q01'])
action_q99 = torch.tensor(stats['action']['q99'])

# Define transforms
def normalize_state(state):
    return (state - state_mean) / state_std

def normalize_action_bounds(action):
    # Robust normalization using percentiles
    return 2 * (action - action_q01) / (action_q99 - action_q01) - 1

def compute_delta_action(action, state, delta_mask):
    """Convert absolute action to delta (relative to state)."""
    delta = action.clone()
    delta[..., delta_mask] = action[..., delta_mask] - state[..., delta_mask]
    return delta

# Create DataLoader with transforms
def collate_fn(batch):
    states = torch.stack([normalize_state(b['observation.state']) for b in batch])
    actions = torch.stack([b['action'] for b in batch])

    # Convert to delta
    delta_mask = [True, True, True, True, True, True, True, False]  # Joints only
    delta_actions = compute_delta_action(actions, states, delta_mask)

    # Normalize
    delta_actions = normalize_action_bounds(delta_actions)

    # Process images
    images = {}
    for key in ['observation.images.base_0_rgb', 'observation.images.wrist_0_rgb']:
        if key in batch[0]:
            imgs = torch.stack([b[key] for b in batch])
            imgs = imgs.float() / 255.0  # Normalize to [0, 1]
            images[key] = imgs

    return {
        'state': states,
        'action': delta_actions,
        'images': images,
        'task': [b.get('task', '') for b in batch],
    }

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    collate_fn=collate_fn,
)

# Training loop
for batch in loader:
    state = batch['state']      # [B, T, 8]
    action = batch['action']    # [B, T, 8]
    images = batch['images']    # dict of [B, T, H, W, 3]
    task = batch['task']        # list of str

    # Your model forward pass here
    # ...