Auto-sync: 2026-06-28 22:04:20

dovla_cil/eval/maniskill_policy_rollout.py

@@ -27,6 +27,15 @@ _FIELD_CONDITIONED_RESIDUAL_REDUCERS = {"field_softmax"}
 _RESIDUAL_REDUCERS = (
     {"none"} | _NUMPY_RESIDUAL_REDUCERS | _FIELD_CONDITIONED_RESIDUAL_REDUCERS
 )
+_RETRIEVAL_METRICS = {"raw", "zscore", "task_relative"}
+_MANISKILL_ACTOR_STATE_DIM = 13
+_TASKS_WITH_REFERENCE_ACTOR = {
+    "PickCube-v1",
+    "PushCube-v1",
+    "PullCube-v1",
+    "StackCube-v1",
+    "PegInsertionSide-v1",
+}
 
 
 @dataclass(frozen=True)
@@ -101,7 +110,8 @@ def evaluate_maniskill_policy_rollout(
     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. ``retrieval_metric='zscore'`` standardizes state features by the train-bank
-    statistics for each task before nearest-neighbor lookup; the default ``raw`` metric
+    statistics for each task before nearest-neighbor lookup; ``task_relative`` retrieves by
+    target/reference actor pose blocks rather than full robot state. The default ``raw`` metric
     preserves earlier results exactly.
 
     When ``selection_mode == 'retrieval_residual'`` the evaluator retrieves counterfactual
@@ -159,8 +169,8 @@ def evaluate_maniskill_policy_rollout(
         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"}:
-        raise ValueError("retrieval_metric must be 'raw' or 'zscore'")
+    if retrieval_metric not in _RETRIEVAL_METRICS:
+        raise ValueError("retrieval_metric must be 'raw', 'zscore', or 'task_relative'")
     if retrieval_residual_anchor not in {"expert", "policy"}:
         raise ValueError("retrieval_residual_anchor must be 'expert' or 'policy'")
     if retrieval_residual_reduce not in _RESIDUAL_REDUCERS:
@@ -476,6 +486,7 @@ def _attach_retrieved_lattice_candidates(
             query,
             retrieval_neighbors=retrieval_neighbors,
             retrieval_metric=retrieval_metric,
+            task_id=case.task_id,
         )
         source_group_ids: list[str] = []
         actions: list[list[list[float]]] = []
@@ -571,6 +582,7 @@ def _attach_retrieved_residual_candidates(
             query,
             retrieval_neighbors=retrieval_neighbors,
             retrieval_metric=retrieval_metric,
+            task_id=case.task_id,
         )
         source_group_ids: list[str] = []
         residuals: list[list[list[float]]] = []
@@ -684,6 +696,7 @@ def _nearest_retrieval_entries(
     *,
     retrieval_neighbors: int,
     retrieval_metric: str,
+    task_id: str | None = None,
 ) -> list[tuple[Any, np.ndarray, Any, Any]]:
     return [
         entry
@@ -692,6 +705,7 @@ def _nearest_retrieval_entries(
             query,
             retrieval_neighbors=retrieval_neighbors,
             retrieval_metric=retrieval_metric,
+            task_id=task_id,
         )
     ]
 
@@ -702,6 +716,7 @@ def _nearest_retrieval_entries_with_distances(
     *,
     retrieval_neighbors: int,
     retrieval_metric: str,
+    task_id: str | None = None,
 ) -> list[tuple[tuple[Any, np.ndarray, Any, Any], float]]:
     if retrieval_metric == "raw":
         scored = [
@@ -709,8 +724,26 @@ def _nearest_retrieval_entries_with_distances(
             for item in candidates
         ]
         return sorted(scored, key=lambda item: item[1])[:retrieval_neighbors]
+    if retrieval_metric == "task_relative":
+        normalized_query = _task_relative_retrieval_feature(query, task_id=task_id)
+        scored = [
+            (
+                item,
+                float(
+                    np.mean(
+                        (
+                            _task_relative_retrieval_feature(item[1], task_id=task_id)
+                            - normalized_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'")
+        raise ValueError("retrieval_metric must be 'raw', 'zscore', or 'task_relative'")
     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)
@@ -725,6 +758,38 @@ def _nearest_retrieval_entries_with_distances(
     ]
 
 
+def _task_relative_retrieval_feature(
+    feature: np.ndarray,
+    *,
+    task_id: str | None,
+) -> np.ndarray:
+    values = np.asarray(feature, dtype=np.float32).reshape(-1)
+    if values.size < _MANISKILL_ACTOR_STATE_DIM:
+        return values
+    target = values[:_MANISKILL_ACTOR_STATE_DIM]
+    parts = [
+        4.0 * target[:3],
+        target[3:7],
+    ]
+    has_reference = (
+        task_id in _TASKS_WITH_REFERENCE_ACTOR
+        and values.size >= 2 * _MANISKILL_ACTOR_STATE_DIM
+    )
+    if has_reference:
+        start = _MANISKILL_ACTOR_STATE_DIM
+        reference = values[start : start + _MANISKILL_ACTOR_STATE_DIM]
+        parts.extend(
+            [
+                2.0 * reference[:3],
+                4.0 * (target[:3] - reference[:3]),
+                reference[3:7],
+            ]
+        )
+    if task_id in {"LiftPegUpright-v1", "PegInsertionSide-v1"}:
+        parts.append(2.0 * target[3:7])
+    return np.concatenate(parts).astype(np.float32)
+
+
 def _kernel_weights_from_distances(distances: list[float]) -> list[float]:
     if not distances:
         return []