Auto-sync: 2026-06-28 12:31:29

dovla_cil/eval/maniskill_policy_rollout.py

@@ -56,6 +56,7 @@ def evaluate_maniskill_policy_rollout(
     candidate_sigma: float = 0.2,
     selection_seed: int = 0,
     selection_margin: float = 0.0,
+    prepend_policy_candidate: bool = False,
     field_optim_steps: int = 0,
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
@@ -65,6 +66,7 @@ def evaluate_maniskill_policy_rollout(
     retrieval_type_min_success: float = 0.0,
     retrieval_residual_scale: float = 1.0,
     retrieval_residual_anchor: str = "expert",
+    retrieval_residual_reduce: str = "none",
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> dict[str, Any]:
     """Execute a checkpoint policy from restored ManiSkill CIL states.
@@ -147,6 +149,10 @@ def evaluate_maniskill_policy_rollout(
         raise ValueError("retrieval_metric must be 'raw' or 'zscore'")
     if retrieval_residual_anchor not in {"expert", "policy"}:
         raise ValueError("retrieval_residual_anchor must be 'expert' or 'policy'")
+    if retrieval_residual_reduce not in {"none", "mean_by_type", "median_by_type"}:
+        raise ValueError(
+            "retrieval_residual_reduce must be 'none', 'mean_by_type', or 'median_by_type'"
+        )
     if not 0.0 <= retrieval_type_min_success <= 1.0:
         raise ValueError("retrieval_type_min_success must be in [0, 1]")
     if retrieval_residual_scale < 0:
@@ -212,6 +218,7 @@ def evaluate_maniskill_policy_rollout(
                 retrieval_metric=retrieval_metric,
                 retrieval_type_min_success=retrieval_type_min_success,
                 retrieval_residual_anchor=retrieval_residual_anchor,
+                retrieval_residual_reduce=retrieval_residual_reduce,
             )
     by_task: dict[str, list[_RolloutCase]] = defaultdict(list)
     for case in cases:
@@ -237,6 +244,7 @@ def evaluate_maniskill_policy_rollout(
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed,
                 selection_margin=selection_margin,
+                prepend_policy_candidate=prepend_policy_candidate,
                 field_optim_steps=field_optim_steps,
                 field_optim_step_size=field_optim_step_size,
                 field_optim_trust_radius=field_optim_trust_radius,
@@ -271,6 +279,7 @@ def evaluate_maniskill_policy_rollout(
         and num_candidates > 1
         else 0.0,
         "selection_margin": selection_margin,
+        "prepend_policy_candidate": bool(prepend_policy_candidate),
         "field_optim_steps": field_optim_steps
         if selection_mode == "field_optim"
         else 0,
@@ -298,6 +307,9 @@ def evaluate_maniskill_policy_rollout(
         "retrieval_residual_anchor": retrieval_residual_anchor
         if selection_mode == "retrieval_residual"
         else "none",
+        "retrieval_residual_reduce": retrieval_residual_reduce
+        if selection_mode == "retrieval_residual"
+        else "none",
         "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]),
@@ -460,6 +472,7 @@ def _attach_retrieved_residual_candidates(
     retrieval_metric: str = "raw",
     retrieval_type_min_success: float = 0.0,
     retrieval_residual_anchor: str = "expert",
+    retrieval_residual_reduce: str = "none",
 ) -> list[_RolloutCase]:
     if observation_mode != "state":
         raise ValueError("retrieval_residual currently supports state observations only")
@@ -532,6 +545,12 @@ def _attach_retrieved_residual_candidates(
             source_group_ids.append(source_group_id)
             residuals.extend(source_residuals)
             candidate_types.extend(source_candidate_types)
+        if retrieval_residual_reduce != "none":
+            residuals, candidate_types = _reduce_residual_candidates_by_type(
+                residuals,
+                candidate_types,
+                mode=retrieval_residual_reduce,
+            )
         output.append(
             replace(
                 case,
@@ -543,6 +562,38 @@ def _attach_retrieved_residual_candidates(
     return output
 
 
+def _reduce_residual_candidates_by_type(
+    residuals: list[list[list[float]]],
+    candidate_types: list[str],
+    *,
+    mode: str,
+) -> tuple[list[list[list[float]]], list[str]]:
+    if mode not in {"mean_by_type", "median_by_type"}:
+        raise ValueError("mode must be 'mean_by_type' or 'median_by_type'")
+    if len(residuals) != len(candidate_types):
+        raise ValueError("residuals and candidate_types must have the same length")
+
+    ordered_types = list(dict.fromkeys(candidate_types))
+    reduced_residuals: list[list[list[float]]] = []
+    reduced_types: list[str] = []
+    for candidate_type in ordered_types:
+        values = [
+            np.asarray(residual, dtype=np.float32)
+            for residual, residual_type in zip(residuals, candidate_types)
+            if residual_type == candidate_type
+        ]
+        if not values:
+            continue
+        stack = np.stack(values, axis=0)
+        if mode == "mean_by_type":
+            reduced = np.mean(stack, axis=0)
+        else:
+            reduced = np.median(stack, axis=0)
+        reduced_residuals.append(reduced.astype(np.float32).tolist())
+        reduced_types.append(candidate_type)
+    return reduced_residuals, reduced_types
+
+
 def _candidate_type_success_rates(
     dataset: CILDataset,
     *,
@@ -605,6 +656,7 @@ def _evaluate_task_cases(
     candidate_sigma: float = 0.2,
     selection_seed: int = 0,
     selection_margin: float = 0.0,
+    prepend_policy_candidate: bool = False,
     field_optim_steps: int = 0,
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
@@ -669,6 +721,7 @@ def _evaluate_task_cases(
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed + start,
                 selection_margin=selection_margin,
+                prepend_policy_candidate=prepend_policy_candidate,
                 field_optim_steps=field_optim_steps,
                 field_optim_step_size=field_optim_step_size,
                 field_optim_trust_radius=field_optim_trust_radius,
@@ -742,6 +795,7 @@ def _evaluate_task_cases(
                             case,
                             selected_index=int(candidate_index[index]),
                             selection_mode=selection_mode,
+                            prepended_policy_candidate=prepend_policy_candidate,
                         ),
                         "selected_candidate_index": int(candidate_index[index]),
                         "lattice_candidate_count": _effective_lattice_candidate_count(
@@ -749,6 +803,7 @@ def _evaluate_task_cases(
                             selection_mode=selection_mode,
                             num_candidates=num_candidates,
                             candidate_sigma=candidate_sigma,
+                            prepended_policy_candidate=prepend_policy_candidate,
                         ),
                         "candidate_source_group_id": case.candidate_source_group_id,
                     }
@@ -769,6 +824,7 @@ def _select_action_chunk(
     candidate_sigma: float,
     selection_seed: int,
     selection_margin: float = 0.0,
+    prepend_policy_candidate: bool = False,
     field_optim_steps: int = 0,
     field_optim_step_size: float = 0.05,
     field_optim_trust_radius: float = 0.5,
@@ -793,6 +849,9 @@ def _select_action_chunk(
     if selection_mode in {"lattice", "retrieval_lattice"}:
         if action_candidates is None:
             raise ValueError(f"{selection_mode} selection requires action_candidates")
+        policy_baseline = None
+        if prepend_policy_candidate:
+            policy_baseline = model.forward_policy(observations, instructions)
         return _select_lattice_action_chunk(
             model,
             observations,
@@ -803,6 +862,7 @@ def _select_action_chunk(
             action_high=action_high,
             candidate_mask=candidate_mask,
             selection_margin=selection_margin,
+            baseline_action=policy_baseline,
         )
 
     policy_mean = model.forward_policy(observations, instructions)
@@ -1065,9 +1125,25 @@ def _select_lattice_action_chunk(
     action_high: Any | None,
     candidate_mask: Any | None,
     selection_margin: float = 0.0,
+    baseline_action: Any | None = None,
 ) -> tuple[Any, np.ndarray]:
     if action_candidates.ndim != 4:
         raise ValueError("action_candidates must have shape [B,K,H,D]")
+    if baseline_action is not None:
+        if baseline_action.ndim != 3:
+            raise ValueError("baseline_action must have shape [B,H,D]")
+        action_candidates = torch.cat(
+            [baseline_action.unsqueeze(1), action_candidates],
+            dim=1,
+        )
+        if candidate_mask is not None:
+            baseline_mask = torch.ones(
+                candidate_mask.shape[0],
+                1,
+                dtype=candidate_mask.dtype,
+                device=candidate_mask.device,
+            )
+            candidate_mask = torch.cat([baseline_mask, candidate_mask], dim=1)
     batch_size, candidate_count = action_candidates.shape[:2]
     candidates = _clamp_action_tensor(
         action_candidates,
@@ -1176,8 +1252,11 @@ def _effective_lattice_candidate_count(
     selection_mode: str,
     num_candidates: int,
     candidate_sigma: float,
+    prepended_policy_candidate: bool = False,
 ) -> int:
     count = len(case.candidate_action_values)
+    if prepended_policy_candidate and selection_mode in {"lattice", "retrieval_lattice"}:
+        count += 1
     if selection_mode == "retrieval_residual" and num_candidates > 1 and candidate_sigma > 0:
         count += num_candidates - 1
     return count
@@ -1205,6 +1284,7 @@ def _selected_candidate_type(
     *,
     selected_index: int,
     selection_mode: str,
+    prepended_policy_candidate: bool = False,
 ) -> str:
     if selection_mode == "policy":
         return "policy_continuous"
@@ -1216,6 +1296,10 @@ def _selected_candidate_type(
         if 0 <= selected_index < len(case.candidate_types):
             return f"retrieval_residual_{case.candidate_types[selected_index]}"
         return "retrieval_residual_gaussian"
+    if prepended_policy_candidate and selection_mode in {"lattice", "retrieval_lattice"}:
+        if selected_index == 0:
+            return "policy_continuous"
+        selected_index -= 1
     if 0 <= selected_index < len(case.candidate_types):
         return f"lattice_{case.candidate_types[selected_index]}"
     return "lattice_unknown"

logs/auto_sync_hf.log

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