Auto-sync: 2026-06-28 20:44:54

dovla_cil/eval/maniskill_policy_rollout.py

@@ -99,9 +99,10 @@ def evaluate_maniskill_policy_rollout(
 
     When ``selection_mode == 'retrieval_residual'`` the evaluator retrieves counterfactual
     action residuals (candidate minus expert action) from the nearest training-split state(s),
-    adds those residuals around the current policy mean, scores the resulting local proposal
-    lattice with the learned field, and executes the best chunk. This keeps the proposal
-    geometry counterfactual while avoiding same-state validation candidates.
+    optionally reduces them into type-wise tangent consensus proposals, adds those residuals
+    around the current policy mean, scores the resulting local proposal lattice with the learned
+    field, and executes the best chunk. This keeps the proposal geometry counterfactual while
+    avoiding same-state validation candidates.
 
     When ``selection_mode == 'field_optim'`` the evaluator starts from the policy mean plus
     optional Gaussian multi-start proposals, performs projected gradient ascent on the learned
@@ -155,9 +156,15 @@ 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"}:
+    if retrieval_residual_reduce not in {
+        "none",
+        "mean_by_type",
+        "median_by_type",
+        "kernel_mean_by_type",
+    }:
         raise ValueError(
-            "retrieval_residual_reduce must be 'none', 'mean_by_type', or 'median_by_type'"
+            "retrieval_residual_reduce must be 'none', 'mean_by_type', "
+            "'median_by_type', or 'kernel_mean_by_type'"
         )
     if not 0.0 <= retrieval_type_min_success <= 1.0:
         raise ValueError("retrieval_type_min_success must be in [0, 1]")
@@ -217,7 +224,6 @@ def evaluate_maniskill_policy_rollout(
                 observation_mode=model_config.observation_mode,
                 retrieval_neighbors=retrieval_neighbors,
                 retrieval_metric=retrieval_metric,
-                retrieval_type_min_success=retrieval_type_min_success,
             )
         else:
             cases = _attach_retrieved_residual_candidates(
@@ -557,7 +563,7 @@ def _attach_retrieved_residual_candidates(
             vectorize_toy_observation(case.observation, obs_dim=obs_dim),
             dtype=np.float32,
         )
-        nearest = _nearest_retrieval_entries(
+        nearest = _nearest_retrieval_entries_with_distances(
             candidates,
             query,
             retrieval_neighbors=retrieval_neighbors,
@@ -566,15 +572,22 @@ def _attach_retrieved_residual_candidates(
         source_group_ids: list[str] = []
         residuals: list[list[list[float]]] = []
         candidate_types: list[str] = []
-        for source_group_id, _feature, source_residuals, source_candidate_types in nearest:
+        residual_weights: list[float] = []
+        source_weights = _kernel_weights_from_distances(
+            [distance for _entry, distance in nearest]
+        )
+        for (entry, _distance), source_weight in zip(nearest, source_weights):
+            source_group_id, _feature, source_residuals, source_candidate_types = entry
             source_group_ids.append(source_group_id)
             residuals.extend(source_residuals)
             candidate_types.extend(source_candidate_types)
+            residual_weights.extend([float(source_weight)] * len(source_residuals))
         if retrieval_residual_reduce != "none":
             residuals, candidate_types = _reduce_residual_candidates_by_type(
                 residuals,
                 candidate_types,
                 mode=retrieval_residual_reduce,
+                weights=residual_weights,
             )
         output.append(
             replace(
@@ -592,26 +605,48 @@ def _reduce_residual_candidates_by_type(
     candidate_types: list[str],
     *,
     mode: str,
+    weights: list[float] | None = None,
 ) -> 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 mode not in {"mean_by_type", "median_by_type", "kernel_mean_by_type"}:
+        raise ValueError(
+            "mode must be 'mean_by_type', 'median_by_type', or 'kernel_mean_by_type'"
+        )
     if len(residuals) != len(candidate_types):
         raise ValueError("residuals and candidate_types must have the same length")
+    if weights is not None and len(weights) != len(residuals):
+        raise ValueError("weights and residuals 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
-        ]
+        values: list[np.ndarray] = []
+        value_weights: list[float] = []
+        for index, (residual, residual_type) in enumerate(zip(residuals, candidate_types)):
+            if residual_type != candidate_type:
+                continue
+            values.append(np.asarray(residual, dtype=np.float32))
+            if weights is not None:
+                value_weights.append(float(weights[index]))
         if not values:
             continue
         stack = np.stack(values, axis=0)
         if mode == "mean_by_type":
             reduced = np.mean(stack, axis=0)
+        elif mode == "kernel_mean_by_type":
+            if value_weights:
+                np_weights = np.asarray(value_weights, dtype=np.float32)
+                weight_sum = float(np.sum(np_weights))
+                if weight_sum > 1e-12:
+                    np_weights = np_weights / weight_sum
+                    reduced = np.sum(
+                        stack * np_weights.reshape((-1,) + (1,) * (stack.ndim - 1)),
+                        axis=0,
+                    )
+                else:
+                    reduced = np.mean(stack, axis=0)
+            else:
+                reduced = np.mean(stack, axis=0)
         else:
             reduced = np.median(stack, axis=0)
         reduced_residuals.append(reduced.astype(np.float32).tolist())
@@ -647,11 +682,30 @@ def _nearest_retrieval_entries(
     retrieval_neighbors: int,
     retrieval_metric: str,
 ) -> list[tuple[Any, np.ndarray, Any, Any]]:
-    if retrieval_metric == "raw":
-        return sorted(
+    return [
+        entry
+        for entry, _distance in _nearest_retrieval_entries_with_distances(
             candidates,
-            key=lambda item: float(np.mean((item[1] - query) ** 2)),
-        )[:retrieval_neighbors]
+            query,
+            retrieval_neighbors=retrieval_neighbors,
+            retrieval_metric=retrieval_metric,
+        )
+    ]
+
+
+def _nearest_retrieval_entries_with_distances(
+    candidates: list[tuple[Any, np.ndarray, Any, Any]],
+    query: np.ndarray,
+    *,
+    retrieval_neighbors: int,
+    retrieval_metric: str,
+) -> list[tuple[tuple[Any, np.ndarray, Any, Any], float]]:
+    if retrieval_metric == "raw":
+        scored = [
+            (item, float(np.mean((item[1] - query) ** 2)))
+            for item in candidates
+        ]
+        return sorted(scored, key=lambda item: item[1])[:retrieval_neighbors]
     if retrieval_metric != "zscore":
         raise ValueError("retrieval_metric must be 'raw' or 'zscore'")
     features = np.stack([np.asarray(item[1], dtype=np.float32) for item in candidates], axis=0)
@@ -660,8 +714,35 @@ def _nearest_retrieval_entries(
     scale = np.where(std > 1e-6, std, 1.0).astype(np.float32)
     normalized_features = (features - mean) / scale
     normalized_query = (np.asarray(query, dtype=np.float32) - mean) / scale
-    order = np.argsort(np.mean((normalized_features - normalized_query) ** 2, axis=1))
-    return [candidates[int(index)] for index in order[:retrieval_neighbors]]
+    distances = np.mean((normalized_features - normalized_query) ** 2, axis=1)
+    order = np.argsort(distances)
+    return [
+        (candidates[int(index)], float(distances[int(index)]))
+        for index in order[:retrieval_neighbors]
+    ]
+
+
+def _kernel_weights_from_distances(distances: list[float]) -> list[float]:
+    if not distances:
+        return []
+    values = np.asarray(distances, dtype=np.float32)
+    values = np.maximum(np.nan_to_num(values, nan=np.inf, posinf=np.inf, neginf=0.0), 0.0)
+    finite = values[np.isfinite(values)]
+    if finite.size == 0:
+        return [1.0] * len(distances)
+    positive = finite[finite > 1e-12]
+    if positive.size == 0:
+        return [1.0] * len(distances)
+    bandwidth = float(np.median(positive))
+    if bandwidth <= 1e-12:
+        bandwidth = float(np.max(positive))
+    if bandwidth <= 1e-12:
+        return [1.0] * len(distances)
+    weights = np.exp(-np.minimum(values, bandwidth * 50.0) / bandwidth)
+    weights = np.where(np.isfinite(weights), weights, 0.0)
+    if float(np.sum(weights)) <= 1e-12:
+        return [1.0] * len(distances)
+    return weights.astype(np.float32).tolist()
 
 
 def _evaluate_task_cases(