eval(rewards): evaluate value function

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

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])