Auto-sync: 2026-06-28 17:30:45

dovla_cil/eval/maniskill_policy_rollout.py

@@ -69,6 +69,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_residual_anchor: str = "expert",
     retrieval_residual_reduce: str = "none",
     lattice_exclude_types: tuple[str, ...] = (),
+    candidate_type_bonuses: dict[str, float] | None = None,
 ) -> dict[str, Any]:
     """Execute a checkpoint policy from restored ManiSkill CIL states.
 
@@ -106,6 +107,10 @@ def evaluate_maniskill_policy_rollout(
     optional Gaussian multi-start proposals, performs projected gradient ascent on the learned
     field potential in action space, and executes the best optimized chunk. This is still
     deployment-clean: no dataset action candidates or rewards are consulted.
+
+    ``candidate_type_bonuses`` adds small type-specific priors to field potentials before
+    selection. The default empty mapping preserves previous behavior; non-empty values test
+    typed sparse-intervention hypotheses without reading validation rewards.
     """
 
     try:
@@ -160,6 +165,10 @@ def evaluate_maniskill_policy_rollout(
         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")
+    candidate_type_bonuses = {
+        str(candidate_type): float(bonus)
+        for candidate_type, bonus in (candidate_type_bonuses or {}).items()
+    }
     if selection_mode == "policy":
         num_candidates = 1
     checkpoint = torch.load(
@@ -255,6 +264,7 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_residual_scale=retrieval_residual_scale,
                 retrieval_residual_scales=retrieval_residual_scales,
                 lattice_exclude_types=lattice_exclude_types,
+                candidate_type_bonuses=candidate_type_bonuses,
             )
             rows.extend(task_rows)
             task_summaries[task_id] = _summarize_rows(task_rows)
@@ -318,6 +328,7 @@ def evaluate_maniskill_policy_rollout(
         if selection_mode == "retrieval_residual"
         else "none",
         "lattice_exclude_types": list(lattice_exclude_types),
+        "candidate_type_bonuses": candidate_type_bonuses,
         "policy_rollout_success_rate": _mean([row["success"] for row in rows]),
         "policy_rollout_progress": _mean([row["progress"] for row in rows]),
         "oracle_success_rate": _mean([row["oracle_success"] for row in rows]),
@@ -678,6 +689,7 @@ def _evaluate_task_cases(
     retrieval_residual_scale: float = 1.0,
     retrieval_residual_scales: tuple[float, ...] = (),
     lattice_exclude_types: tuple[str, ...] = (),
+    candidate_type_bonuses: dict[str, float] | None = None,
 ) -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
     for start in range(0, len(cases), group_batch_size):
@@ -767,6 +779,17 @@ def _evaluate_task_cases(
                     and lattice_exclude_types
                     else None
                 ),
+                candidate_type_bonus=(
+                    _lattice_candidate_type_bonus(
+                        batch,
+                        torch=torch,
+                        device=device,
+                        candidate_type_bonuses=candidate_type_bonuses or {},
+                    )
+                    if selection_mode in {"lattice", "retrieval_lattice", "retrieval_residual"}
+                    and candidate_type_bonuses
+                    else None
+                ),
             )
             predicted_np = predicted.detach().cpu().numpy().astype(np.float32)
             adapted = _adapt_action_dim(predicted_np, env_dim)
@@ -864,6 +887,7 @@ def _select_action_chunk(
     action_high: Any | None = None,
     action_candidates: Any | None = None,
     candidate_mask: Any | None = None,
+    candidate_type_bonus: Any | None = None,
 ) -> tuple[Any, np.ndarray]:
     """Return the executed action chunk per state and the chosen candidate index.
 
@@ -891,6 +915,7 @@ def _select_action_chunk(
             action_low=action_low,
             action_high=action_high,
             candidate_mask=candidate_mask,
+            candidate_type_bonus=candidate_type_bonus,
             selection_margin=selection_margin,
             baseline_action=policy_baseline,
         )
@@ -910,6 +935,7 @@ def _select_action_chunk(
             action_low=action_low,
             action_high=action_high,
             candidate_mask=candidate_mask,
+            candidate_type_bonus=candidate_type_bonus,
             residual_scale=retrieval_residual_scale,
             residual_scales=retrieval_residual_scales,
             num_gaussian_candidates=num_candidates,
@@ -1155,6 +1181,7 @@ def _select_lattice_action_chunk(
     action_low: Any | None,
     action_high: Any | None,
     candidate_mask: Any | None,
+    candidate_type_bonus: Any | None = None,
     selection_margin: float = 0.0,
     baseline_action: Any | None = None,
 ) -> tuple[Any, np.ndarray]:
@@ -1175,6 +1202,14 @@ def _select_lattice_action_chunk(
                 device=candidate_mask.device,
             )
             candidate_mask = torch.cat([baseline_mask, candidate_mask], dim=1)
+        if candidate_type_bonus is not None:
+            baseline_bonus = torch.zeros(
+                candidate_type_bonus.shape[0],
+                1,
+                dtype=candidate_type_bonus.dtype,
+                device=candidate_type_bonus.device,
+            )
+            candidate_type_bonus = torch.cat([baseline_bonus, candidate_type_bonus], dim=1)
     batch_size, candidate_count = action_candidates.shape[:2]
     candidates = _clamp_action_tensor(
         action_candidates,
@@ -1194,6 +1229,13 @@ def _select_lattice_action_chunk(
     ]
     field = model.forward_field(flat_observations, flat_instructions, flat_candidates)
     potentials = field["potential"].reshape(batch_size, candidate_count)
+    if candidate_type_bonus is not None:
+        if candidate_type_bonus.shape != potentials.shape:
+            raise ValueError("candidate_type_bonus must have shape [B,K]")
+        potentials = potentials + candidate_type_bonus.to(
+            device=potentials.device,
+            dtype=potentials.dtype,
+        )
     if candidate_mask is not None:
         potentials = potentials.masked_fill(~candidate_mask, float("-inf"))
     best_index = torch.argmax(potentials, dim=1)
@@ -1220,6 +1262,7 @@ def _select_residual_lattice_action_chunk(
     action_low: Any | None,
     action_high: Any | None,
     candidate_mask: Any | None,
+    candidate_type_bonus: Any | None = None,
     residual_scale: float,
     num_gaussian_candidates: int,
     candidate_sigma: float,
@@ -1241,6 +1284,8 @@ def _select_residual_lattice_action_chunk(
     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 candidate_type_bonus is not None and len(scales) > 1:
+        candidate_type_bonus = torch.cat([candidate_type_bonus 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))
@@ -1265,6 +1310,14 @@ def _select_residual_lattice_action_chunk(
                 device=candidate_mask.device,
             )
             candidate_mask = torch.cat([candidate_mask, extra_mask], dim=1)
+        if candidate_type_bonus is not None:
+            extra_bonus = torch.zeros(
+                candidate_type_bonus.shape[0],
+                num_gaussian_candidates - 1,
+                dtype=candidate_type_bonus.dtype,
+                device=candidate_type_bonus.device,
+            )
+            candidate_type_bonus = torch.cat([candidate_type_bonus, extra_bonus], dim=1)
     return _select_lattice_action_chunk(
         model,
         observations,
@@ -1282,6 +1335,7 @@ def _select_residual_lattice_action_chunk(
             else action_high
         ),
         candidate_mask=candidate_mask,
+        candidate_type_bonus=candidate_type_bonus,
         selection_margin=selection_margin,
     )
 
@@ -1344,6 +1398,27 @@ def _lattice_candidate_mask(
     return torch.tensor(rows, dtype=torch.bool, device=device)
 
 
+def _lattice_candidate_type_bonus(
+    batch: list[_RolloutCase],
+    *,
+    torch: Any,
+    device: str,
+    candidate_type_bonuses: dict[str, float],
+) -> Any:
+    rows: list[list[float]] = []
+    for case in batch:
+        bonuses: list[float] = []
+        for candidate_type in case.candidate_types:
+            bonus = candidate_type_bonuses.get(candidate_type)
+            if bonus is None:
+                bonus = candidate_type_bonuses.get(f"retrieval_residual_{candidate_type}")
+            if bonus is None:
+                bonus = candidate_type_bonuses.get(f"lattice_{candidate_type}")
+            bonuses.append(float(bonus or 0.0))
+        rows.append(bonuses)
+    return torch.tensor(rows, dtype=torch.float32, device=device)
+
+
 def _selected_candidate_type(
     case: _RolloutCase,
     *,

logs/auto_sync_hf.log

@@ -208,3 +208,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.