From 8a779f1e9a908e20f320bdcd22602e79a03b180b Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Sun, 26 Jul 2026 20:09:53 +0200 Subject: [PATCH] eval(rewards): evaluate value function --- pyproject.toml | 1 + .../annotations/steerable_pipeline/config.py | 13 +- .../steerable_pipeline/modules/advantage.py | 59 ++- .../scripts/lerobot_eval_reward_model.py | 351 ++++++++++++++++++ tests/annotations/test_advantage.py | 33 ++ tests/scripts/test_eval_reward_model.py | 50 +++ 6 files changed, 500 insertions(+), 7 deletions(-) create mode 100644 src/lerobot/scripts/lerobot_eval_reward_model.py create mode 100644 tests/scripts/test_eval_reward_model.py diff --git a/pyproject.toml b/pyproject.toml index 9e2eeec98..bf497d637 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -358,6 +358,7 @@ lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main" lerobot-annotate="lerobot.scripts.lerobot_annotate:main" lerobot-rollout="lerobot.scripts.lerobot_rollout:main" lerobot-compute-returns="lerobot.scripts.lerobot_compute_returns:main" +lerobot-eval-reward-model="lerobot.scripts.lerobot_eval_reward_model:main" # ---------------- Tool Configurations ---------------- diff --git a/src/lerobot/annotations/steerable_pipeline/config.py b/src/lerobot/annotations/steerable_pipeline/config.py index b6878508c..e026a5213 100644 --- a/src/lerobot/annotations/steerable_pipeline/config.py +++ b/src/lerobot/annotations/steerable_pipeline/config.py @@ -183,6 +183,11 @@ class AdvantageConfig: # Ignored when constant_value is set. value_function_path: str = "" + # Optional CSV from ``lerobot-eval-reward-model``. When set, annotation + # consumes its per-frame predictions instead of rerunning VF inference. + # This is the recommended path for temporal and multi-camera value models. + predictions_path: str = "" + # Device to run the value function on. device: str = "cuda" @@ -191,9 +196,11 @@ class AdvantageConfig: # 50 = fine-tuning mode: A_t = Σ r_{t:t+N} + V(s_{t+N}) - V(s_t). n_step: int | None = None - # Per-task percentile for binarization threshold ε_ℓ. - # Actions with advantage > ε_ℓ get I_t = True (positive). - threshold_percentile: float = 0.3 + # Percentile for binarization threshold ε_ℓ. Appendix F uses a threshold + # yielding about 40% positive rollout actions during post-training, i.e. + # the 60th percentile. Pre-training uses 0.7 (~30% positive), while the + # slow-but-successful laundry setting uses 0.9 (~10% positive). + threshold_percentile: float = 0.6 # When True, compute a single global threshold across all episodes (paper behavior). # When False, compute threshold per-episode (faster but less accurate). diff --git a/src/lerobot/annotations/steerable_pipeline/modules/advantage.py b/src/lerobot/annotations/steerable_pipeline/modules/advantage.py index 951a3ebe3..008f1373c 100644 --- a/src/lerobot/annotations/steerable_pipeline/modules/advantage.py +++ b/src/lerobot/annotations/steerable_pipeline/modules/advantage.py @@ -25,8 +25,10 @@ Paper reference: pi*0.6, Section IV-B and Appendix F. from __future__ import annotations +import csv import logging from dataclasses import dataclass, field +from pathlib import Path from typing import Any import numpy as np @@ -58,6 +60,7 @@ class AdvantageModule: _preprocessor: Any = field(default=None, init=False, repr=False) _threshold: float | None = field(default=None, init=False, repr=False) _cache: dict = field(default_factory=dict, init=False, repr=False) + _prediction_lookup: dict[tuple[int, int], float] | None = field(default=None, init=False, repr=False) @property def enabled(self) -> bool: @@ -94,7 +97,10 @@ class AdvantageModule: (advantages, intervention_mask) both shape [num_frames]. advantages[t] = A_t, intervention_mask[t] = True if frame is intervention. """ - self._ensure_model_loaded() + if self.config.predictions_path: + self._ensure_predictions_loaded() + else: + self._ensure_model_loaded() df = record.frames_df() num_frames = len(df) @@ -135,6 +141,27 @@ class AdvantageModule: skip_mask: Optional boolean mask [num_frames]. Frames where True are skipped (left as 0.0) to avoid unnecessary inference. """ + if self.config.predictions_path: + self._ensure_predictions_loaded() + assert self._prediction_lookup is not None + missing = [ + frame_index + for frame_index in record.frame_indices + if (record.episode_index, frame_index) not in self._prediction_lookup + ] + if missing: + raise KeyError( + f"Predictions CSV is missing {len(missing)} frame(s) from episode " + f"{record.episode_index}; first missing frame_index={missing[0]}" + ) + return np.asarray( + [ + self._prediction_lookup[(record.episode_index, frame_index)] + for frame_index in record.frame_indices + ], + dtype=np.float32, + ) + df = record.frames_df() num_frames = len(df) values = np.zeros(num_frames, dtype=np.float32) @@ -194,6 +221,27 @@ class AdvantageModule: return values + def _ensure_predictions_loaded(self) -> None: + if self._prediction_lookup is not None: + return + path = Path(self.config.predictions_path) + if not path.is_file(): + raise FileNotFoundError(f"Advantage predictions CSV not found: {path}") + lookup: dict[tuple[int, int], float] = {} + with path.open(newline="") as handle: + reader = csv.DictReader(handle) + required = {"episode_index", "frame_index", "predicted_value"} + missing_columns = required.difference(reader.fieldnames or ()) + if missing_columns: + raise ValueError(f"Predictions CSV is missing required columns: {sorted(missing_columns)}") + for row in reader: + key = (int(row["episode_index"]), int(row["frame_index"])) + if key in lookup: + raise ValueError(f"Duplicate prediction for episode/frame {key}") + lookup[key] = float(row["predicted_value"]) + self._prediction_lookup = lookup + logger.info("Loaded %d per-frame value predictions from %s", len(lookup), path) + def _decode_video_frames( self, record: EpisodeRecord, infer_indices: np.ndarray ) -> tuple[str | None, torch.Tensor | None]: @@ -261,7 +309,7 @@ class AdvantageModule: if self.config.constant_value: return - if not self.config.value_function_path: + if not self.config.value_function_path and not self.config.predictions_path: return logger.info("Computing global advantage threshold (two-pass mode)...") @@ -294,8 +342,11 @@ class AdvantageModule: self._run_constant_mode(record, staging) return - if not self.config.value_function_path: - logger.warning("No value_function_path or constant_value configured; skipping advantage scoring.") + if not self.config.value_function_path and not self.config.predictions_path: + logger.warning( + "No value_function_path, predictions_path, or constant_value configured; " + "skipping advantage scoring." + ) return if record.episode_index in self._cache: diff --git a/src/lerobot/scripts/lerobot_eval_reward_model.py b/src/lerobot/scripts/lerobot_eval_reward_model.py new file mode 100644 index 000000000..79bb327c2 --- /dev/null +++ b/src/lerobot/scripts/lerobot_eval_reward_model.py @@ -0,0 +1,351 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Evaluate a trained reward model against per-frame MC returns. + +The output CSV is also accepted by ``lerobot-annotate +--advantage.predictions_path=...``. This keeps model inference faithful to the +normal LeRobot dataset and processor path, including multi-camera observations, +state normalization, and temporal delta windows. +""" + +from __future__ import annotations + +import argparse +import logging +import math +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from huggingface_hub import snapshot_download +from torch.utils.data import DataLoader +from tqdm import tqdm + +from lerobot.configs.rewards import RewardModelConfig +from lerobot.datasets import LeRobotDataset, LeRobotDatasetMetadata +from lerobot.datasets.factory import resolve_delta_timestamps +from lerobot.processor import PolicyProcessorPipeline, batch_to_transition, transition_to_batch +from lerobot.rewards import make_reward_model +from lerobot.utils.collate import lerobot_collate_fn +from lerobot.utils.constants import ACTION, OBS_STATE, POLICY_PREPROCESSOR_DEFAULT_NAME +from lerobot.utils.hub import find_latest_hub_checkpoint + +logger = logging.getLogger(__name__) + + +def _resolve_model_path(path_or_repo: str) -> Path: + path = Path(path_or_repo) + if path.is_dir(): + if (path / "pretrained_model").is_dir(): + return path / "pretrained_model" + return path + + latest = find_latest_hub_checkpoint(path_or_repo) + if latest is None: + snapshot = Path(snapshot_download(repo_id=path_or_repo, repo_type="model")) + return snapshot + snapshot = Path( + snapshot_download( + repo_id=path_or_repo, + repo_type="model", + allow_patterns=f"{latest}/pretrained_model/*", + ) + ) + return snapshot / latest / "pretrained_model" + + +def _correlation(x: np.ndarray, y: np.ndarray) -> float: + if len(x) < 2 or np.std(x) == 0 or np.std(y) == 0: + return float("nan") + return float(np.corrcoef(x, y)[0, 1]) + + +def _spearman(x: np.ndarray, y: np.ndarray) -> float: + return _correlation( + pd.Series(x).rank(method="average").to_numpy(), pd.Series(y).rank(method="average").to_numpy() + ) + + +def _binary_auc(scores: np.ndarray, labels: np.ndarray) -> float: + labels = labels.astype(bool) + num_positive = int(labels.sum()) + num_negative = len(labels) - num_positive + if num_positive == 0 or num_negative == 0: + return float("nan") + ranks = pd.Series(scores).rank(method="average").to_numpy() + rank_sum_positive = float(ranks[labels].sum()) + return (rank_sum_positive - num_positive * (num_positive + 1) / 2) / (num_positive * num_negative) + + +def _held_out_episodes(metadata: LeRobotDatasetMetadata, eval_split: float) -> list[int] | None: + if eval_split == 0: + return None + if not 0 < eval_split < 1: + raise ValueError(f"eval_split must be in [0,1), got {eval_split}") + task_to_episodes: dict[str, list[int]] = {} + episode_tasks = metadata.episodes["tasks"] + for episode_index in range(metadata.total_episodes): + task = episode_tasks[episode_index][0] if episode_tasks[episode_index] else "" + task_to_episodes.setdefault(task, []).append(episode_index) + held_out: list[int] = [] + for episodes in task_to_episodes.values(): + count = math.ceil(len(episodes) * eval_split) + held_out.extend(episodes[-count:]) + return held_out + + +def _compute_advantages( + target: np.ndarray, + prediction: np.ndarray, + episode_index: np.ndarray, + n_step: int | None, +) -> np.ndarray: + if n_step is None: + return target - prediction + advantage = np.empty_like(target) + for episode in np.unique(episode_index): + indices = np.flatnonzero(episode_index == episode) + for local_index, index in enumerate(indices): + bootstrap_index = local_index + n_step + if bootstrap_index >= len(indices): + advantage[index] = target[index] - prediction[index] + else: + future = indices[bootstrap_index] + advantage[index] = target[index] - target[future] + prediction[future] - prediction[index] + return advantage + + +def _target_distribution(model, target: torch.Tensor, is_terminal: torch.Tensor) -> torch.Tensor: + try: + return model.compute_target_distribution( + target, + is_terminal, + method=model.config.target_method, + use_one_hot_terminal=model.config.use_one_hot_terminal, + ) + except TypeError: + return model.compute_target_distribution(target, is_terminal) + + +def _predict_logits_and_value(model, batch: dict) -> tuple[torch.Tensor, torch.Tensor]: + if hasattr(model, "_vlm_forward"): + logits, predicted = model._vlm_forward(batch) + return logits, predicted.reshape(-1) + if hasattr(model, "_get_value_readout"): + logits = model.value_head(model._get_value_readout(batch)) + probabilities = logits.softmax(-1) + centers = model.value_head.bin_centers.to(probabilities.dtype) + return logits, (probabilities * centers).sum(-1) + raise TypeError(f"{type(model).__name__} does not expose a distributional value readout") + + +def evaluate(args: argparse.Namespace) -> None: + device = torch.device(args.device) + model_path = _resolve_model_path(args.reward_model_path) + logger.info("Loading reward model from %s", model_path) + + config = RewardModelConfig.from_pretrained(model_path) + config.pretrained_path = str(model_path) + config.device = device.type + + metadata = LeRobotDatasetMetadata(args.dataset_repo_id, root=args.root) + delta_timestamps = resolve_delta_timestamps(config, metadata) + episodes = _held_out_episodes(metadata, args.eval_split) + if episodes is not None: + logger.info("Evaluating %d held-out episode(s) (eval_split=%s)", len(episodes), args.eval_split) + dataset = LeRobotDataset( + args.dataset_repo_id, + root=args.root, + episodes=episodes, + delta_timestamps=delta_timestamps, + video_backend=args.video_backend, + return_uint8=True, + ) + dataloader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.num_workers, + collate_fn=lerobot_collate_fn if dataset.meta.has_language_columns else None, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + ) + + model = make_reward_model(config, dataset_meta=metadata).eval() + preprocessor = PolicyProcessorPipeline.from_pretrained( + model_path, + config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json", + overrides={"device_processor": {"device": device.type}}, + to_transition=batch_to_transition, + to_output=transition_to_batch, + ) + + predictions: list[np.ndarray] = [] + targets: list[np.ndarray] = [] + terminals: list[np.ndarray] = [] + episode_indices: list[np.ndarray] = [] + frame_indices: list[np.ndarray] = [] + actions: list[np.ndarray] = [] + states: list[np.ndarray] = [] + interventions: list[np.ndarray] = [] + nll_sum = 0.0 + target_entropy_sum = 0.0 + count = 0 + + with torch.inference_mode(): + for batch in tqdm(dataloader, desc="Evaluating reward model"): + if args.max_samples > 0 and count >= args.max_samples: + break + for camera_key in metadata.camera_keys: + if camera_key in batch and batch[camera_key].dtype == torch.uint8: + batch[camera_key] = batch[camera_key].float().div_(255) + + raw_target = batch["mc_return"].reshape(-1) + raw_terminal = batch["is_terminal"].reshape(-1).bool() + raw_episode = batch["episode_index"].reshape(-1) + raw_frame = batch["frame_index"].reshape(-1) + if args.max_samples > 0: + keep = min(len(raw_target), args.max_samples - count) + if keep < len(raw_target): + batch = { + key: value[:keep] + if isinstance(value, torch.Tensor) + else value[:keep] + if isinstance(value, list) + else value + for key, value in batch.items() + } + raw_target = raw_target[:keep] + raw_terminal = raw_terminal[:keep] + raw_episode = raw_episode[:keep] + raw_frame = raw_frame[:keep] + + if ACTION in batch and isinstance(batch[ACTION], torch.Tensor): + actions.append(batch[ACTION].reshape(len(raw_target), -1).float().cpu().numpy()) + if OBS_STATE in batch and isinstance(batch[OBS_STATE], torch.Tensor): + state = batch[OBS_STATE] + if state.ndim >= 3: + state = state[:, -1] + states.append(state.reshape(len(raw_target), -1).float().cpu().numpy()) + if "intervention" in batch and isinstance(batch["intervention"], torch.Tensor): + interventions.append(batch["intervention"].reshape(-1).bool().cpu().numpy()) + + processed = preprocessor(batch) + with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + logits, predicted = _predict_logits_and_value(model, processed) + + target_on_device = processed["mc_return"].reshape(-1) + terminal_on_device = processed["is_terminal"].reshape(-1).bool() + target_dist = _target_distribution(model, target_on_device, terminal_on_device) + per_sample_nll = -(target_dist * logits.log_softmax(-1)).sum(-1) + entropy = -(target_dist * target_dist.clamp_min(1e-12).log()).sum(-1) + + batch_count = len(predicted) + nll_sum += float(per_sample_nll.sum()) + target_entropy_sum += float(entropy.sum()) + count += batch_count + predictions.append(predicted.float().cpu().numpy()) + targets.append(raw_target.float().cpu().numpy()) + terminals.append(raw_terminal.cpu().numpy()) + episode_indices.append(raw_episode.cpu().numpy()) + frame_indices.append(raw_frame.cpu().numpy()) + + prediction = np.concatenate(predictions) + target = np.concatenate(targets) + terminal = np.concatenate(terminals) + episode_index = np.concatenate(episode_indices).astype(np.int64) + frame_index = np.concatenate(frame_indices).astype(np.int64) + residual = prediction - target + advantage = _compute_advantages(target, prediction, episode_index, args.n_step) + threshold = float(np.percentile(advantage, args.threshold_percentile * 100)) + advantage_label = np.where(advantage > threshold, "positive", "negative") + + terminal_success = np.isclose(target[terminal], 0.0, atol=1e-6) + metrics = { + "samples": len(prediction), + "nll": nll_sum / count, + "target_entropy": target_entropy_sum / count, + "excess_nll": (nll_sum - target_entropy_sum) / count, + "mae": float(np.mean(np.abs(residual))), + "rmse": float(np.sqrt(np.mean(np.square(residual)))), + "bias": float(np.mean(residual)), + "prediction_std": float(np.std(prediction)), + "target_std": float(np.std(target)), + "pearson": _correlation(prediction, target), + "spearman": _spearman(prediction, target), + "terminal_success_auc": _binary_auc(prediction[terminal], terminal_success), + "advantage_threshold": threshold, + "positive_fraction": float(np.mean(advantage_label == "positive")), + } + for name, value in metrics.items(): + logger.info("%s: %s", name, f"{value:.6f}" if isinstance(value, float) else value) + + output_path = Path(args.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output = { + "episode_index": episode_index, + "frame_index": frame_index, + "mc_return": target, + "predicted_value": prediction, + "residual": residual, + "advantage": advantage, + "advantage_label": advantage_label, + "is_terminal": terminal, + } + if actions and states: + action = np.concatenate(actions) + state = np.concatenate(states) + common_dim = min(action.shape[1], state.shape[1]) + output["command_delta_norm"] = np.linalg.norm( + action[:, :common_dim] - state[:, :common_dim], + axis=1, + ) + state_motion = np.zeros(len(state), dtype=np.float32) + same_episode_indices = np.flatnonzero(episode_index[1:] == episode_index[:-1]) + 1 + state_motion[same_episode_indices] = np.linalg.norm( + state[same_episode_indices] - state[same_episode_indices - 1], + axis=1, + ) + output["state_motion_norm"] = state_motion + if interventions and sum(map(len, interventions)) == len(prediction): + output["intervention"] = np.concatenate(interventions) + pd.DataFrame(output).to_csv(output_path, index=False) + logger.info("Wrote per-frame predictions to %s", output_path) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reward-model-path", required=True) + parser.add_argument("--dataset-repo-id", required=True) + parser.add_argument("--root") + parser.add_argument("--output-path", required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--video-backend", default="torchcodec") + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument("--max-samples", type=int, default=0) + parser.add_argument("--eval-split", type=float, default=0.0) + parser.add_argument("--n-step", type=int) + parser.add_argument("--threshold-percentile", type=float, default=0.6) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + evaluate(args) + + +if __name__ == "__main__": + main() diff --git a/tests/annotations/test_advantage.py b/tests/annotations/test_advantage.py index c0f1f6817..ecb87fce7 100644 --- a/tests/annotations/test_advantage.py +++ b/tests/annotations/test_advantage.py @@ -89,6 +89,7 @@ def test_advantage_module_enabled_by_default(): cfg = AdvantageConfig() module = AdvantageModule(config=cfg) assert module.enabled + assert cfg.threshold_percentile == 0.6 def test_run_episode_skips_without_value_function_path(staging: EpisodeStaging): @@ -103,6 +104,38 @@ def test_run_episode_skips_without_value_function_path(staging: EpisodeStaging): assert rows == [] +def test_run_episode_uses_precomputed_predictions(staging: EpisodeStaging, tmp_path: Path): + predictions_path = tmp_path / "predictions.csv" + predictions_path.write_text( + "episode_index,frame_index,predicted_value\n0,0,-0.4\n0,1,-0.4\n0,2,-0.4\n0,3,-0.4\n" + ) + record = _make_record( + num_frames=4, + mc_returns=np.array([-0.7, -0.4, -0.2, -0.1], dtype=np.float32), + ) + module = AdvantageModule( + config=AdvantageConfig( + predictions_path=str(predictions_path), + threshold_percentile=0.5, + ) + ) + + module.run_episode(record, staging) + + rows = staging.read("advantage") + assert [row["content"] for row in rows] == ["negative", "negative", "positive", "positive"] + assert module._model is None + + +def test_predictions_csv_requires_every_frame(staging: EpisodeStaging, tmp_path: Path): + predictions_path = tmp_path / "predictions.csv" + predictions_path.write_text("episode_index,frame_index,predicted_value\n0,0,-0.4\n") + module = AdvantageModule(config=AdvantageConfig(predictions_path=str(predictions_path))) + + with pytest.raises(KeyError, match="missing 1 frame"): + module.run_episode(_make_record(num_frames=2), staging) + + def test_binarization_with_mock_values(staging: EpisodeStaging): """Advantage binarization produces positive/negative labels based on threshold.""" num_frames = 10 diff --git a/tests/scripts/test_eval_reward_model.py b/tests/scripts/test_eval_reward_model.py new file mode 100644 index 000000000..002fb3626 --- /dev/null +++ b/tests/scripts/test_eval_reward_model.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +from lerobot.scripts.lerobot_eval_reward_model import ( + _binary_auc, + _compute_advantages, + _correlation, + _held_out_episodes, + _spearman, +) + + +def test_reward_evaluation_correlations(): + target = np.array([-1.0, -0.5, 0.0]) + prediction = np.array([-0.9, -0.5, -0.1]) + + assert _correlation(prediction, target) == pytest.approx(1.0) + assert _spearman(prediction, target) == pytest.approx(1.0) + + +def test_reward_evaluation_auc(): + labels = np.array([False, False, True, True]) + scores = np.array([-0.9, -0.7, -0.2, -0.1]) + + assert _binary_auc(scores, labels) == pytest.approx(1.0) + + +def test_reward_evaluation_reproduces_training_episode_split(): + metadata = SimpleNamespace( + total_episodes=5, + episodes={"tasks": [["a"], ["a"], ["a"], ["b"], ["b"]]}, + ) + + assert _held_out_episodes(metadata, 0.34) == [1, 2, 4] + assert _held_out_episodes(metadata, 0.0) is None + + +def test_reward_evaluation_n_step_advantages(): + target = np.array([-0.9, -0.6, -0.3, -0.8, -0.4], dtype=np.float32) + prediction = np.array([-0.8, -0.5, -0.2, -0.7, -0.3], dtype=np.float32) + episodes = np.array([0, 0, 0, 1, 1]) + + advantage = _compute_advantages(target, prediction, episodes, n_step=2) + + expected_first = target[0] - target[2] + prediction[2] - prediction[0] + assert advantage[0] == pytest.approx(expected_first) + assert advantage[1] == pytest.approx(target[1] - prediction[1]) + assert advantage[3] == pytest.approx(target[3] - prediction[3])