Auto-sync: 2026-06-30 20:06:30

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_type_margins: dict[str, float] | None = None,
     retrieval_residual_challenger_tasks: tuple[str, ...] = (),
     retrieval_residual_challenger_type_tasks: dict[str, tuple[str, ...]] | None = None,
     lattice_exclude_types: tuple[str, ...] = (),
@@ -272,6 +273,15 @@ def evaluate_maniskill_policy_rollout(
         ).items()
         if candidate_type.strip()
     }
+    retrieval_residual_challenger_type_margins = {
+        _normalize_residual_candidate_type_name(candidate_type): float(margin)
+        for candidate_type, margin in (
+            retrieval_residual_challenger_type_margins or {}
+        ).items()
+        if candidate_type.strip()
+    }
+    if any(margin < 0 for margin in retrieval_residual_challenger_type_margins.values()):
+        raise ValueError("retrieval_residual_challenger_type_margins must be non-negative")
     retrieval_residual_challenger_scales = tuple(
         float(scale) for scale in retrieval_residual_challenger_scales
     )
@@ -428,6 +438,9 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_residual_challenger_margin=(
                     retrieval_residual_challenger_margin
                 ),
+                retrieval_residual_challenger_type_margins=(
+                    retrieval_residual_challenger_type_margins
+                ),
                 retrieval_residual_challenger_tasks=(
                     retrieval_residual_challenger_tasks
                 ),
@@ -576,6 +589,11 @@ def evaluate_maniskill_policy_rollout(
             if selection_mode == "retrieval_residual"
             else 0.0
         ),
+        "retrieval_residual_challenger_type_margins": (
+            retrieval_residual_challenger_type_margins
+            if selection_mode == "retrieval_residual"
+            else {}
+        ),
         "retrieval_residual_challenger_tasks": list(
             retrieval_residual_challenger_tasks
         )
@@ -1278,6 +1296,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_type_margins: dict[str, float] | None = None,
     retrieval_residual_challenger_tasks: tuple[str, ...] = (),
     retrieval_residual_challenger_type_tasks: dict[str, tuple[str, ...]] | None = None,
     lattice_exclude_types: tuple[str, ...] = (),
