diff --git a/README.md b/README.md index 4aa071d..8125fcf 100644 --- a/README.md +++ b/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 ), diff --git a/TESTING.md b/TESTING.md index fb7c87e..427c691 100644 --- a/TESTING.md +++ b/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 diff --git a/daft_physical_ai/_render_rewards.py b/daft_physical_ai/_render_rewards.py index d464658..6dd8e8a 100644 --- a/daft_physical_ai/_render_rewards.py +++ b/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,7 +154,7 @@ 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" ")", ), @@ -175,16 +162,16 @@ def _demo_cells(config: RewardsDemoConfig) -> list[tuple[str, str]]: "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" ")", diff --git a/daft_physical_ai/cli/rewards.py b/daft_physical_ai/cli/rewards.py index 3a88977..5c2ddf1 100644 --- a/daft_physical_ai/cli/rewards.py +++ b/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: diff --git a/daft_physical_ai/rewards/__init__.py b/daft_physical_ai/rewards/__init__.py index f001225..5f1c358 100644 --- a/daft_physical_ai/rewards/__init__.py +++ b/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) diff --git a/daft_physical_ai/rewards/_robometer.py b/daft_physical_ai/rewards/_robometer.py index 33f158d..290a5d4 100644 --- a/daft_physical_ai/rewards/_robometer.py +++ b/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): diff --git a/examples/README.md b/examples/README.md index df03660..2ce9cfd 100644 --- a/examples/README.md +++ b/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): diff --git a/examples/rewards/demo.ipynb b/examples/rewards/demo.ipynb index b310651..633ffd8 100644 --- a/examples/rewards/demo.ipynb +++ b/examples/rewards/demo.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "05fd343e", + "id": "c39f29fe", "metadata": {}, "source": [ "# Reward scoring demo - Robometer on nvidia/LIBERO_LeRobot_v3\n", @@ -14,30 +14,30 @@ }, { "cell_type": "markdown", - "id": "b1de9c6f", + "id": "89eb8e9a", "metadata": {}, "source": [ "## Setup\n", "\n", - "Install with `pip install daft-physical-ai huggingface_hub matplotlib`, then import." + "Install with `pip install daft-physical-ai matplotlib`, then import." ] }, { "cell_type": "code", "execution_count": 1, - "id": "c3ee2165", + "id": "1804ed83", "metadata": {}, "outputs": [], "source": [ - "import daft\n", - "from daft import col, lit\n", + "from daft import col\n", + "from daft.datasets import lerobot\n", "\n", "from daft_physical_ai.rewards import score_rewards" ] }, { "cell_type": "markdown", - "id": "bcb7395f", + "id": "c6af108f", "metadata": {}, "source": [ "## Configure\n", @@ -48,7 +48,7 @@ { "cell_type": "code", "execution_count": 2, - "id": "99363ac7", + "id": "ab9e109c", "metadata": {}, "outputs": [], "source": [ @@ -61,7 +61,7 @@ }, { "cell_type": "markdown", - "id": "c1fc3296", + "id": "63c2f54f", "metadata": {}, "source": [ "## Point at your Robometer server\n", @@ -72,7 +72,7 @@ { "cell_type": "code", "execution_count": 3, - "id": "8362829c", + "id": "33bbfb3d", "metadata": {}, "outputs": [], "source": [ @@ -92,46 +92,23 @@ }, { "cell_type": "markdown", - "id": "02029718", - "metadata": {}, - "source": [ - "## Fetch the episode metadata and video\n", - "\n", - "LeRobot 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." - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "935ffc55", - "metadata": {}, - "outputs": [], - "source": [ - "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\")" - ] - }, - { - "cell_type": "markdown", - "id": "c262cef5", + "id": "9b804376", "metadata": {}, "source": [ "## Build the episode DataFrame\n", "\n", - "One row per episode: the task text (from the episode's own LeRobot metadata), its length, and where its frames live in the video." + "One 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." ] }, { "cell_type": "code", - "execution_count": 5, - "id": "a0aeae28", + "execution_count": 4, + "id": "fffe3fce", "metadata": {}, "outputs": [], "source": [ "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", @@ -140,32 +117,32 @@ " \"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", ")" ] }, { "cell_type": "markdown", - "id": "fc0b4347", + "id": "a4e7c6c1", "metadata": {}, "source": [ "## 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." + "`score_rewards` returns a reward column: it samples `MAX_FRAMES` frames per episode, decodes them from the episode's segment of the 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." ] }, { "cell_type": "code", - "execution_count": 6, - "id": "c00c6997", + "execution_count": 5, + "id": "339d37c7", "metadata": {}, "outputs": [], "source": [ "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", ")" @@ -173,7 +150,7 @@ }, { "cell_type": "markdown", - "id": "e0666076", + "id": "cac78225", "metadata": {}, "source": [ "## Read the curves\n", @@ -183,8 +160,8 @@ }, { "cell_type": "code", - "execution_count": 7, - "id": "0cfbd636", + "execution_count": 6, + "id": "79d9505e", "metadata": {}, "outputs": [ { @@ -220,7 +197,7 @@ }, { "cell_type": "markdown", - "id": "c352f5ba", + "id": "65131aed", "metadata": {}, "source": [ "## Plot the progress curves\n", @@ -230,8 +207,8 @@ }, { "cell_type": "code", - "execution_count": 8, - "id": "9d2bde97", + "execution_count": 7, + "id": "3353800c", "metadata": {}, "outputs": [ { @@ -261,7 +238,7 @@ }, { "cell_type": "markdown", - "id": "9168b140", + "id": "31440793", "metadata": {}, "source": [ "## Filter with a Daft query\n", @@ -271,8 +248,8 @@ }, { "cell_type": "code", - "execution_count": 9, - "id": "1c82f06f", + "execution_count": 8, + "id": "40889b82", "metadata": {}, "outputs": [ { @@ -334,7 +311,7 @@ " \n", "
\n", "
\n", - "
\n", + "
\n", "\n", "\n", "\n", @@ -342,12 +319,12 @@ "\n", "
episode_index
Int64		task
String
1
put the white bowl on top of the cabinet
\n", "
\n", - "
\n", + "
\n", "
\n", - "
Cell Details
\n", - " \n", + "
Cell Details
\n", + " \n", "
\n", - "
\n", + "
\n", "

Click on a cell to view its full content

\n", "
\n", "
\n", @@ -356,7 +333,7 @@ "