# Training with PyTorch

> Stream a slice into a PyTorch DataLoader with the SDK, without exporting files first.

Source: https://rayform.world/docs/guides/pytorch
Last updated: 2026-10-18
Status: Rayform is in private preview. APIs, formats and the SDK described here may change before general availability.

`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.

```bash
pip install "rayform[torch]"
```

## A dataset in five lines

```python title="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})
```