Auto-sync: 2026-06-27 10:19:52

dovla_cil/eval/maniskill_policy_rollout.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import pickle
 from collections import Counter, defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any
 
@@ -37,6 +37,7 @@ class _RolloutCase:
     best_action_values: list[list[float]]
     candidate_action_values: list[list[list[float]]]
     candidate_types: list[str]
+    candidate_source_group_id: str | None = None
 
 
 def evaluate_maniskill_policy_rollout(
@@ -74,6 +75,11 @@ def evaluate_maniskill_policy_rollout(
     When ``selection_mode == 'lattice'`` the evaluator scores the pre-generated CIL action
     lattice for the current state with the learned field and executes the selected action. It
     never reads candidate rewards during selection; it only uses the action proposals.
+
+    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.
     """
 
     try:
@@ -88,8 +94,10 @@ def evaluate_maniskill_policy_rollout(
 
     if group_batch_size <= 0:
         raise ValueError("group_batch_size must be positive")
-    if selection_mode not in {"policy", "field", "lattice"}:
-        raise ValueError("selection_mode must be 'policy', 'field', or 'lattice'")
+    if selection_mode not in {"policy", "field", "lattice", "retrieval_lattice"}:
+        raise ValueError(
+            "selection_mode must be 'policy', 'field', 'lattice', or 'retrieval_lattice'"
+        )
     if num_candidates <= 0:
         raise ValueError("num_candidates must be positive")
     if selection_mode == "policy":
@@ -125,6 +133,14 @@ def evaluate_maniskill_policy_rollout(
         group_ids,
         observation_mode=model_config.observation_mode,
     )
+    if selection_mode == "retrieval_lattice":
+        cases = _attach_retrieved_lattice_candidates(
+            dataset,
+            cases,
+            eval_group_ids=group_ids,
+            obs_dim=model_config.obs_dim,
+            observation_mode=model_config.observation_mode,
+        )
     by_task: dict[str, list[_RolloutCase]] = defaultdict(list)
     for case in cases:
         by_task[case.task_id].append(case)
@@ -154,7 +170,7 @@ def evaluate_maniskill_policy_rollout(
             task_summaries[task_id] = _summarize_rows(task_rows)
 
     effective_num_candidates = num_candidates
-    if selection_mode == "lattice":
+    if selection_mode in {"lattice", "retrieval_lattice"}:
         effective_num_candidates = max(
             [int(row.get("lattice_candidate_count", 0)) for row in rows],
             default=0,
@@ -335,6 +351,7 @@ def _evaluate_task_cases(
                         device=device,
                     )
                     if selection_mode == "lattice"
+                    or selection_mode == "retrieval_lattice"
                     else None
                 ),
                 candidate_mask=(
@@ -344,7 +361,8 @@ def _evaluate_task_cases(
                         device=device,
                         exclude_types=lattice_exclude_types,
                     )
-                    if selection_mode == "lattice" and lattice_exclude_types
+                    if selection_mode in {"lattice", "retrieval_lattice"}
+                    and lattice_exclude_types
                     else None
                 ),
             )
@@ -394,6 +412,7 @@ def _evaluate_task_cases(
                         ),
                         "selected_candidate_index": int(candidate_index[index]),
                         "lattice_candidate_count": len(case.candidate_action_values),
+                        "candidate_source_group_id": case.candidate_source_group_id,
                     }
                 )
         finally:
@@ -424,9 +443,9 @@ def _select_action_chunk(
     scored, so no dataset candidate ever leaks into the deployed action.
     """
 
-    if selection_mode == "lattice":
+    if selection_mode in {"lattice", "retrieval_lattice"}:
         if action_candidates is None:
-            raise ValueError("lattice selection requires action_candidates")
+            raise ValueError(f"{selection_mode} selection requires action_candidates")
         return _select_lattice_action_chunk(
             model,
             observations,