chore: import upstream snapshot with attribution

examples/python/dataloader/.gitignore

@@ -0,0 +1,3 @@
+data/
+act_checkpoint/
+.venv/

examples/python/dataloader/README.md

@@ -0,0 +1,77 @@
+Train a [LeRobot](https://github.com/huggingface/lerobot) ACT policy using Rerun's experimental PyTorch dataloader, streaming trajectory data directly from a Rerun catalog.
+
+For an explanation of the dataloader API and how the example fits together, see the [Train PyTorch models with the Rerun dataloader](https://rerun.io/docs/howto/train) how-to guide.
+
+## Run the code
+
+### 1. Install dependencies
+
+This example has its own `uv` project, separate from the workspace `.venv`, because LeRobot requires
+Python >=3.12 while the workspace supports older versions.
+
+**Standalone** (sparse-checkout of just this directory, no local Rerun build):
+
+```bash
+uv sync --no-sources --no-dev
+```
+
+**Monorepo dev** (full repo checkout, editable local `rerun-sdk`):
+
+```bash
+cd examples/python/dataloader
+RERUN_ALLOW_MISSING_BIN=1 uv sync
+uv pip install ../../../rerun_py/rerun_dev_fixup
+```
+
+Then either `source .venv/bin/activate` or prefix subsequent commands with `uv run`.
+
+### 2. Start a local Rerun server
+
+In a separate terminal:
+
+```bash
+rerun server
+```
+
+This serves a Rerun server at `rerun+http://127.0.0.1:51234` (the default used by the scripts).
+
+### 3. Prepare and register the dataset
+
+Downloads a LeRobot dataset from HuggingFace, splits it into per-episode RRDs, and registers them as a dataset in the catalog:
+
+```bash
+uv run python prepare_dataset.py
+```
+
+Pass `--repo-id user/other_lerobot_ds` to use a different dataset, or `--catalog-url ""` to skip registration and only write local RRDs.
+
+### 4. Train
+
+```bash
+uv run python train.py
+```
+
+The script streams batches from the catalog, trains an ACT policy for a few epochs, and saves a checkpoint to `act_checkpoint/`.
+
+It accepts a few CLI flags (run `uv run python train.py --help` for the full list):
+
+```bash
+uv run python train.py \
+    --catalog-url rerun+http://127.0.0.1:51234 \
+    --dataset rerun_so101-pick-and-place \
+    --num-segments 3 \
+    --epochs 5 \
+    --batch-size 8 \
+    --num-workers 8 \
+    --lr 1e-5 \
+    --checkpoint-dir act_checkpoint \
+    --dataset-style iterable  # or "map"
+```
+
+Pass `--num-segments 0` to train on all segments in the dataset.
+
+### Training with traces
+
+```sh
+TELEMETRY_ENABLED=true OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://localhost:4317 uv run python train.py
+```

examples/python/dataloader/prepare_dataset.py

@@ -0,0 +1,152 @@
+"""Download a LeRobot dataset from HuggingFace Hub and prepare it for the dataloader.
+
+This script:
+1. Downloads a LeRobot dataset from HuggingFace Hub.
+2. Loads it into Rerun via the built-in LeRobot importer (`log_file_from_path`).
+3. Splits the resulting archive into one RRD per episode.
+4. Registers the per-episode RRDs to a catalog server instance.
+
+"""
+
+from __future__ import annotations
+
+import argparse
+import re
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+import rerun as rr
+
+DEFAULT_REPO_ID = "rerun/so101-pick-and-place"
+DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parent / "data"
+APPLICATION_ID = "lerobot"
+
+_EPISODE_RE = re.compile(r"^(episode_)(\d+)$")
+
+
+def _zero_pad_episode_id(rec_id: str, width: int = 5) -> str:
+    """Turn `episode_1` into `episode_00001` so segments sort lexicographically."""
+    m = _EPISODE_RE.match(rec_id)
+    if m:
+        return f"{m.group(1)}{int(m.group(2)):0{width}d}"
+    return rec_id
+
+
+def download_dataset(repo_id: str, dest: Path) -> Path:
+    """Download a LeRobot dataset from HuggingFace Hub into *dest* and return its path."""
+    print(f"Downloading {repo_id} to {dest} …")
+    local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", local_dir=dest)
+    return Path(local_dir)
+
+
+def lerobot_to_combined_rrd(dataset_dir: Path, combined_rrd: Path) -> None:
+    """Use Rerun's built-in LeRobot importer to turn the dataset into a single RRD."""
+    print(f"Converting {dataset_dir} -> {combined_rrd}")
+    with rr.RecordingStream(APPLICATION_ID) as rec:
+        rec.save(str(combined_rrd))
+        rec.log_file_from_path(str(dataset_dir))
+
+
+def split_into_episode_rrds(combined_rrd: Path, rrd_dir: Path) -> list[Path]:
+    """Split a combined RRD archive into one RRD per episode.
+
+    Returns the paths of the written per-episode RRDs.
+    """
+    rrd_dir.mkdir(parents=True, exist_ok=True)
+
+    reader = rr.experimental.RrdReader(str(combined_rrd))
+    recordings = reader.recordings()
+    print(f"Archive contains {len(recordings)} recordings")
+
+    episode_paths: list[Path] = []
+    for entry in recordings:
+        store = reader.store(store=entry)
+        # Skip metadata-only recordings (e.g. the "root" recording that only carries properties).
+        if not store.schema().entity_paths():
+            continue
+
+        episode_id = _zero_pad_episode_id(entry.recording_id)
+        rrd_path = rrd_dir / f"{episode_id}.rrd"
+
+        with rr.RecordingStream(APPLICATION_ID, recording_id=episode_id, send_properties=False) as rec:
+            rec.save(str(rrd_path))
+            rec.send_chunks(store)
+        episode_paths.append(rrd_path)
+        print(f"  wrote {rrd_path} ({rrd_path.stat().st_size / (1024 * 1024):.1f} MB)")
+
+    return episode_paths
+
+
+def register_to_catalog(
+    rrd_paths: list[Path],
+    *,
+    catalog_url: str,
+    dataset_name: str,
+) -> None:
+    """Register per-episode RRDs to a catalog server instance.
+
+    Uses absolute file:// URIs so the catalog can read the RRDs directly from the local filesystem.
+    """
+    print(f"\nRegistering {len(rrd_paths)} episodes to {catalog_url} as dataset '{dataset_name}' …")
+    client = rr.catalog.CatalogClient(catalog_url)
+    dataset = client.create_dataset(dataset_name, exist_ok=True)
+
+    uris = [f"file://{p.resolve()}" for p in rrd_paths]
+    on_duplicate = rr.catalog.OnDuplicateSegmentLayer(rr.catalog.OnDuplicateSegmentLayer.REPLACE)
+    dataset.register(uris, on_duplicate=on_duplicate).wait()
+    print("  registration done")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        "--repo-id",
+        default=DEFAULT_REPO_ID,
+        help=f"HuggingFace dataset repo id (default: {DEFAULT_REPO_ID}).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=DEFAULT_OUTPUT_DIR,
+        help=f"Directory to store downloaded dataset and output RRDs (default: {DEFAULT_OUTPUT_DIR}).",
+    )
+    parser.add_argument(
+        "--catalog-url",
+        default="rerun+http://127.0.0.1:51234",
+        help="Rerun catalog URL to register episodes with. Pass an empty string to skip registration.",
+    )
+    parser.add_argument(
+        "--dataset-name",
+        default=None,
+        help="Name of the dataset to create/use in the catalog (default: derived from --repo-id).",
+    )
+    parser.add_argument(
+        "--keep-combined",
+        action="store_true",
+        help="Keep the intermediate combined RRD after splitting (useful for debugging).",
+    )
+    args = parser.parse_args()
+
+    output_dir = args.output_dir
+    repo_slug = args.repo_id.replace("/", "_")
+    dataset_dir = output_dir / "lerobot" / repo_slug
+    rrd_dir = output_dir / "rrds" / repo_slug
+    combined_rrd = output_dir / f"{repo_slug}_combined.rrd"
+
+    download_dataset(args.repo_id, dataset_dir)
+    lerobot_to_combined_rrd(dataset_dir, combined_rrd)
+    episode_paths = split_into_episode_rrds(combined_rrd, rrd_dir)
+
+    if not args.keep_combined:
+        combined_rrd.unlink()
+
+    print(f"\nWrote {len(episode_paths)} per-episode RRDs to {rrd_dir}")
+
+    if args.catalog_url:
+        dataset_name = args.dataset_name or repo_slug
+        register_to_catalog(episode_paths, catalog_url=args.catalog_url, dataset_name=dataset_name)
+
+
+if __name__ == "__main__":
+    main()

examples/python/dataloader/pyproject.toml

@@ -0,0 +1,45 @@
+[project]
+name = "dataloader"
+version = "0.1.0"
+readme = "README.md"
+requires-python = ">=3.12,<3.13"
+dependencies = [
+  "rerun-sdk[dataloader,catalog,tracing]",
+  "huggingface-hub>=1.0",
+  "lerobot[dataset]==0.6.0",
+]
+
+[dependency-groups]
+dev = ["mypy==1.19.1"]
+
+[tool.rerun-example]
+# Picked up by scripts/ci/isolated_examples.py and the `py-lint-isolated-examples` pixi task.
+isolated = true
+
+[tool.uv]
+# The example is flat scripts, not a wheel — skip project build, just sync deps.
+package = false
+# pyarrow 24.0.0 segfaults rerun on import and ships an incomplete py.typed that breaks mypy.
+# Isolated examples don't inherit the workspace root's constraint-dependencies, so pin it here.
+# pyarrow arrives transitively (via rerun-sdk), hence a constraint rather than a direct dependency.
+constraint-dependencies = ["pyarrow>=23.0.1,<24"]
+
+# Default `uv sync` uses the in-repo editable rerun-sdk (monorepo dev mode).
+# After syncing, also run `uv pip install ../../../rerun_py/rerun_dev_fixup` to
+# install the .pth shim that makes `import rerun` resolve to the editable source tree.
+#
+# Standalone users (e.g. sparse-checkout of just this example) run instead:
+#   uv sync --no-sources --no-dev
+# That ignores the path source below and resolves `rerun-sdk` from PyPI.
+# rerun-dev-fixup is intentionally absent from this file: uv 0.7.x resolves all
+# dependency groups and extras unconditionally, so any path-only package here would
+# block standalone `--no-sources` resolution.
+[tool.uv.sources]
+rerun-sdk = { path = "../../../rerun_py", editable = true }
+
+# Merged onto the shared base at `../_isolated/mypy.ini` by
+# scripts/ci/isolated_examples.py — list the untyped third-party libs this
+# example actually imports.
+[[tool.mypy.overrides]]
+module = ["lerobot.*", "torch.*", "torchvision.*", "diffusers.*", "accelerate.*"]
+ignore_missing_imports = true

examples/python/dataloader/train.py

@@ -0,0 +1,286 @@
+"""Train a LeRobot ACT policy using the Rerun dataloader.
+
+Demonstrates how to stream robot trajectory data from Rerun's catalog
+into an imitation learning policy (Action Chunking Transformers).
+
+The Rerun dataloader's Field.window feature fetches future action chunks in a single query per batch.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import cast
+
+import torch
+import torch.nn.functional as F
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.act.configuration_act import ACTConfig
+from lerobot.policies.act.modeling_act import ACTPolicy
+from torch.utils.data import DataLoader
+
+from rerun._tracing import tracing_scope, with_tracing
+from rerun.catalog import CatalogClient
+from rerun.experimental.dataloader import (
+    DataSource,
+    Field,
+    NumericDecoder,
+    RerunIterableDataset,
+    RerunMapDataset,
+    VideoFrameDecoder,
+)
+
+CHECKPOINT_DIR = Path(__file__).resolve().parent / "act_checkpoint"
+
+IMAGE_H = 32
+IMAGE_W = 128
+CAMERAS = ("laptop", "phone", "side")
+IMAGE_KEYS = tuple(f"observation.images.{cam}" for cam in CAMERAS)
+
+CHUNK_SIZE = 50
+EPOCHS = 5
+BATCH_SIZE = 8
+LR = 1e-5
+NUM_WORKERS = 4
+FETCH_SIZE = 256
+
+
+class CollateFn:
+    """Picklable collate callable for PyTorch DataLoader multiprocessing."""
+
+    def __init__(self, chunk_size: int, state_dim: int) -> None:
+        self.chunk_size = chunk_size
+        self.state_dim = state_dim
+
+    @with_tracing("CollateFn")
+    def __call__(self, samples: list[dict[str, torch.Tensor | None]]) -> dict[str, torch.Tensor]:
+        # `VideoFrameDecoder` returns `None` when a target precedes the first keyframe; filter those out.
+        complete: list[dict[str, torch.Tensor]] = [
+            cast("dict[str, torch.Tensor]", s) for s in samples if all(s[f"image_{cam}"] is not None for cam in CAMERAS)
+        ]
+        batch_size = len(complete)
+
+        states = torch.stack([s["state"] for s in complete]).float()
+
+        # Future action chunks: reshape windowed flat tensors
+        actions = torch.stack([s["action"].reshape(self.chunk_size, self.state_dim) for s in complete]).float()
+
+        batch: dict[str, torch.Tensor] = {
+            "observation.state": states,
+            "action": actions,
+            "action_is_pad": torch.zeros(batch_size, self.chunk_size, dtype=torch.bool),
+        }
+        # Per-camera images: (3, H, W) uint8 -> float in [0, 1], resized to (IMAGE_H, IMAGE_W)
+        for cam, key in zip(CAMERAS, IMAGE_KEYS):
+            imgs = torch.stack([s[f"image_{cam}"] for s in complete]).float() / 255.0
+            batch[key] = F.interpolate(imgs, size=(IMAGE_H, IMAGE_W), mode="bilinear", align_corners=False)
+        return batch
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        "--catalog-url",
+        default="rerun+http://127.0.0.1:51234",
+        help="Rerun catalog URL",
+    )
+    parser.add_argument(
+        "--dataset",
+        default="rerun_so101-pick-and-place",
+        help="Dataset name in the catalog",
+    )
+    parser.add_argument("--num-segments", type=int, default=3, help="Number of segments to use (0 for all)")
+    parser.add_argument("--epochs", type=int, default=EPOCHS, help="Number of training epochs")
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Training batch size")
+    parser.add_argument("--num-workers", type=int, default=NUM_WORKERS, help="DataLoader worker processes")
+    parser.add_argument(
+        "--fetch-size",
+        type=int,
+        default=FETCH_SIZE,
+        help="Samples fetched per server query for the iterable dataset",
+    )
+    parser.add_argument("--lr", type=float, default=LR, help="Learning rate")
+    parser.add_argument(
+        "--dataset-style",
+        choices=("iterable", "map"),
+        default="iterable",
+        help="Which Rerun dataset class to use: 'iterable' (RerunIterableDataset, internal shuffling) "
+        "or 'map' (RerunMapDataset, random access via DataLoader samplers).",
+    )
+    parser.add_argument(
+        "--checkpoint-dir",
+        type=Path,
+        default=CHECKPOINT_DIR,
+        help="Directory to save the trained policy checkpoint",
+    )
+    return parser.parse_args()
+
+
+@with_tracing("main")
+def main() -> None:
+    args = parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    client = CatalogClient(args.catalog_url)
+    dataset_entry = client.get_dataset(args.dataset)
+
+    all_segments = dataset_entry.segment_ids()
+    segments = all_segments if args.num_segments == 0 else all_segments[: args.num_segments]
+    print(f"Using {len(segments)} segments")
+
+    source = DataSource(dataset_entry, segments=segments)
+
+    fields = {
+        "state": Field("/observation.state:Scalars:scalars", decode=NumericDecoder()),
+        "action": Field(
+            "/action:Scalars:scalars",
+            decode=NumericDecoder(),
+            window=(1, CHUNK_SIZE),
+        ),
+        "image_laptop": Field(
+            "/observation.images.laptop:VideoStream:sample",
+            decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
+        ),
+        "image_phone": Field(
+            "/observation.images.phone:VideoStream:sample",
+            decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
+        ),
+        "image_side": Field(
+            "/observation.images.side:VideoStream:sample",
+            decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
+        ),
+    }
+
+    ds: RerunIterableDataset | RerunMapDataset
+    if args.dataset_style == "map":
+        ds = RerunMapDataset(source=source, index="frame_index", fields=fields)
+    else:
+        ds = RerunIterableDataset(source=source, index="frame_index", fields=fields, fetch_size=args.fetch_size)
+    print(f"Using {args.dataset_style} dataset with {len(ds)} samples (after window trimming)")
+
+    # IterableDataset doesn't support indexing, so probe shape via iteration.
+    state_tensor = next(iter(ds))["state"]
+    assert state_tensor is not None  # NumericDecoder never returns None
+    state_dim = state_tensor.shape[0]
+    action_dim = state_dim
+    print(f"Dimensions: {state_dim=}, {action_dim=}")
+
+    config = ACTConfig(
+        chunk_size=CHUNK_SIZE,
+        n_action_steps=CHUNK_SIZE,
+        use_vae=True,
+        kl_weight=10.0,
+        dim_model=256,
+        n_heads=8,
+        dim_feedforward=1024,
+        n_encoder_layers=4,
+        n_decoder_layers=1,
+        latent_dim=32,
+        n_vae_encoder_layers=4,
+        dropout=0.1,
+        vision_backbone="resnet18",
+        pretrained_backbone_weights=None,
+        normalization_mapping={
+            "STATE": NormalizationMode.MEAN_STD,
+            "VISUAL": NormalizationMode.MEAN_STD,
+            "ACTION": NormalizationMode.MEAN_STD,
+        },
+        input_features={
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
+            **{key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_H, IMAGE_W)) for key in IMAGE_KEYS},
+        },
+        output_features={
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,)),
+        },
+    )
+
+    policy = ACTPolicy(config)
+    policy.train()
+    policy.to(device)
+    print(f"ACT policy created ({sum(p.numel() for p in policy.parameters()):,} parameters, device={device})")
+
+    optimizer = torch.optim.AdamW(
+        policy.get_optim_params(),
+        lr=args.lr,
+        weight_decay=1e-4,
+    )
+
+    collate_fn = CollateFn(CHUNK_SIZE, state_dim)
+    # For the map-style dataset, shuffling is driven by the DataLoader's default RandomSampler.
+    # Swap in `sampler=DistributedSampler(ds)` (and call `sampler.set_epoch(epoch)` each epoch)
+    # for multi-node training, or plug in any other PyTorch sampler.
+    loader = DataLoader(
+        ds,
+        batch_size=args.batch_size,
+        shuffle=isinstance(ds, RerunMapDataset),
+        num_workers=args.num_workers,
+        collate_fn=collate_fn,
+        persistent_workers=True,
+        prefetch_factor=8,
+    )
+
+    num_batches = len(loader)
+    print(f"\nTraining for {args.epochs} epochs, {num_batches} batches/epoch, batch_size={args.batch_size}\n")
+
+    for epoch in range(args.epochs):
+        with tracing_scope(f"epoch {epoch}"):
+            if isinstance(ds, RerunIterableDataset):
+                ds.set_epoch(epoch)
+
+            total_loss = 0.0
+            total_l1 = 0.0
+            total_kld = 0.0
+            n = 0
+
+            t_last_print = time.perf_counter()
+            data_sum = 0.0
+            model_sum = 0.0
+            t_data_start = time.perf_counter()
+            for batch in loader:
+                data_time = time.perf_counter() - t_data_start
+
+                t_model_start = time.perf_counter()
+                batch = {k: v.to(device) for k, v in batch.items()}
+                loss, loss_dict = policy.forward(batch)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                model_time = time.perf_counter() - t_model_start
+
+                total_loss += loss.item()
+                total_l1 += loss_dict["l1_loss"]
+                total_kld += loss_dict.get("kld_loss", 0.0)
+                data_sum += data_time
+                model_sum += model_time
+                n += 1
+
+                if n % 10 == 0 or n == 1:
+                    now = time.perf_counter()
+                    since_last = now - t_last_print
+                    t_last_print = now
+                    print(
+                        f"  epoch {epoch + 1}/{args.epochs}  batch {n}/{num_batches}"
+                        f"  loss={loss.item():.4f}"
+                        f"  data={data_sum:.1f}s model={model_sum:.1f}s"
+                        f"  since_last={since_last:.1f}s",
+                        flush=True,
+                    )
+                    data_sum = 0.0
+                    model_sum = 0.0
+
+                t_data_start = time.perf_counter()
+
+            avg_loss = total_loss / max(n, 1)
+            avg_l1 = total_l1 / max(n, 1)
+            avg_kld = total_kld / max(n, 1)
+            print(f"Epoch {epoch + 1}/{args.epochs}  loss={avg_loss:.4f}  l1={avg_l1:.4f}  kld={avg_kld:.4f}")
+
+    with tracing_scope("save_pretrained"):
+        policy.save_pretrained(str(args.checkpoint_dir))
+    print(f"\nSaved checkpoint to {args.checkpoint_dir}")
+
+
+if __name__ == "__main__":
+    main()