eval(rewards): evaluate value function

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

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

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:

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

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