auto-sync 2026-07-02T15:44:03Z workspace

workspace/dovla_cil/eval/maniskill_policy_rollout.py

@@ -138,6 +138,7 @@ def evaluate_maniskill_policy_rollout(
     candidate_type_bonuses_by_task: dict[str, dict[str, float]] | None = None,
     candidate_type_bonus_components: bool = False,
     field_rank_biases_by_task: dict[str, list[float]] | None = None,
+    residual_scale_bonuses_by_task: dict[str, dict[str, float]] | None = None,
     candidate_oracle_rollouts: int = 0,
     candidate_oracle_unique_tolerance: float = 1.0e-6,
 ) -> dict[str, Any]:
@@ -338,6 +339,10 @@ def evaluate_maniskill_policy_rollout(
         str(task_id): [float(value) for value in values]
         for task_id, values in (field_rank_biases_by_task or {}).items()
     }
+    residual_scale_bonuses_by_task = {
+        str(task_id): {str(scale): float(bonus) for scale, bonus in bonuses.items()}
+        for task_id, bonuses in (residual_scale_bonuses_by_task or {}).items()
+    }
     if selection_mode == "policy":
         num_candidates = 1
     lattice_exclude_types = tuple(
@@ -461,6 +466,10 @@ def evaluate_maniskill_policy_rollout(
                 candidate_type_bonuses_by_task,
                 task_id=task_id,
             )
+            task_residual_scale_bonuses = _task_residual_scale_bonuses(
+                residual_scale_bonuses_by_task,
+                task_id=task_id,
+            )
             task_rows = _evaluate_task_cases(
                 task_id,
                 task_cases,
@@ -509,6 +518,7 @@ def evaluate_maniskill_policy_rollout(
                 lattice_exclude_type_tasks=lattice_exclude_type_tasks,
                 candidate_type_bonuses=task_candidate_type_bonuses,
                 candidate_type_bonus_components=candidate_type_bonus_components,
+                residual_scale_bonuses=task_residual_scale_bonuses,
                 field_rank_biases=(
                     field_rank_biases_by_task.get(task_id)
                     or field_rank_biases_by_task.get("*")
@@ -681,6 +691,7 @@ def evaluate_maniskill_policy_rollout(
         "candidate_type_bonuses_by_task": candidate_type_bonuses_by_task,
         "candidate_type_bonus_components": bool(candidate_type_bonus_components),
         "field_rank_biases_by_task": field_rank_biases_by_task,
+        "residual_scale_bonuses_by_task": residual_scale_bonuses_by_task,
         "candidate_oracle_rollouts": int(candidate_oracle_rollouts),
         "candidate_oracle_unique_tolerance": float(candidate_oracle_unique_tolerance),
         "policy_rollout_success_rate": _mean([row["success"] for row in rows]),
@@ -1403,6 +1414,7 @@ def _evaluate_task_cases(
     lattice_exclude_type_tasks: dict[str, tuple[str, ...]] | None = None,
     candidate_type_bonuses: dict[str, float] | None = None,
     candidate_type_bonus_components: bool = False,
+    residual_scale_bonuses: dict[str, float] | None = None,
     field_rank_biases: list[float] | None = None,
     candidate_oracle_rollouts: int = 0,
     candidate_oracle_unique_tolerance: float = 1.0e-6,
@@ -1534,6 +1546,7 @@ def _evaluate_task_cases(
                     )
                     else None
                 ),
+                residual_scale_bonuses=residual_scale_bonuses or {},
                 field_rank_bias=(
                     torch.tensor(field_rank_biases, dtype=torch.float32, device=device)
                     if field_rank_biases
@@ -1613,6 +1626,7 @@ def _evaluate_task_cases(
                     lattice_exclude_types=task_lattice_exclude_types,
                     candidate_type_bonuses=candidate_type_bonuses or {},
                     candidate_type_bonus_components=candidate_type_bonus_components,
+                    residual_scale_bonuses=residual_scale_bonuses or {},
                     candidate_oracle_unique_tolerance=candidate_oracle_unique_tolerance,
                 )
             for index, case in enumerate(batch):
@@ -2060,6 +2074,7 @@ def _select_action_chunk(
     action_candidates: Any | None = None,
     candidate_mask: Any | None = None,
     candidate_type_bonus: Any | None = None,
+    residual_scale_bonuses: dict[str, float] | None = None,
     field_rank_bias: Any | None = None,
     residual_aggregate_mask: Any | None = None,
     challenger_mask: Any | None = None,
@@ -2609,6 +2624,8 @@ def _select_residual_lattice_action_chunk(
     )
     scales = tuple(float(scale) for scale in residual_scales) or (float(residual_scale),)
     if residual_reduce == "field_softmax":
+        if residual_scale_bonuses:
+            raise ValueError("residual_scale_bonuses are not supported with field_softmax")
         if challenger_mask is not None:
             raise ValueError("challenger_mask is not supported with field_softmax")
         return _select_field_softmax_residual_action_chunk(
@@ -2639,6 +2656,7 @@ def _select_residual_lattice_action_chunk(
         residual_scales=residual_scales,
         candidate_mask=candidate_mask,
         candidate_type_bonus=candidate_type_bonus,
+        residual_scale_bonuses=residual_scale_bonuses or {},
         num_gaussian_candidates=num_gaussian_candidates,
         candidate_sigma=candidate_sigma,
         selection_seed=selection_seed,
@@ -2696,6 +2714,7 @@ def _build_residual_lattice_candidates(
     residual_scales: tuple[float, ...],
     candidate_mask: Any | None,
     candidate_type_bonus: Any | None,
+    residual_scale_bonuses: dict[str, float],
     num_gaussian_candidates: int,
     candidate_sigma: float,
     selection_seed: int,
@@ -2711,6 +2730,29 @@ def _build_residual_lattice_candidates(
         candidate_mask = torch.cat([candidate_mask for _ in scales], dim=1)
     if candidate_type_bonus is not None and len(scales) > 1:
         candidate_type_bonus = torch.cat([candidate_type_bonus for _ in scales], dim=1)
+    if residual_scale_bonuses:
+        residual_count = int(residuals.shape[1])
+        scale_blocks = []
+        for scale in scales:
+            bonus = float(
+                residual_scale_bonuses.get(str(scale), residual_scale_bonuses.get(f"{scale:g}", 0.0))
+            )
+            scale_blocks.append(
+                torch.full(
+                    (policy_mean.shape[0], residual_count),
+                    bonus,
+                    dtype=policy_mean.dtype,
+                    device=policy_mean.device,
+                )
+            )
+        scale_bonus = torch.cat(scale_blocks, dim=1)
+        if candidate_type_bonus is None:
+            candidate_type_bonus = scale_bonus
+        else:
+            candidate_type_bonus = candidate_type_bonus + scale_bonus.to(
+                device=candidate_type_bonus.device,
+                dtype=candidate_type_bonus.dtype,
+            )
     if num_gaussian_candidates > 1 and candidate_sigma > 0:
         generator = torch.Generator(device=policy_mean.device)
         generator.manual_seed(int(selection_seed))
@@ -3193,6 +3235,16 @@ def _task_candidate_type_bonuses(
     return bonuses
 
 
+def _task_residual_scale_bonuses(
+    bonuses_by_task: dict[str, dict[str, float]],
+    *,
+    task_id: str,
+) -> dict[str, float]:
+    bonuses = dict(bonuses_by_task.get("*", {}))
+    bonuses.update(bonuses_by_task.get(task_id, {}))
+    return bonuses
+
+
 def _lattice_candidate_type_bonus(
     batch: list[_RolloutCase],
     *,