Portando grandes conjuntos de dados para o LeRobot Dataset v3.0

Este tutorial explica como portar conjuntos de dados robóticos em larga escala para o formato LeRobot Dataset v3.0. Usaremos o conjunto de dados DROID 1.0.1 como exemplo principal, que demonstra o processamento de conjuntos de dados com múltiplos terabytes e milhares de shards em clusters SLURM.

Organização de arquivos: v2.1 vs v3.0

O Dataset v3.0 muda fundamentalmente a forma como os dados são organizados e armazenados:

Estrutura v2.1 (Baseada em episódios):

dataset/
├── data/chunk-000/episode_000000.parquet
├── data/chunk-000/episode_000001.parquet
├── videos/chunk-000/camera/episode_000000.mp4
└── meta/episodes.jsonl

Estrutura da v3.0 (baseada em arquivos):

dataset/
├── data/chunk-000/file-000.parquet        # Multiple episodes per file
├── videos/camera/chunk-000/file-000.mp4   # Consolidated video chunks
└── meta/episodes/chunk-000/file-000.parquet  # Structured metadata

Essa transição de arquivos de episódios individuais para blocos baseados em arquivos melhora drasticamente o desempenho e reduz o sobrecarga de armazenamento.

Novidades no Dataset v3.0

O Dataset v3.0 introduz melhorias significativas para o tratamento de grandes conjuntos de dados:

🏗️ Organização de arquivos aprimorada

• Estrutura baseada em arquivos: Os episódios agora são agrupados em arquivos segmentados, em vez de arquivos individuais para cada episódio
• Tamanhos de arquivo configuráveis: para arquivos de dados e vídeos
• Eficiência de armazenamento aprimorada: Melhor compressão e redução de sobrecarga

📊 Gerenciamento moderno de metadados

• Metadados baseados em Parquet: Substituiu JSON Lines pelo formato parquet eficiente
• Acesso estruturado a episódios: Acesso direto a um DataFrame do pandas via dataset.meta.episodes
• Estatísticas por episódio: Rastreamento aprimorado de estatísticas no nível do episódio

🚀 Melhorias de desempenho

• Acesso mapeado em memória: Uso de RAM aprimorado por meio do mapeamento de memória do PyArrow
• Carregamento mais rápido: tempo de inicialização do conjunto de dados reduzido significativamente
• Escala aprimorada: Projetado para conjuntos de dados com milhões de episódios

Pré-requisitos

Antes de portar grandes conjuntos de dados, verifique se você tem:

• LeRobot instalado com suporte à v3.0. Siga nosso Guia de Instalação.
• Armazenamento suficiente: conjuntos de dados brutos podem ser muito grandes (por exemplo, o DROID requer 2 TB)
• Acesso ao cluster (recomendado para grandes conjuntos de dados): SLURM ou agendador de tarefas semelhante
• Dependências específicas do conjunto de dados: Para o DROID, você precisará das utilidades de Dataset do TensorFlow

Entendendo o Conjunto de Dados DROID

DROID 1.0.1 é um excelente exemplo de um conjunto de dados robótico em larga escala:

• Tamanho: 1,7 TB (formato RLDS), 8,7 TB (dados brutos)
• Estrutura: 2048 shards pré-definidos do dataset TensorFlow
• Conteúdo: mais de 76.000 trajetórias de manipulação robótica de robôs Franka Emika Panda
• Escopo: Tarefas de manipulação do mundo real em vários ambientes e objetos
• Formato: Originalmente no formato TensorFlow Records/RLDS, exigindo conversão para o formato LeRobot
• Hospedagem: Google Cloud Storage com acesso público via gsutil

O conjunto de dados contém demonstrações diversas de manipulação com:

• Várias visões de câmera (câmera de pulso, câmeras externas)
• Descrições de tarefas de linguagem natural
• Estado proprioceptivo e ações do robô
• Anotações de sucesso/falha

Esquema de recursos do DROID

DROID_FEATURES = {
    # Episode markers
    "is_first": {"dtype": "bool", "shape": (1,)},
    "is_last": {"dtype": "bool", "shape": (1,)},
    "is_terminal": {"dtype": "bool", "shape": (1,)},

    # Language instructions
    "language_instruction": {"dtype": "string", "shape": (1,)},
    "language_instruction_2": {"dtype": "string", "shape": (1,)},
    "language_instruction_3": {"dtype": "string", "shape": (1,)},

    # Robot state
    "observation.state.gripper_position": {"dtype": "float32", "shape": (1,)},
    "observation.state.cartesian_position": {"dtype": "float32", "shape": (6,)},
    "observation.state.joint_position": {"dtype": "float32", "shape": (7,)},

    # Camera observations
    "observation.images.wrist_left": {"dtype": "video", "shape": (180, 320, 3), "names": ["height", "width", "channels"]},
    "observation.images.exterior_1_left": {"dtype": "video", "shape": (180, 320, 3), "names": ["height", "width", "channels"]},
    "observation.images.exterior_2_left": {"dtype": "video", "shape": (180, 320, 3), "names": ["height", "width", "channels"]},

    # Actions
    "action.gripper_position": {"dtype": "float32", "shape": (1,)},
    "action.cartesian_position": {"dtype": "float32", "shape": (6,)},
    "action.joint_position": {"dtype": "float32", "shape": (7,)},

    # Standard LeRobot format
    "observation.state": {"dtype": "float32", "shape": (8,)},  # joints + gripper
    "action": {"dtype": "float32", "shape": (8,)},  # joints + gripper
}

