PyTorch

PyTorch is not where robot data gets collected or labeled — it is where a captured dataset turns into batches a policy can train on. That step is deceptively easy to get wrong: robot episodes are long, multimodal, and dominated by video, and the naive way to write a Dataset for them (open every file, index every frame, decode on demand) works correctly and trains at a fraction of the speed the hardware allows.

How the data gets there

The two entry points are torch.utils.data.Dataset (map-style, random access via __getitem__) and torch.utils.data.IterableDataset (sequential, via __iter__). Episode-indexed datasets on local or fast network storage — LeRobot's LeRobotDataset is the reference example — use the map-style interface: episodes and frames get a stable global index, __getitem__(i) returns a single training sample, and the DataLoader handles shuffling and batching. LeRobotDataset gives O(1) access to any frame by index, but "O(1)" only holds because it caches a decoder per worker (via torchcodec, falling back to pyav) instead of opening a fresh one on every call — the naive version of the same interface, without that cache, pays a full seek-and-decode on nearly every __getitem__.

For datasets sharded across many files or streamed from object storage, IterableDataset with WebDataset-style tar shards is the better fit — sequential reads instead of random seeks, at the cost of losing per-sample random access and needing more care to keep shuffling and epoch boundaries well defined across multiple dataloader workers.

A minimal episode-chunk sampling Dataset — the shape actual training code takes when the goal is windows of consecutive frames for action chunking rather than single frames:

import torch
from torch.utils.data import Dataset
 
class EpisodeChunkDataset(Dataset):
    """Samples fixed-length, contiguous chunks from within a single episode.
 
    episode_index: dict mapping episode_id -> (video_path, num_frames, action_path)
    chunk_len: number of consecutive frames per training sample
    """
 
    def __init__(self, episode_index, chunk_len=16):
        self.chunk_len = chunk_len
        self.samples = []  # (episode_id, start_frame)
        for ep_id, meta in episode_index.items():
            n = meta["num_frames"]
            for start in range(0, n - chunk_len + 1, chunk_len):
                self.samples.append((ep_id, start))
        self.episode_index = episode_index
        self._decoders = {}  # per-worker decoder cache, keyed by episode_id
 
    def _get_decoder(self, ep_id, video_path):
        # one decoder instance per (worker, episode) pair, reused across calls
        if ep_id not in self._decoders:
            from torchcodec.decoders import VideoDecoder
            self._decoders[ep_id] = VideoDecoder(video_path)
        return self._decoders[ep_id]
 
    def __len__(self):
        return len(self.samples)
 
    def __getitem__(self, idx):
        ep_id, start = self.samples[idx]
        meta = self.episode_index[ep_id]
        decoder = self._get_decoder(ep_id, meta["video_path"])
 
        frames = decoder[start : start + self.chunk_len]         # decode once, contiguous
        actions = torch.load(meta["action_path"])[start : start + self.chunk_len]
        return {"frames": frames, "actions": actions, "episode_id": ep_id}

Each worker process gets its own copy of self._decoders (it is populated lazily after forking), so decoders are never shared across worker boundaries — the pattern that keeps num_workers > 0 safe with stateful decoders.

The workflow in practice

Day to day, the loop looks less like model iteration and more like data-loading iteration until throughput stops being the bottleneck:

• Split by episode before anything else. Compute the train/val split on episode IDs, then derive frame- or chunk-level sample lists from each split independently. Doing it the other way around — splitting a flat list of frames — leaks near-duplicate adjacent frames across the split.
• Compute normalization statistics once, over the training split only. Per-modality mean/std (or min/max for action bounds) gets computed once and cached to disk; recomputing it per run, or computing it over the full dataset including validation episodes, is a common and quiet source of eval-vs-train mismatch.
• Chunk sampling for action-chunked policies. Instead of returning a single timestep, __getitem__ returns a window of chunk_len consecutive frames and actions — the shape ACT, diffusion policy, and most current VLA fine-tuning pipelines expect.
• pin_memory=True and non-blocking transfers. Pinning host memory in the DataLoader and calling .to(device, non_blocking=True) on the batch overlaps the host-to-GPU copy with compute, which matters more for video-heavy batches than typical vision datasets.

Reproducible seeding across workers

A dataloader with num_workers > 0 forks the sampling process, and each fork needs a distinct but deterministic seed — otherwise every worker either reproduces the exact same "random" augmentation (defeating the point of augmenting) or the run becomes unreproducible across restarts because the seed depends on process scheduling. The fix is a worker_init_fn that derives each worker's seed from a single base seed plus its worker ID:

import random
import numpy as np
import torch
from torch.utils.data import DataLoader
 
def worker_init_fn(worker_id):
    worker_seed = (torch.initial_seed() + worker_id) % (2**32)
    np.random.seed(worker_seed)
    random.seed(worker_seed)
 
g = torch.Generator()
g.manual_seed(42)
 
loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
    worker_init_fn=worker_init_fn,
    generator=g,          # controls the main-process shuffling order
)

The generator argument fixes the shuffling order the main process uses to hand indices to workers; worker_init_fn fixes what each worker does with randomness once it has an index. Both are needed — setting only one leaves either the sample order or the augmentation non-reproducible between runs with the same nominal seed.

Gotchas

Video decoders are not fork-safe across num_workers. Opening a decoder in __init__ (before the DataLoader forks worker processes) and sharing it across workers causes corruption or crashes on most backends. Open decoders lazily inside __getitem__ or a per-worker init hook instead, as in the snippet above.

persistent_workers=False (the default) re-spawns the worker pool every epoch. For datasets where each worker's warm state — an open decoder, a cached index — is expensive to rebuild, persistent_workers=True avoids paying that cost every epoch, at the cost of workers holding onto memory between epochs.

Distributed training must shard by episode, not by global sample index. Splitting a flat sample list evenly across DDP ranks can put frames from the same episode on different GPUs in the same step, which is usually harmless for i.i.d. vision data but can interact badly with any batch-level normalization or contrastive objective that assumes independence.

Prefetch factor trades memory for latency hiding. prefetch_factor (samples/batches pre-loaded per worker) defaults to a modest value; raising it helps hide decode latency behind GPU compute but scales host memory use with num_workers × prefetch_factor, which matters when each sample is a chunk of decoded video frames rather than a single small tensor.

MAP-STYLE DATASET

Dataset + __getitem__, random access, shuffling

STREAMING DATASET

IterableDataset, sequential shard reads

SPLIT KEY

episode_id

GPU TRANSFER

pin_memory=True + non_blocking=True