Auto-sync: 2026-06-30 14:47:01

dovla_cil/eval/maniskill_policy_rollout.py

@@ -105,6 +105,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_residual_challenger_types: tuple[str, ...] = (),
     retrieval_residual_challenger_scales: tuple[float, ...] = (),
     retrieval_residual_challenger_margin: float = 0.0,
+    retrieval_residual_challenger_tasks: tuple[str, ...] = (),
     lattice_exclude_types: tuple[str, ...] = (),
     candidate_type_bonuses: dict[str, float] | None = None,
     candidate_type_bonus_components: bool = False,
@@ -255,6 +256,11 @@ def evaluate_maniskill_policy_rollout(
         for candidate_type in retrieval_residual_challenger_types
         if candidate_type.strip()
     )
+    retrieval_residual_challenger_tasks = tuple(
+        task_id.strip()
+        for task_id in retrieval_residual_challenger_tasks
+        if task_id.strip()
+    )
     retrieval_residual_challenger_scales = tuple(
         float(scale) for scale in retrieval_residual_challenger_scales
     )
@@ -388,6 +394,9 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_residual_challenger_margin=(
                     retrieval_residual_challenger_margin
                 ),
+                retrieval_residual_challenger_tasks=(
+                    retrieval_residual_challenger_tasks
+                ),
                 lattice_exclude_types=lattice_exclude_types,
                 candidate_type_bonuses=candidate_type_bonuses,
                 candidate_type_bonus_components=candidate_type_bonus_components,
@@ -529,6 +538,11 @@ def evaluate_maniskill_policy_rollout(
             if selection_mode == "retrieval_residual"
             else 0.0
         ),
+        "retrieval_residual_challenger_tasks": list(
+            retrieval_residual_challenger_tasks
+        )
+        if selection_mode == "retrieval_residual"
+        else [],
         "lattice_exclude_types": list(lattice_exclude_types),
         "candidate_type_bonuses": candidate_type_bonuses,
         "candidate_type_bonus_components": bool(candidate_type_bonus_components),
@@ -1216,6 +1230,7 @@ def _evaluate_task_cases(
     retrieval_residual_challenger_types: tuple[str, ...] = (),
     retrieval_residual_challenger_scales: tuple[float, ...] = (),
     retrieval_residual_challenger_margin: float = 0.0,
+    retrieval_residual_challenger_tasks: tuple[str, ...] = (),
     lattice_exclude_types: tuple[str, ...] = (),
     candidate_type_bonuses: dict[str, float] | None = None,
     candidate_type_bonus_components: bool = False,
@@ -1343,14 +1358,15 @@ def _evaluate_task_cases(
                     else None
                 ),
                 challenger_mask=(
-                    _lattice_candidate_include_mask(
+                    _task_limited_challenger_mask(
                         batch,
+                        task_id=task_id,
                         torch=torch,
                         device=device,
                         include_types=retrieval_residual_challenger_types,
+                        include_tasks=retrieval_residual_challenger_tasks,
                     )
                     if selection_mode == "retrieval_residual"
-                    and retrieval_residual_challenger_types
                     else None
                 ),
                 challenger_margin=retrieval_residual_challenger_margin,
@@ -2798,6 +2814,28 @@ def _lattice_candidate_include_mask(
     return torch.tensor(rows, dtype=torch.bool, device=device)
 
 
+def _task_limited_challenger_mask(
+    batch: list[_RolloutCase],
+    *,
+    task_id: str,
+    torch: Any,
+    device: str,
+    include_types: tuple[str, ...],
+    include_tasks: tuple[str, ...],
+) -> Any | None:
+    if not any(candidate_type for candidate_type in include_types):
+        return None
+    allowed_tasks = {task for task in include_tasks if task}
+    if allowed_tasks and task_id not in allowed_tasks:
+        return None
+    return _lattice_candidate_include_mask(
+        batch,
+        torch=torch,
+        device=device,
+        include_types=include_types,
+    )
+
+
 def _candidate_type_matches_exclusion(candidate_type: str, excluded: set[str]) -> bool:
     if candidate_type in excluded:
         return True