Auto-sync: 2026-06-28 13:02:56

dovla_cil/eval/maniskill_policy_rollout.py

@@ -65,6 +65,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_metric: str = "raw",
     retrieval_type_min_success: float = 0.0,
     retrieval_residual_scale: float = 1.0,
+    retrieval_residual_scales: tuple[float, ...] = (),
     retrieval_residual_anchor: str = "expert",
     retrieval_residual_reduce: str = "none",
     lattice_exclude_types: tuple[str, ...] = (),
@@ -157,6 +158,8 @@ def evaluate_maniskill_policy_rollout(
         raise ValueError("retrieval_type_min_success must be in [0, 1]")
     if retrieval_residual_scale < 0:
         raise ValueError("retrieval_residual_scale must be non-negative")
+    if any(scale < 0 for scale in retrieval_residual_scales):
+        raise ValueError("retrieval_residual_scales must be non-negative")
     if selection_mode == "policy":
         num_candidates = 1
     checkpoint = torch.load(
@@ -250,6 +253,7 @@ def evaluate_maniskill_policy_rollout(
                 field_optim_trust_radius=field_optim_trust_radius,
                 field_optim_l2_penalty=field_optim_l2_penalty,
                 retrieval_residual_scale=retrieval_residual_scale,
+                retrieval_residual_scales=retrieval_residual_scales,
                 lattice_exclude_types=lattice_exclude_types,
             )
             rows.extend(task_rows)
@@ -304,6 +308,9 @@ def evaluate_maniskill_policy_rollout(
         "retrieval_residual_scale": retrieval_residual_scale
         if selection_mode == "retrieval_residual"
         else 0.0,
+        "retrieval_residual_scales": list(retrieval_residual_scales)
+        if selection_mode == "retrieval_residual"
+        else [],
         "retrieval_residual_anchor": retrieval_residual_anchor
         if selection_mode == "retrieval_residual"
         else "none",
@@ -662,6 +669,7 @@ def _evaluate_task_cases(
     field_optim_trust_radius: float = 0.5,
     field_optim_l2_penalty: float = 0.0,
     retrieval_residual_scale: float = 1.0,
+    retrieval_residual_scales: tuple[float, ...] = (),
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
@@ -727,6 +735,7 @@ def _evaluate_task_cases(
                 field_optim_trust_radius=field_optim_trust_radius,
                 field_optim_l2_penalty=field_optim_l2_penalty,
                 retrieval_residual_scale=retrieval_residual_scale,
+                retrieval_residual_scales=retrieval_residual_scales,
                 action_low=action_low,
                 action_high=action_high,
                 action_candidates=(
@@ -796,6 +805,9 @@ def _evaluate_task_cases(
                             selected_index=int(candidate_index[index]),
                             selection_mode=selection_mode,
                             prepended_policy_candidate=prepend_policy_candidate,
+                            residual_scale_count=_residual_scale_count(
+                                retrieval_residual_scales
+                            ),
                         ),
                         "selected_candidate_index": int(candidate_index[index]),
                         "lattice_candidate_count": _effective_lattice_candidate_count(
@@ -804,6 +816,9 @@ def _evaluate_task_cases(
                             num_candidates=num_candidates,
                             candidate_sigma=candidate_sigma,
                             prepended_policy_candidate=prepend_policy_candidate,
+                            residual_scale_count=_residual_scale_count(
+                                retrieval_residual_scales
+                            ),
                         ),
                         "candidate_source_group_id": case.candidate_source_group_id,
                     }
@@ -830,6 +845,7 @@ def _select_action_chunk(
     field_optim_trust_radius: float = 0.5,
     field_optim_l2_penalty: float = 0.0,
     retrieval_residual_scale: float = 1.0,
+    retrieval_residual_scales: tuple[float, ...] = (),
     action_low: Any | None = None,
     action_high: Any | None = None,
     action_candidates: Any | None = None,
@@ -881,6 +897,7 @@ def _select_action_chunk(
             action_high=action_high,
             candidate_mask=candidate_mask,
             residual_scale=retrieval_residual_scale,
+            residual_scales=retrieval_residual_scales,
             num_gaussian_candidates=num_candidates,
             candidate_sigma=candidate_sigma,
             selection_seed=selection_seed,
@@ -1194,13 +1211,22 @@ def _select_residual_lattice_action_chunk(
     candidate_sigma: float,
     selection_seed: int,
     selection_margin: float = 0.0,
+    residual_scales: tuple[float, ...] = (),
 ) -> tuple[Any, np.ndarray]:
     if action_residuals.ndim != 4:
         raise ValueError("action_residuals must have shape [B,K,H,D]")
-    candidates = policy_mean.unsqueeze(1) + float(residual_scale) * action_residuals.to(
+    residuals = action_residuals.to(
         device=policy_mean.device,
         dtype=policy_mean.dtype,
     )
+    scales = tuple(float(scale) for scale in residual_scales) or (float(residual_scale),)
+    candidate_blocks = [
+        policy_mean.unsqueeze(1) + scale * residuals
+        for scale in scales
+    ]
+    candidates = torch.cat(candidate_blocks, dim=1)
+    if candidate_mask is not None and len(scales) > 1:
+        candidate_mask = torch.cat([candidate_mask for _ in scales], dim=1)
     if num_gaussian_candidates > 1 and candidate_sigma > 0:
         generator = torch.Generator(device=policy_mean.device)
         generator.manual_seed(int(selection_seed))
@@ -1253,15 +1279,22 @@ def _effective_lattice_candidate_count(
     num_candidates: int,
     candidate_sigma: float,
     prepended_policy_candidate: bool = False,
+    residual_scale_count: int = 1,
 ) -> int:
     count = len(case.candidate_action_values)
     if prepended_policy_candidate and selection_mode in {"lattice", "retrieval_lattice"}:
         count += 1
+    if selection_mode == "retrieval_residual":
+        count *= max(1, int(residual_scale_count))
     if selection_mode == "retrieval_residual" and num_candidates > 1 and candidate_sigma > 0:
         count += num_candidates - 1
     return count
 
 
+def _residual_scale_count(residual_scales: tuple[float, ...]) -> int:
+    return max(1, len(residual_scales))
+
+
 def _lattice_candidate_mask(
     batch: list[_RolloutCase],
     *,
@@ -1285,6 +1318,7 @@ def _selected_candidate_type(
     selected_index: int,
     selection_mode: str,
     prepended_policy_candidate: bool = False,
+    residual_scale_count: int = 1,
 ) -> str:
     if selection_mode == "policy":
         return "policy_continuous"
@@ -1293,8 +1327,10 @@ def _selected_candidate_type(
     if selection_mode == "field_optim":
         return "field_optim_selected"
     if selection_mode == "retrieval_residual":
-        if 0 <= selected_index < len(case.candidate_types):
-            return f"retrieval_residual_{case.candidate_types[selected_index]}"
+        residual_count = len(case.candidate_types)
+        expanded_count = residual_count * max(1, int(residual_scale_count))
+        if 0 <= selected_index < expanded_count and residual_count > 0:
+            return f"retrieval_residual_{case.candidate_types[selected_index % residual_count]}"
         return "retrieval_residual_gaussian"
     if prepended_policy_candidate and selection_mode in {"lattice", "retrieval_lattice"}:
         if selected_index == 0: