Auto-sync: 2026-06-28 01:28:06

dovla_cil/eval/maniskill_policy_rollout.py

@@ -60,6 +60,7 @@ def evaluate_maniskill_policy_rollout(
     field_optim_trust_radius: float = 0.5,
     field_optim_l2_penalty: float = 0.0,
     retrieval_neighbors: int = 1,
+    retrieval_residual_scale: float = 1.0,
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> dict[str, Any]:
     """Execute a checkpoint policy from restored ManiSkill CIL states.
@@ -134,6 +135,8 @@ def evaluate_maniskill_policy_rollout(
         raise ValueError("field_optim_l2_penalty must be non-negative")
     if retrieval_neighbors <= 0:
         raise ValueError("retrieval_neighbors must be positive")
+    if retrieval_residual_scale < 0:
+        raise ValueError("retrieval_residual_scale must be non-negative")
     if selection_mode == "policy":
         num_candidates = 1
     checkpoint = torch.load(
@@ -218,6 +221,7 @@ def evaluate_maniskill_policy_rollout(
                 field_optim_step_size=field_optim_step_size,
                 field_optim_trust_radius=field_optim_trust_radius,
                 field_optim_l2_penalty=field_optim_l2_penalty,
+                retrieval_residual_scale=retrieval_residual_scale,
                 lattice_exclude_types=lattice_exclude_types,
             )
             rows.extend(task_rows)
@@ -243,7 +247,8 @@ def evaluate_maniskill_policy_rollout(
         "selection_mode": selection_mode,
         "num_candidates": effective_num_candidates,
         "candidate_sigma": candidate_sigma
-        if selection_mode in {"field", "field_optim"}
+        if selection_mode in {"field", "field_optim", "retrieval_residual"}
+        and num_candidates > 1
         else 0.0,
         "field_optim_steps": field_optim_steps
         if selection_mode == "field_optim"
@@ -260,6 +265,9 @@ def evaluate_maniskill_policy_rollout(
         "retrieval_neighbors": retrieval_neighbors
         if selection_mode in {"retrieval_lattice", "retrieval_residual"}
         else 0,
+        "retrieval_residual_scale": retrieval_residual_scale
+        if selection_mode == "retrieval_residual"
+        else 0.0,
         "lattice_exclude_types": list(lattice_exclude_types),
         "policy_rollout_success_rate": _mean([row["success"] for row in rows]),
         "policy_rollout_progress": _mean([row["progress"] for row in rows]),
@@ -508,6 +516,7 @@ def _evaluate_task_cases(
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
     field_optim_l2_penalty: float = 0.0,
+    retrieval_residual_scale: float = 1.0,
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
@@ -570,6 +579,7 @@ def _evaluate_task_cases(
                 field_optim_step_size=field_optim_step_size,
                 field_optim_trust_radius=field_optim_trust_radius,
                 field_optim_l2_penalty=field_optim_l2_penalty,
+                retrieval_residual_scale=retrieval_residual_scale,
                 action_low=action_low,
                 action_high=action_high,
                 action_candidates=(
@@ -640,7 +650,12 @@ def _evaluate_task_cases(
                             selection_mode=selection_mode,
                         ),
                         "selected_candidate_index": int(candidate_index[index]),
-                        "lattice_candidate_count": len(case.candidate_action_values),
+                        "lattice_candidate_count": _effective_lattice_candidate_count(
+                            case,
+                            selection_mode=selection_mode,
+                            num_candidates=num_candidates,
+                            candidate_sigma=candidate_sigma,
+                        ),
                         "candidate_source_group_id": case.candidate_source_group_id,
                     }
                 )
@@ -663,6 +678,7 @@ def _select_action_chunk(
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
     field_optim_l2_penalty: float = 0.0,
+    retrieval_residual_scale: float = 1.0,
     action_low: Any | None = None,
     action_high: Any | None = None,
     action_candidates: Any | None = None,
@@ -708,6 +724,10 @@ def _select_action_chunk(
             action_low=action_low,
             action_high=action_high,
             candidate_mask=candidate_mask,
+            residual_scale=retrieval_residual_scale,
+            num_gaussian_candidates=num_candidates,
+            candidate_sigma=candidate_sigma,
+            selection_seed=selection_seed,
         )
 
     if selection_mode == "policy" or num_candidates <= 1:
@@ -988,13 +1008,41 @@ def _select_residual_lattice_action_chunk(
     action_low: Any | None,
     action_high: Any | None,
     candidate_mask: Any | None,
+    residual_scale: float,
+    num_gaussian_candidates: int,
+    candidate_sigma: float,
+    selection_seed: int,
 ) -> 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) + action_residuals.to(
+    candidates = policy_mean.unsqueeze(1) + float(residual_scale) * action_residuals.to(
         device=policy_mean.device,
         dtype=policy_mean.dtype,
     )
+    if num_gaussian_candidates > 1 and candidate_sigma > 0:
+        generator = torch.Generator(device=policy_mean.device)
+        generator.manual_seed(int(selection_seed))
+        noise = torch.randn(
+            (
+                policy_mean.shape[0],
+                num_gaussian_candidates - 1,
+                policy_mean.shape[1],
+                policy_mean.shape[2],
+            ),
+            generator=generator,
+            device=policy_mean.device,
+            dtype=policy_mean.dtype,
+        )
+        gaussian = policy_mean.unsqueeze(1) + float(candidate_sigma) * noise
+        candidates = torch.cat([candidates, gaussian], dim=1)
+        if candidate_mask is not None:
+            extra_mask = torch.ones(
+                candidate_mask.shape[0],
+                num_gaussian_candidates - 1,
+                dtype=candidate_mask.dtype,
+                device=candidate_mask.device,
+            )
+            candidate_mask = torch.cat([candidate_mask, extra_mask], dim=1)
     return _select_lattice_action_chunk(
         model,
         observations,
@@ -1015,6 +1063,19 @@ def _select_residual_lattice_action_chunk(
     )
 
 
+def _effective_lattice_candidate_count(
+    case: _RolloutCase,
+    *,
+    selection_mode: str,
+    num_candidates: int,
+    candidate_sigma: float,
+) -> int:
+    count = len(case.candidate_action_values)
+    if selection_mode == "retrieval_residual" and num_candidates > 1 and candidate_sigma > 0:
+        count += num_candidates - 1
+    return count
+
+
 def _lattice_candidate_mask(
     batch: list[_RolloutCase],
     *,
@@ -1047,7 +1108,7 @@ def _selected_candidate_type(
     if selection_mode == "retrieval_residual":
         if 0 <= selected_index < len(case.candidate_types):
             return f"retrieval_residual_{case.candidate_types[selected_index]}"
-        return "retrieval_residual_unknown"
+        return "retrieval_residual_gaussian"
     if 0 <= selected_index < len(case.candidate_types):
         return f"lattice_{case.candidate_types[selected_index]}"
     return "lattice_unknown"