Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
),
Expand Down
7 changes: 7 additions & 0 deletions TESTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 13 additions & 26 deletions daft_physical_ai/_render_rewards.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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"
Expand All @@ -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"
")",
Expand Down
2 changes: 1 addition & 1 deletion daft_physical_ai/cli/rewards.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
14 changes: 9 additions & 5 deletions daft_physical_ai/rewards/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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)
34 changes: 28 additions & 6 deletions daft_physical_ai/rewards/_robometer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion examples/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading
Loading