@@ -1360,6 +1379,21 @@ def _evaluate_task_cases(
                 retrieval_residual_reduce=retrieval_residual_reduce,
                 retrieval_residual_action_l2_penalty=retrieval_residual_action_l2_penalty,
                 retrieval_residual_challenger_scales=retrieval_residual_challenger_scales,
+                retrieval_residual_challenger_margin_by_candidate=(
+                    _lattice_candidate_type_margin(
+                        batch,
+                        torch=torch,
+                        device=device,
+                        default_margin=retrieval_residual_challenger_margin,
+                        candidate_type_margins=(
+                            retrieval_residual_challenger_type_margins or {}
+                        ),
+                    )
+                    if selection_mode == "retrieval_residual"
+                    and retrieval_residual_challenger_types
+                    and retrieval_residual_challenger_type_margins
+                    else None
+                ),
                 action_low=action_low,
                 action_high=action_high,
                 action_candidates=(
@@ -1914,6 +1948,7 @@ def _select_action_chunk(
     retrieval_residual_reduce: str = "none",
     retrieval_residual_action_l2_penalty: float = 0.0,
     retrieval_residual_challenger_scales: tuple[float, ...] = (),
+    retrieval_residual_challenger_margin_by_candidate: Any | None = None,
     action_low: Any | None = None,
     action_high: Any | None = None,
     action_candidates: Any | None = None,
@@ -2010,6 +2045,9 @@ def _select_action_chunk(
             retrieval_residual_challenger_scales=retrieval_residual_challenger_scales,
             challenger_mask=challenger_mask,
             challenger_margin=challenger_margin,
+            challenger_margin_by_candidate=(
+                retrieval_residual_challenger_margin_by_candidate
+            ),
         )
 
     if selection_mode == "policy" or num_candidates <= 1:
@@ -2256,6 +2294,7 @@ def _select_lattice_action_chunk(
     action_l2_penalty_base: Any | None = None,
     challenger_mask: Any | None = None,
     challenger_margin: float = 0.0,
+    challenger_margin_by_candidate: Any | None = None,
 ) -> tuple[Any, np.ndarray]:
     if action_candidates.ndim != 4:
         raise ValueError("action_candidates must have shape [B,K,H,D]")
@@ -2309,6 +2348,11 @@ def _select_lattice_action_chunk(
             raise ValueError("challenger_mask must have shape [B,K]")
         if challenger_margin < 0:
             raise ValueError("challenger_margin must be non-negative")
+        if challenger_margin_by_candidate is not None:
+            if challenger_margin_by_candidate.shape != potentials.shape:
+                raise ValueError("challenger_margin_by_candidate must have shape [B,K]")
+            if torch.any(challenger_margin_by_candidate < 0):
+                raise ValueError("challenger_margin_by_candidate must be non-negative")
         challenger_potentials = potentials.masked_fill(~challenger_mask, float("-inf"))
         challenger_index = torch.argmax(challenger_potentials, dim=1)
         challenger_best = challenger_potentials.gather(
@@ -2319,8 +2363,18 @@ def _select_lattice_action_chunk(
             1,
             best_index.reshape(batch_size, 1),
         ).reshape(batch_size)
+        if challenger_margin_by_candidate is None:
+            challenger_margin_threshold = torch.full_like(
+                challenger_best,
+                float(challenger_margin),
+            )
+        else:
+            challenger_margin_threshold = challenger_margin_by_candidate.to(
+                device=potentials.device,
+                dtype=potentials.dtype,
+            ).gather(1, challenger_index.reshape(batch_size, 1)).reshape(batch_size)
         use_challenger = torch.isfinite(challenger_best) & (
-            challenger_best > primary_potential + float(challenger_margin)
+            challenger_best > primary_potential + challenger_margin_threshold
         )
         best_index = torch.where(use_challenger, challenger_index, best_index)
     batch_index = torch.arange(batch_size, device=candidates.device)
@@ -2412,6 +2466,7 @@ def _select_residual_lattice_action_chunk(
     retrieval_residual_challenger_scales: tuple[float, ...] = (),
     challenger_mask: Any | None = None,
     challenger_margin: float = 0.0,
+    challenger_margin_by_candidate: Any | None = None,
 ) -> tuple[Any, np.ndarray]:
     if action_residuals.ndim != 4:
         raise ValueError("action_residuals must have shape [B,K,H,D]")
@@ -2462,6 +2517,14 @@ def _select_residual_lattice_action_chunk(
         num_gaussian_candidates=num_gaussian_candidates,
         candidate_sigma=candidate_sigma,
     )
+    challenger_margin_by_candidate = _expand_residual_lattice_values(
+        challenger_margin_by_candidate,
+        torch=torch,
+        scales=scales,
+        fill_value=float(challenger_margin),
+        num_gaussian_candidates=num_gaussian_candidates,
+        candidate_sigma=candidate_sigma,
+    )
     return _select_lattice_action_chunk(
         model,
         observations,
@@ -2485,6 +2548,7 @@ def _select_residual_lattice_action_chunk(
         action_l2_penalty_base=policy_mean,
         challenger_mask=challenger_mask,
         challenger_margin=challenger_margin,
+        challenger_margin_by_candidate=challenger_margin_by_candidate,
     )
 
 
@@ -2578,6 +2642,33 @@ def _expand_residual_lattice_mask(
     return expanded
 
 
+def _expand_residual_lattice_values(
+    values: Any | None,
+    *,
+    torch: Any,
+    scales: tuple[float, ...],
+    fill_value: float,
+    num_gaussian_candidates: int,
+    candidate_sigma: float,
+) -> Any | None:
+    if values is None:
+        return None
+    scale_values = tuple(float(scale) for scale in scales)
+    if not scale_values:
+        raise ValueError("scales must not be empty")
+    blocks = [values for _ 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.full(
+            (expanded.shape[0], num_gaussian_candidates - 1),
+            float(fill_value),
+            dtype=expanded.dtype,
+            device=expanded.device,
+        )
+        expanded = torch.cat([expanded, extra], dim=1)
+    return expanded
+
+
 def _scale_is_allowed(scale: float, allowed_scales: tuple[float, ...]) -> bool:
     if not allowed_scales:
         return True
@@ -2978,6 +3069,25 @@ def _lattice_candidate_type_bonus(
     return torch.tensor(rows, dtype=torch.float32, device=device)
 
 
+def _lattice_candidate_type_margin(
+    batch: list[_RolloutCase],
+    *,
+    torch: Any,
+    device: str,
+    default_margin: float,
+    candidate_type_margins: dict[str, float],
+) -> Any:
+    rows: list[list[float]] = []
+    for case in batch:
+        rows.append(
+            [
+                float(candidate_type_margins.get(candidate_type, default_margin))
+                for candidate_type in case.candidate_types
+            ]
+        )
+    return torch.tensor(rows, dtype=torch.float32, device=device)
+
+
 def _candidate_type_bonus_for_type(
     candidate_type: str,
     candidate_type_bonuses: dict[str, float],

logs/auto_sync_hf.log

@@ -278,3 +278,4 @@ No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
+No files have been modified since last commit. Skipping to prevent empty commit.