Auto-sync: 2026-06-29 07:04:48

dovla_cil/eval/maniskill_policy_rollout.py

@@ -81,6 +81,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_neighbors: int = 1,
     retrieval_metric: str = "raw",
     retrieval_type_min_success: float = 0.0,
+    retrieval_type_success_bonus_scale: float = 0.0,
     retrieval_residual_min_source_progress: float = 0.0,
     retrieval_residual_min_source_advantage: float = -1.0e9,
     retrieval_residual_source_progress_bonus_scale: float = 0.0,
@@ -187,6 +188,8 @@ def evaluate_maniskill_policy_rollout(
         )
     if not 0.0 <= retrieval_type_min_success <= 1.0:
         raise ValueError("retrieval_type_min_success must be in [0, 1]")
+    if retrieval_type_success_bonus_scale < 0:
+        raise ValueError("retrieval_type_success_bonus_scale must be non-negative")
     if not 0.0 <= retrieval_residual_min_source_progress <= 1.0:
         raise ValueError("retrieval_residual_min_source_progress must be in [0, 1]")
     if retrieval_residual_source_progress_bonus_scale < 0:
@@ -264,6 +267,7 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_neighbors=retrieval_neighbors,
                 retrieval_metric=retrieval_metric,
                 retrieval_type_min_success=retrieval_type_min_success,
+                retrieval_type_success_bonus_scale=retrieval_type_success_bonus_scale,
                 retrieval_residual_min_source_progress=retrieval_residual_min_source_progress,
                 retrieval_residual_min_source_advantage=(
                     retrieval_residual_min_source_advantage
@@ -365,6 +369,11 @@ def evaluate_maniskill_policy_rollout(
         "retrieval_type_min_success": retrieval_type_min_success
         if selection_mode == "retrieval_residual"
         else 0.0,
+        "retrieval_type_success_bonus_scale": (
+            retrieval_type_success_bonus_scale
+            if selection_mode == "retrieval_residual"
+            else 0.0
+        ),
         "retrieval_residual_min_source_progress": retrieval_residual_min_source_progress
         if selection_mode == "retrieval_residual"
         else 0.0,
@@ -575,6 +584,7 @@ def _attach_retrieved_residual_candidates(
     retrieval_neighbors: int,
     retrieval_metric: str = "raw",
     retrieval_type_min_success: float = 0.0,
+    retrieval_type_success_bonus_scale: float = 0.0,
     retrieval_residual_min_source_progress: float = 0.0,
     retrieval_residual_min_source_advantage: float = -1.0e9,
     retrieval_residual_source_progress_bonus_scale: float = 0.0,
@@ -587,7 +597,8 @@ def _attach_retrieved_residual_candidates(
         raise ValueError("retrieval_residual currently supports state observations only")
     heldout = set(heldout_group_ids)
     uses_source_bonus = (
-        retrieval_residual_source_progress_bonus_scale > 0
+        retrieval_type_success_bonus_scale > 0
+        or retrieval_residual_source_progress_bonus_scale > 0
         or retrieval_residual_source_score_bonus_scale > 0
         or retrieval_residual_source_advantage_bonus_scale > 0
     )
@@ -644,7 +655,8 @@ def _attach_retrieved_residual_candidates(
             residuals.append(residual.tolist())
             candidate_types.append(f"residual_{record.candidate_type}")
             residual_bonuses.append(
-                float(retrieval_residual_source_progress_bonus_scale) * source_progress
+                float(retrieval_type_success_bonus_scale) * type_success
+                + float(retrieval_residual_source_progress_bonus_scale) * source_progress
                 + float(retrieval_residual_source_score_bonus_scale) * source_score
                 + float(retrieval_residual_source_advantage_bonus_scale) * source_advantage
             )