Auto-sync: 2026-06-30 04:59:14

dovla_cil/eval/maniskill_policy_rollout.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 import pickle
 from collections import Counter, defaultdict
 from dataclasses import dataclass, replace
@@ -101,6 +102,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_residual_direction: str = "candidate_minus_anchor",
     retrieval_residual_reduce: str = "none",
     retrieval_residual_challenger_types: tuple[str, ...] = (),
+    retrieval_residual_challenger_scales: tuple[float, ...] = (),
     retrieval_residual_challenger_margin: float = 0.0,
     lattice_exclude_types: tuple[str, ...] = (),
     candidate_type_bonuses: dict[str, float] | None = None,
@@ -244,6 +246,9 @@ def evaluate_maniskill_policy_rollout(
         for candidate_type in retrieval_residual_challenger_types
         if candidate_type.strip()
     )
+    retrieval_residual_challenger_scales = tuple(
+        float(scale) for scale in retrieval_residual_challenger_scales
+    )
     candidate_type_bonuses = {
         str(candidate_type): float(bonus)
         for candidate_type, bonus in (candidate_type_bonuses or {}).items()
@@ -367,6 +372,9 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_residual_challenger_types=(
                     retrieval_residual_challenger_types
                 ),
+                retrieval_residual_challenger_scales=(
+                    retrieval_residual_challenger_scales
+                ),
                 retrieval_residual_challenger_margin=(
                     retrieval_residual_challenger_margin
                 ),
@@ -490,6 +498,11 @@ def evaluate_maniskill_policy_rollout(
         )
         if selection_mode == "retrieval_residual"
         else [],
+        "retrieval_residual_challenger_scales": list(
+            retrieval_residual_challenger_scales
+        )
+        if selection_mode == "retrieval_residual"
+        else [],
         "retrieval_residual_challenger_margin": (
             retrieval_residual_challenger_margin
             if selection_mode == "retrieval_residual"
@@ -1179,6 +1192,7 @@ def _evaluate_task_cases(
     retrieval_residual_reduce: str = "none",
     retrieval_residual_action_l2_penalty: float = 0.0,
     retrieval_residual_challenger_types: tuple[str, ...] = (),
+    retrieval_residual_challenger_scales: tuple[float, ...] = (),
     retrieval_residual_challenger_margin: float = 0.0,
     lattice_exclude_types: tuple[str, ...] = (),
     candidate_type_bonuses: dict[str, float] | None = None,
@@ -1252,6 +1266,7 @@ def _evaluate_task_cases(
                 retrieval_residual_scales=retrieval_residual_scales,
                 retrieval_residual_reduce=retrieval_residual_reduce,
                 retrieval_residual_action_l2_penalty=retrieval_residual_action_l2_penalty,
+                retrieval_residual_challenger_scales=retrieval_residual_challenger_scales,
                 action_low=action_low,
                 action_high=action_high,
                 action_candidates=(
@@ -1790,6 +1805,7 @@ def _select_action_chunk(
     retrieval_residual_scales: tuple[float, ...] = (),
     retrieval_residual_reduce: str = "none",
     retrieval_residual_action_l2_penalty: float = 0.0,
+    retrieval_residual_challenger_scales: tuple[float, ...] = (),
     action_low: Any | None = None,
     action_high: Any | None = None,
     action_candidates: Any | None = None,
@@ -1855,6 +1871,7 @@ def _select_action_chunk(
             selection_seed=selection_seed,
             selection_margin=selection_margin,
             residual_aggregate_mask=residual_aggregate_mask,
+            retrieval_residual_challenger_scales=retrieval_residual_challenger_scales,
             challenger_mask=challenger_mask,
             challenger_margin=challenger_margin,
         )
@@ -2256,6 +2273,7 @@ def _select_residual_lattice_action_chunk(
     residual_reduce: str = "none",
     residual_aggregate_mask: Any | None = None,
     action_l2_penalty: float = 0.0,
+    retrieval_residual_challenger_scales: tuple[float, ...] = (),
     challenger_mask: Any | None = None,
     challenger_margin: float = 0.0,
 ) -> tuple[Any, np.ndarray]:
@@ -2303,7 +2321,8 @@ def _select_residual_lattice_action_chunk(
     challenger_mask = _expand_residual_lattice_mask(
         challenger_mask,
         torch=torch,
-        scale_count=len(tuple(float(scale) for scale in residual_scales) or (float(residual_scale),)),
+        scales=scales,
+        include_scales=retrieval_residual_challenger_scales,
         num_gaussian_candidates=num_gaussian_candidates,
         candidate_sigma=candidate_sigma,
     )
@@ -2396,15 +2415,22 @@ def _expand_residual_lattice_mask(
     mask: Any | None,
     *,
     torch: Any,
-    scale_count: int,
+    scales: tuple[float, ...],
+    include_scales: tuple[float, ...] = (),
     num_gaussian_candidates: int,
     candidate_sigma: float,
 ) -> Any | None:
     if mask is None:
         return None
-    expanded = mask
-    if scale_count > 1:
-        expanded = torch.cat([expanded for _ in range(scale_count)], dim=1)
+    scale_values = tuple(float(scale) for scale in scales)
+    if not scale_values:
+        raise ValueError("scales must not be empty")
+    allowed_scales = tuple(float(scale) for scale in include_scales)
+    blocks = [
+        mask if _scale_is_allowed(scale, allowed_scales) else torch.zeros_like(mask)
+        for scale in scale_values
+    ]
+    expanded = torch.cat(blocks, dim=1) if len(blocks) > 1 else blocks[0]
     if num_gaussian_candidates > 1 and candidate_sigma > 0:
         extra = torch.zeros(
             expanded.shape[0],
@@ -2416,6 +2442,15 @@ def _expand_residual_lattice_mask(
     return expanded
 
 
+def _scale_is_allowed(scale: float, allowed_scales: tuple[float, ...]) -> bool:
+    if not allowed_scales:
+        return True
+    return any(
+        math.isclose(float(scale), float(allowed), rel_tol=1.0e-6, abs_tol=1.0e-6)
+        for allowed in allowed_scales
+    )
+
+
 def _align_residual_horizon_to_policy(policy_mean: Any, residuals: Any, *, torch: Any) -> Any:
     if residuals.ndim != 4:
         raise ValueError("residuals must have shape [B,K,H,D]")