Merge pull request #38 from Eventual-Inc/rewards-demo-read-episodes

feat(rewards): rebuild the demo on lerobot.read_episodes; accept Daft file handles

Author: ykdojo

README.md

@@ -88,10 +88,12 @@ for RL post-training, or to catch mislabeled tasks.
 from daft_physical_ai.rewards import score_rewards
 
 # one row per episode: task text, length, and where its frames live in the video
+# (e.g. from daft.datasets.lerobot.read_episodes - the video column can be a
+# Daft file handle or a local path string)
 df = df.with_column(
     "rewards",
     score_rewards(
-        df["task"], df["length"], df["from_ts"], df["to_ts"], df["video_path"],
+        df["task"], df["length"], df["from_ts"], df["to_ts"], df["video"],
         url="http://localhost:8001",   # any running Robometer eval server
         max_frames=8,                  # frames sampled per episode
     ),

TESTING.md

@@ -146,6 +146,13 @@ exit codes (0/1/2) and `--force`.
   example's server script, the committed example regenerated against it
   (scores identical to the 07-14 run), and a fresh scaffold's `demo.py` run
   end to end (3 episodes, values identical, the closing filter flags ep1).
+- **read_episodes rewrite (2026-07-16)** - the demo rebuilt on
+  `daft.datasets.lerobot.read_episodes` (episode rows + video file handles
+  streamed from the Hub; no `hf_hub_download`, no hardcoded chunk paths) and
+  `score_rewards` extended to accept a Daft file handle: fresh `modal deploy`,
+  the committed example regenerated against it - all 5 episodes' per-frame
+  progress and success identical to the 07-15 run, closing filter still flags
+  ep1 + ep3.
 
 **Known limitation (not a bug):** executing the rewards demo or its regen
 needs a live Robometer eval server (`ROBOMETER_URL`); CI exercises the

daft_physical_ai/_render_rewards.py

@@ -114,11 +114,11 @@ def _demo_cells(config: RewardsDemoConfig) -> list[tuple[str, str]]:
         ("markdown", intro),
         (
             "markdown",
-            "## Setup\n\nInstall with `pip install daft-physical-ai huggingface_hub matplotlib`, then import.",
+            "## Setup\n\nInstall with `pip install daft-physical-ai matplotlib`, then import.",
         ),
         (
             "code",
-            "import daft\nfrom daft import col, lit\n\nfrom daft_physical_ai.rewards import score_rewards",
+            "from daft import col\nfrom daft.datasets import lerobot\n\nfrom daft_physical_ai.rewards import score_rewards",
         ),
         (
             "markdown",
@@ -136,29 +136,16 @@ def _demo_cells(config: RewardsDemoConfig) -> list[tuple[str, str]]:
         ("code", _SERVER_CELL),
         (
             "markdown",
-            "## Fetch the episode metadata and video\n\nLeRobot v3 stores episode metadata as "
-            "parquet and concatenates episodes into shared mp4 files. The first metadata and "
-            "video files cover the first episodes, which is all this demo scores.",
-        ),
-        (
-            "code",
-            "from huggingface_hub import hf_hub_download\n"
-            "\n"
-            'meta_path = hf_hub_download(DATASET, f"{SPLIT}/meta/episodes/chunk-000/file-000.parquet", '
-            'repo_type="dataset")\n'
-            'video_path = hf_hub_download(DATASET, f"{SPLIT}/videos/{VIDEO_KEY}/chunk-000/file-000.mp4", '
-            'repo_type="dataset")',
-        ),
-        (
-            "markdown",
-            "## Build the episode DataFrame\n\nOne row per episode: the task text (from the "
-            "episode's own LeRobot metadata), its length, and where its "
-            "frames live in the video.",
+            "## Build the episode DataFrame\n\nOne row per episode, straight from Daft's LeRobot "
+            "reader: `read_episodes` reads the episode metadata and resolves which shared mp4 "
+            "holds each episode's footage; `include_video_metadata=True` keeps where in that "
+            "file the episode lives (`from_timestamp`/`to_timestamp`). Everything streams from "
+            "the Hub - nothing to download first.",
         ),
         (
             "code",
             "df = (\n"
-            "    daft.read_parquet(meta_path)\n"
+            '    lerobot.read_episodes(f"hf://datasets/{DATASET}/{SPLIT}", include_video_metadata=True)\n'
             '    .sort("episode_index")\n'
             "    .limit(EPISODES)\n"
             "    .select(\n"
@@ -167,24 +154,24 @@ def _demo_cells(config: RewardsDemoConfig) -> list[tuple[str, str]]:
             '        "length",\n'
             '        col(f"videos/{VIDEO_KEY}/from_timestamp").alias("from_ts"),\n'
             '        col(f"videos/{VIDEO_KEY}/to_timestamp").alias("to_ts"),\n'
-            '        lit(video_path).alias("video_path"),\n'
+            '        col(f"videos/{VIDEO_KEY}/video").alias("video"),\n'
             "    )\n"
             ")",
         ),
         (
             "markdown",
             "## Score the episodes\n\n`score_rewards` returns a reward column: it samples "
             "`MAX_FRAMES` frames per episode, decodes them from the episode's segment of the "
-            "video, and asks the server for per-frame progress + success. It's a lazy async "
-            "Daft UDF, so nothing runs until we materialize below - and episodes score "
-            "concurrently when they do.",
+            "video (streamed through the file handle), and asks the server for per-frame "
+            "progress + success. It's a lazy async Daft UDF, so nothing runs until we "
+            "materialize below - and episodes score concurrently when they do.",
         ),
         (
             "code",
             "df = df.with_column(\n"
             '    "rewards",\n'
             "    score_rewards(\n"
-            '        df["task"], df["length"], df["from_ts"], df["to_ts"], df["video_path"],\n'
+            '        df["task"], df["length"], df["from_ts"], df["to_ts"], df["video"],\n'
             "        url=ROBOMETER_URL, max_frames=MAX_FRAMES, headers=HEADERS,\n"
             "    ),\n"
             ")",

daft_physical_ai/cli/rewards.py

@@ -171,7 +171,7 @@ def run(args: argparse.Namespace) -> int:
         print(f"  python {out_dir / 'run_robometer_server.py'}   # any NVIDIA GPU (A10G/L4 fits the 4B bf16)")
         print(f"  uvx modal deploy {out_dir / 'modal_eval_server.py'}   # Modal (uvx modal setup first)")
     print("\nThen run the demo against it (deps fetched on the fly, nothing to install):")
-    withs = "--with daft-physical-ai --with huggingface_hub --with matplotlib"
+    withs = "--with daft-physical-ai --with matplotlib"
     if have_script:
         print(f"  ROBOMETER_URL=http://... uv run {withs} {script_path}")
     if have_nb:

daft_physical_ai/rewards/__init__.py

@@ -25,7 +25,7 @@ def score_rewards(
     length: Expression,
     from_ts: Expression,
     to_ts: Expression,
-    video_path: Expression,
+    video: Expression,
     *,
     url: str,
     max_frames: int = 8,
@@ -46,7 +46,11 @@ def score_rewards(
         length: episode length column (frame count).
         from_ts: episode start timestamp column (seconds, in the video).
         to_ts: episode end timestamp column (seconds, in the video).
-        video_path: path column for the video file holding the episode.
+        video: video column for the file holding the episode - a local path
+            string or a Daft file handle (e.g. the ``videos/{key}/video``
+            column from ``daft.datasets.lerobot.read_episodes``; handles are
+            streamed through Daft's IO layer, so remote ``hf://`` datasets
+            work without downloading).
         url: base URL of a running Robometer eval server (local or remote);
             the pipeline doesn't care what's behind it.
         max_frames: how many frames to sample per episode (default 8, matching
@@ -62,13 +66,13 @@ def score_rewards(
     """
 
     @daft.func(return_dtype=REWARD_DTYPE)
-    async def _score(task: str, length: int, from_ts: float, to_ts: float, video_path: str) -> dict:
+    async def _score(task: str, length: int, from_ts: float, to_ts: float, video) -> dict:
         idxs = sample_indexes(int(length), max_frames)
         # av decode is blocking; keep it off the event loop so requests overlap.
-        frames, refs = await asyncio.to_thread(decode_frames, video_path, float(from_ts), float(to_ts), idxs)
+        frames, refs = await asyncio.to_thread(decode_frames, video, float(from_ts), float(to_ts), idxs)
         npy, sample_json = build_request(frames, task)
         out = await post_request(url, npy, sample_json, headers=headers, timeout_s=timeout_s)
         progress, success = parse_response(out)
         return {"reward_score": progress, "robometer_success": success, "reward_frames": refs}
 
-    return _score(task, length, from_ts, to_ts, video_path)
+    return _score(task, length, from_ts, to_ts, video)

daft_physical_ai/rewards/_robometer.py

@@ -12,9 +12,14 @@
 
 import io
 import json
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 def sample_indexes(length: int, max_frames: int = 8) -> list[int]:
     """Uniformly sample frame indexes across an episode, first + last always included.
@@ -32,22 +37,39 @@ def sample_indexes(length: int, max_frames: int = 8) -> list[int]:
     return sorted({round(i * float(length - 1) / float(max_frames - 1)) for i in range(max_frames)})
 
 
-def decode_frames(video_path: str, from_ts: float, to_ts: float, want: list[int]) -> tuple[np.ndarray, list[dict]]:
+@contextmanager
+def _open_container(video: Any) -> Iterator[Any]:
+    """Yield an av container for a local path string or a Daft file handle.
+
+    A handle (e.g. ``VideoFile`` from ``lerobot.read_episodes``) is streamed
+    through Daft's IO layer via ``.open()``, so remote schemes like ``hf://``
+    work; PyAV alone can only open local paths and its own protocols.
+    """
+    import av
+
+    if isinstance(video, str):
+        with av.open(video) as container:
+            yield container
+    else:
+        with video.open() as f, av.open(f) as container:
+            yield container
+
+
+def decode_frames(video: Any, from_ts: float, to_ts: float, want: list[int]) -> tuple[np.ndarray, list[dict]]:
     """Decode an episode's segment of a concatenated LeRobot mp4 and pick the wanted frames.
 
-    ``want`` holds frame indexes relative to the episode start (``from_ts``).
+    ``video`` is a local path string or a Daft file handle. ``want`` holds
+    frame indexes relative to the episode start (``from_ts``).
     Returns ``(frames [N, H, W, 3] uint8, refs)`` where each ref records the
     frame's relative index and absolute timestamp in seconds.
     """
-    import av
-
     want_set = set(want)
     frames, refs = [], []
-    with av.open(video_path) as container:
+    with _open_container(video) as container:
         stream = container.streams.video[0]
         time_base = stream.time_base
         if time_base is None:
-            raise ValueError(f"video stream in {video_path} has no time base")
+            raise ValueError(f"video stream in {getattr(video, 'path', video)} has no time base")
         container.seek(max(0, int((from_ts - 1.0) / time_base)), stream=stream)
         rel = None
         for frame in container.decode(stream):

examples/README.md

@@ -48,7 +48,7 @@ talks to a Robometer eval server you run - the committed
 ```bash
 # serve (pick one), then run against it:
 ROBOMETER_URL=http://localhost:8001 \
-uv run --with daft-physical-ai --with huggingface_hub --with matplotlib examples/rewards/demo.py
+uv run --with daft-physical-ai --with matplotlib examples/rewards/demo.py
 ```
 
 Generate your own (different dataset, episode count, frame budget):