Auto-sync: 2026-06-28 07:49:25

dovla_cil/eval/maniskill_policy_rollout.py

@@ -55,6 +55,7 @@ def evaluate_maniskill_policy_rollout(
     num_candidates: int = 1,
     candidate_sigma: float = 0.2,
     selection_seed: int = 0,
+    selection_margin: float = 0.0,
     field_optim_steps: int = 0,
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
@@ -138,6 +139,8 @@ def evaluate_maniskill_policy_rollout(
         raise ValueError("field_optim_trust_radius must be non-negative")
     if field_optim_l2_penalty < 0:
         raise ValueError("field_optim_l2_penalty must be non-negative")
+    if selection_margin < 0:
+        raise ValueError("selection_margin must be non-negative")
     if retrieval_neighbors <= 0:
         raise ValueError("retrieval_neighbors must be positive")
     if retrieval_metric not in {"raw", "zscore"}:
@@ -233,6 +236,7 @@ def evaluate_maniskill_policy_rollout(
                 num_candidates=num_candidates,
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed,
+                selection_margin=selection_margin,
                 field_optim_steps=field_optim_steps,
                 field_optim_step_size=field_optim_step_size,
                 field_optim_trust_radius=field_optim_trust_radius,
@@ -266,6 +270,7 @@ def evaluate_maniskill_policy_rollout(
         if selection_mode in {"field", "field_optim", "retrieval_residual"}
         and num_candidates > 1
         else 0.0,
+        "selection_margin": selection_margin,
         "field_optim_steps": field_optim_steps
         if selection_mode == "field_optim"
         else 0,
@@ -599,6 +604,7 @@ def _evaluate_task_cases(
     num_candidates: int = 1,
     candidate_sigma: float = 0.2,
     selection_seed: int = 0,
+    selection_margin: float = 0.0,
     field_optim_steps: int = 0,
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
@@ -662,6 +668,7 @@ def _evaluate_task_cases(
                 num_candidates=num_candidates,
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed + start,
+                selection_margin=selection_margin,
                 field_optim_steps=field_optim_steps,
                 field_optim_step_size=field_optim_step_size,
                 field_optim_trust_radius=field_optim_trust_radius,
@@ -761,6 +768,7 @@ def _select_action_chunk(
     num_candidates: int,
     candidate_sigma: float,
     selection_seed: int,
+    selection_margin: float = 0.0,
     field_optim_steps: int = 0,
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
@@ -794,6 +802,7 @@ def _select_action_chunk(
             action_low=action_low,
             action_high=action_high,
             candidate_mask=candidate_mask,
+            selection_margin=selection_margin,
         )
 
     policy_mean = model.forward_policy(observations, instructions)
@@ -815,6 +824,7 @@ def _select_action_chunk(
             num_gaussian_candidates=num_candidates,
             candidate_sigma=candidate_sigma,
             selection_seed=selection_seed,
+            selection_margin=selection_margin,
         )
 
     if selection_mode == "policy" or num_candidates <= 1:
@@ -1054,6 +1064,7 @@ def _select_lattice_action_chunk(
     action_low: Any | None,
     action_high: Any | None,
     candidate_mask: Any | None,
+    selection_margin: float = 0.0,
 ) -> tuple[Any, np.ndarray]:
     if action_candidates.ndim != 4:
         raise ValueError("action_candidates must have shape [B,K,H,D]")
@@ -1079,6 +1090,13 @@ def _select_lattice_action_chunk(
     if candidate_mask is not None:
         potentials = potentials.masked_fill(~candidate_mask, float("-inf"))
     best_index = torch.argmax(potentials, dim=1)
+    if selection_margin > 0.0 and candidate_count > 0:
+        baseline = potentials[:, 0]
+        best_potential = potentials.gather(1, best_index.reshape(batch_size, 1)).reshape(
+            batch_size
+        )
+        keep_baseline = best_potential <= baseline + float(selection_margin)
+        best_index = torch.where(keep_baseline, torch.zeros_like(best_index), best_index)
     batch_index = torch.arange(batch_size, device=candidates.device)
     best_actions = candidates[batch_index, best_index]
     return best_actions, best_index.detach().cpu().numpy()
@@ -1099,6 +1117,7 @@ def _select_residual_lattice_action_chunk(
     num_gaussian_candidates: int,
     candidate_sigma: float,
     selection_seed: int,
+    selection_margin: float = 0.0,
 ) -> tuple[Any, np.ndarray]:
     if action_residuals.ndim != 4:
         raise ValueError("action_residuals must have shape [B,K,H,D]")
@@ -1147,6 +1166,7 @@ def _select_residual_lattice_action_chunk(
             else action_high
         ),
         candidate_mask=candidate_mask,
+        selection_margin=selection_margin,
     )