Training with PyTorch | Docs

rayform.torch.EpisodeDataset reads a slice directly. Data is fetched on demand and kept in a local cache, so the second epoch reads from disk.

Shell

pip install "rayform[torch]"

A dataset in five lines

train.py

import rayform
from torch.utils.data import DataLoader

client = rayform.Client()
slice_ = client.slices.get("slc_7c1e04aa")

dataset = rayform.torch.EpisodeDataset(
    slice_,
    obs=["cam0/rgb", "wrist/rgb", "joints/position"],
    action="joints/position_delta",
    fps=15,
    horizon=16,            # action chunk length
    frame="ros",
)
loader = DataLoader(dataset, batch_size=64, num_workers=8, shuffle=True, persistent_workers=True)

What a sample contains

Key	Shape	Type	Notes
`cam0/rgb`	`(3, H, W)`	`uint8`	Converted to float in your transforms
`joints/position`	`(J,)`	`float32`	Radians
`action`	`(horizon, J)`	`float32`	Future deltas from the sample time
`t_ns`	`()`	`int64`	Sample time in nanoseconds
`meta`	dict		`slice`, `segment`, `segment_version`, `site`

Workers and caching

The dataset is safe to use with forked workers: each worker opens its own connections after the fork.
The cache defaults to ~/.cache/rayform and can be shared over NFS or Lustre with RAYFORM_CACHE_DIR. Cached files are content-addressed, so they never go stale.
Prefetch a slice before a long run with rayform pull slc_7c1e04aa, then set RAYFORM_OFFLINE=1 to forbid network access during training.

Distributed training

EpisodeDataset shards by rank when torch.distributed is initialised, so each process sees a disjoint part of every epoch:

Python

dataset = rayform.torch.EpisodeDataset(slice_, obs=["cam0/rgb"], action="joints/position_delta", fps=15)
sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=8)

Reproducibility

Log the slice ID with every run. Sampling is deterministic for a given slice, seed and epoch:

Python

dataset.set_epoch(epoch)   # same order for the same seed and epoch
wandb.config.update({"rayform_slice": slice_.id})