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

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_metric: str = "raw",
     retrieval_residual_scale: float = 1.0,
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> dict[str, Any]:
@@ -85,7 +86,9 @@ def evaluate_maniskill_policy_rollout(
     When ``selection_mode == 'retrieval_lattice'`` action proposals come from the nearest
     training-split state with the same task rather than the evaluated state's own lattice. This
     tests whether the field can use reusable intervention proposals without same-state proposal
-    leakage.
+    leakage. ``retrieval_metric='zscore'`` standardizes state features by the train-bank
+    statistics for each task before nearest-neighbor lookup; the default ``raw`` metric
+    preserves earlier results exactly.
 
     When ``selection_mode == 'retrieval_residual'`` the evaluator retrieves counterfactual
     action residuals (candidate minus expert action) from the nearest training-split state(s),
@@ -135,6 +138,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_metric not in {"raw", "zscore"}:
+        raise ValueError("retrieval_metric must be 'raw' or 'zscore'")
     if retrieval_residual_scale < 0:
         raise ValueError("retrieval_residual_scale must be non-negative")
     if selection_mode == "policy":
@@ -184,6 +189,7 @@ def evaluate_maniskill_policy_rollout(
                 obs_dim=model_config.obs_dim,
                 observation_mode=model_config.observation_mode,
                 retrieval_neighbors=retrieval_neighbors,
+                retrieval_metric=retrieval_metric,
             )
         else:
             cases = _attach_retrieved_residual_candidates(
@@ -193,6 +199,7 @@ def evaluate_maniskill_policy_rollout(
                 obs_dim=model_config.obs_dim,
                 observation_mode=model_config.observation_mode,
                 retrieval_neighbors=retrieval_neighbors,
+                retrieval_metric=retrieval_metric,
             )
     by_task: dict[str, list[_RolloutCase]] = defaultdict(list)
     for case in cases:
@@ -265,6 +272,9 @@ def evaluate_maniskill_policy_rollout(
         "retrieval_neighbors": retrieval_neighbors
         if selection_mode in {"retrieval_lattice", "retrieval_residual"}
         else 0,
+        "retrieval_metric": retrieval_metric
+        if selection_mode in {"retrieval_lattice", "retrieval_residual"}
+        else "none",
         "retrieval_residual_scale": retrieval_residual_scale
         if selection_mode == "retrieval_residual"
         else 0.0,
@@ -355,6 +365,7 @@ def _attach_retrieved_lattice_candidates(
     obs_dim: int,
     observation_mode: str,
     retrieval_neighbors: int,
+    retrieval_metric: str = "raw",
 ) -> list[_RolloutCase]:
     if observation_mode != "state":
         raise ValueError("retrieval_lattice currently supports state observations only")
@@ -394,10 +405,12 @@ def _attach_retrieved_lattice_candidates(
             vectorize_toy_observation(case.observation, obs_dim=obs_dim),
             dtype=np.float32,
         )
-        nearest = sorted(
+        nearest = _nearest_retrieval_entries(
             candidates,
-            key=lambda item: float(np.mean((item[1] - query) ** 2)),
-        )[:retrieval_neighbors]
+            query,
+            retrieval_neighbors=retrieval_neighbors,
+            retrieval_metric=retrieval_metric,
+        )
         source_group_ids: list[str] = []
         actions: list[list[list[float]]] = []
         candidate_types: list[str] = []
@@ -424,6 +437,7 @@ def _attach_retrieved_residual_candidates(
     obs_dim: int,
     observation_mode: str,
     retrieval_neighbors: int,
+    retrieval_metric: str = "raw",
 ) -> list[_RolloutCase]:
     if observation_mode != "state":
         raise ValueError("retrieval_residual currently supports state observations only")
@@ -474,10 +488,12 @@ def _attach_retrieved_residual_candidates(
             vectorize_toy_observation(case.observation, obs_dim=obs_dim),
             dtype=np.float32,
         )
-        nearest = sorted(
+        nearest = _nearest_retrieval_entries(
             candidates,
-            key=lambda item: float(np.mean((item[1] - query) ** 2)),
-        )[:retrieval_neighbors]
+            query,
+            retrieval_neighbors=retrieval_neighbors,
+            retrieval_metric=retrieval_metric,
+        )
         source_group_ids: list[str] = []
         residuals: list[list[list[float]]] = []
         candidate_types: list[str] = []
@@ -496,6 +512,30 @@ def _attach_retrieved_residual_candidates(
     return output
 
 
+def _nearest_retrieval_entries(
+    candidates: list[tuple[Any, np.ndarray, Any, Any]],
+    query: np.ndarray,
+    *,
+    retrieval_neighbors: int,
+    retrieval_metric: str,
+) -> list[tuple[Any, np.ndarray, Any, Any]]:
+    if retrieval_metric == "raw":
+        return sorted(
+            candidates,
+            key=lambda item: float(np.mean((item[1] - query) ** 2)),
+        )[: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)
+    mean = features.mean(axis=0, dtype=np.float64).astype(np.float32)
+    std = features.std(axis=0, dtype=np.float64).astype(np.float32)
+    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]]
+
+
 def _evaluate_task_cases(
     task_id: str,
     cases: list[_RolloutCase],