Auto-sync: 2026-06-29 01:34:49

dovla_cil/eval/maniskill_policy_rollout.py

@@ -83,6 +83,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_type_min_success: float = 0.0,
     retrieval_residual_min_source_progress: float = 0.0,
     retrieval_residual_source_progress_bonus_scale: float = 0.0,
+    retrieval_residual_source_score_bonus_scale: float = 0.0,
     retrieval_residual_scale: float = 1.0,
     retrieval_residual_scales: tuple[float, ...] = (),
     retrieval_residual_anchor: str = "expert",
@@ -187,6 +188,8 @@ def evaluate_maniskill_policy_rollout(
         raise ValueError("retrieval_residual_min_source_progress must be in [0, 1]")
     if retrieval_residual_source_progress_bonus_scale < 0:
         raise ValueError("retrieval_residual_source_progress_bonus_scale must be non-negative")
+    if retrieval_residual_source_score_bonus_scale < 0:
+        raise ValueError("retrieval_residual_source_score_bonus_scale must be non-negative")
     if retrieval_residual_scale < 0:
         raise ValueError("retrieval_residual_scale must be non-negative")
     if any(scale < 0 for scale in retrieval_residual_scales):
@@ -258,6 +261,9 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_residual_source_progress_bonus_scale=(
                     retrieval_residual_source_progress_bonus_scale
                 ),
+                retrieval_residual_source_score_bonus_scale=(
+                    retrieval_residual_source_score_bonus_scale
+                ),
                 retrieval_residual_anchor=retrieval_residual_anchor,
                 retrieval_residual_reduce=retrieval_residual_reduce,
             )
@@ -353,6 +359,11 @@ def evaluate_maniskill_policy_rollout(
             if selection_mode == "retrieval_residual"
             else 0.0
         ),
+        "retrieval_residual_source_score_bonus_scale": (
+            retrieval_residual_source_score_bonus_scale
+            if selection_mode == "retrieval_residual"
+            else 0.0
+        ),
         "retrieval_residual_scale": retrieval_residual_scale
         if selection_mode == "retrieval_residual"
         else 0.0,
@@ -537,12 +548,17 @@ def _attach_retrieved_residual_candidates(
     retrieval_type_min_success: float = 0.0,
     retrieval_residual_min_source_progress: float = 0.0,
     retrieval_residual_source_progress_bonus_scale: float = 0.0,
+    retrieval_residual_source_score_bonus_scale: float = 0.0,
     retrieval_residual_anchor: str = "expert",
     retrieval_residual_reduce: str = "none",
 ) -> list[_RolloutCase]:
     if observation_mode != "state":
         raise ValueError("retrieval_residual currently supports state observations only")
     heldout = set(heldout_group_ids)
+    uses_source_bonus = (
+        retrieval_residual_source_progress_bonus_scale > 0
+        or retrieval_residual_source_score_bonus_scale > 0
+    )
     type_success_rates = _candidate_type_success_rates(dataset, heldout_group_ids=heldout)
     bank: dict[
         str,
@@ -583,6 +599,7 @@ def _attach_retrieved_residual_candidates(
                 continue
             reward = getattr(record, "reward", None)
             source_progress = float(getattr(reward, "progress", 0.0))
+            source_score = _source_reward_score(reward, progress=source_progress)
             if source_progress < retrieval_residual_min_source_progress:
                 continue
             residual = np.asarray(_numeric_action_values(record), dtype=np.float32) - anchor_action
@@ -590,6 +607,7 @@ def _attach_retrieved_residual_candidates(
             candidate_types.append(f"residual_{record.candidate_type}")
             residual_bonuses.append(
                 float(retrieval_residual_source_progress_bonus_scale) * source_progress
+                + float(retrieval_residual_source_score_bonus_scale) * source_score
             )
         feature = np.asarray(
             vectorize_toy_observation(records[0].observation_inline or {}, obs_dim=obs_dim),
@@ -610,9 +628,7 @@ def _attach_retrieved_residual_candidates(
                     candidate_action_values=[zero],
                     candidate_types=["policy_residual"],
                     candidate_score_bonuses=(
-                        [0.0]
-                        if retrieval_residual_source_progress_bonus_scale > 0
-                        else None
+                        [0.0] if uses_source_bonus else None
                     ),
                     candidate_source_group_id=None,
                 )
@@ -664,9 +680,7 @@ def _attach_retrieved_residual_candidates(
                 candidate_action_values=residuals,
                 candidate_types=candidate_types,
                 candidate_score_bonuses=(
-                    residual_bonuses
-                    if retrieval_residual_source_progress_bonus_scale > 0
-                    else None
+                    residual_bonuses if uses_source_bonus else None
                 ),
                 candidate_source_group_id=";".join(source_group_ids),
             )
@@ -801,6 +815,15 @@ def _candidate_type_success_rates(
     }
 
 
+def _source_reward_score(reward: Any, *, progress: float) -> float:
+    if reward is None:
+        return float(progress)
+    score = getattr(reward, "score", None)
+    if score is not None:
+        return float(score)
+    return float(progress) + (1.0 if bool(getattr(reward, "terminal_success", False)) else 0.0)
+
+
 def _nearest_retrieval_entries(
     candidates: list[tuple[Any, np.ndarray, Any, Any]],
     query: np.ndarray,