Abordagem 1: Portabilidade para um único computador

Etapa 1: Instalar Dependências

Especificamente para o DROID:

pip install tensorflow
pip install tensorflow_datasets

Para outros conjuntos de dados, instale os leitores apropriados para o formato da sua fonte.

Etapa 2: Baixar os dados brutos

Baixe o DROID do Google Cloud Storage usando gsutil:

# Install Google Cloud SDK if not already installed
# https://cloud.google.com/sdk/docs/install

# Download the full RLDS dataset (1.7TB)
gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 /your/data/

# Or download just the 100-episode sample (2GB) for testing
gsutil -m cp -r gs://gresearch/robotics/droid_100 /your/data/

Conjuntos de dados grandes exigem tempo e armazenamento consideráveis:

• DROID completo (1,7 TB): vários dias para baixar, dependendo da largura de banda
• Tempo de processamento: 7+ dias para a portabilidade local do conjunto de dados completo
• Tempo de upload: 3+ dias para enviar ao Hugging Face Hub
• Armazenamento local: ~400GB para o formato LeRobot processado

Etapa 3: Portar o Conjunto de Dados

python examples/port_datasets/port_droid.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1 \
    --push-to-hub

Desenvolvimento e Testes

Para desenvolvimento, você pode portar um único shard:

python examples/port_datasets/port_droid.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1_test \
    --num-shards 2048 \
    --shard-index 0

Essa abordagem funciona para conjuntos de dados menores ou testes, mas conjuntos de dados grandes exigem computação em cluster.

Abordagem 2: Portabilidade para cluster SLURM (Recomendado)

Para conjuntos de dados grandes, como o DROID, o processamento paralelo em vários nós reduz drasticamente o tempo de processamento.

Etapa 1: Instale as dependências do cluster

pip install datatrove  # Hugging Face's distributed processing library

Etapa 2: Configure seu ambiente SLURM

Encontre as informações da sua partição:

sinfo --format="%R"  # List available partitions
sinfo -N -p your_partition -h -o "%N cpus=%c mem=%m"  # Check resources

Escolha uma partição de CPU - nenhum GPU é necessário para a portabilidade do conjunto de dados.

Etapa 3: Iniciar trabalhos de portabilidade paralelos

python examples/port_datasets/slurm_port_shards.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name port_droid \
    --partition your_partition \
    --workers 2048 \
    --cpus-per-task 8 \
    --mem-per-cpu 1950M

Diretrizes de Parâmetros

• --workers: Número de jobs paralelos (máximo de 2048 para a contagem de shards do DROID)
• --cpus-per-task: 8 CPUs recomendados para a paralelização da codificação de quadros
• --mem-per-cpu: ~16GB de RAM total (8×1950M) para carregar os quadros brutos

Comece com menos workers (por exemplo, 100) para testar a configuração do seu cluster antes de lançar milhares de jobs.

Etapa 4: Monitorar o progresso

Verifique os jobs em execução:

squeue -u $USER

Monitore o progresso geral:

jobs_status /your/logs

Inspecionar logs de trabalhos individuais:

less /your/logs/port_droid/slurm_jobs/JOB_ID_WORKER_ID.out

Depurar trabalhos com falha:

failed_logs /your/logs/port_droid

Etapa 5: Agregue os Shards

Quando todos os trabalhos de portabilidade forem concluídos:

python examples/port_datasets/slurm_aggregate_shards.py \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name aggr_droid \
    --partition your_partition \
    --workers 2048 \
    --cpus-per-task 8 \
    --mem-per-cpu 1950M

Etapa 6: Enviar para o Hub

python examples/port_datasets/slurm_upload.py \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name upload_droid \
    --partition your_partition \
    --workers 50 \
    --cpus-per-task 4 \
    --mem-per-cpu 1950M

O upload usa menos workers (50), pois é limitado pela rede e não pela computação.

Estrutura de Arquivos do Dataset v3.0

Seu conjunto de dados concluído terá esta estrutura moderna:

dataset/
├── meta/
│   ├── episodes/
│   │   └── chunk-000/
│   │       └── file-000.parquet    # Episode metadata
│   ├── tasks.parquet               # Task definitions
│   ├── stats.json                  # Aggregated statistics
│   └── info.json                   # Dataset information
├── data/
│   └── chunk-000/
│       └── file-000.parquet        # Consolidated episode data
└── videos/
    └── camera_key/
        └── chunk-000/
            └── file-000.mp4        # Consolidated video files

Isso substitui a antiga estrutura de episódio por arquivo por blocos eficientes e de tamanho otimizado.

Migrando do Dataset v2.1

Se você tiver conjuntos de dados existentes no formato v2.1, use a ferramenta de migração:

python src/lerobot/scripts/convert_dataset_v21_to_v30.py \
    --repo-id your_id/existing_dataset

Isso faz automaticamente:

• Converte a estrutura de arquivos para o formato v3.0
• Migra metadados de JSON Lines para parquet
• Agrega estatísticas e cria estatísticas por episódio
• Atualiza as informações de versão

Benefícios de desempenho

O Dataset v3.0 oferece melhorias significativas para grandes conjuntos de dados:

• Carregamento mais rápido: redução de 3 a 5 vezes no tempo de inicialização
• Eficiência de memória: Melhor uso de RAM por meio de mapeamento de memória
• Processamento escalável: Trata milhões de episódios de forma eficiente
• Otimização de armazenamento: Redução do número de arquivos e melhoria na compressão