Auto-sync: 2026-06-27 10:26:03

dovla_cil/eval/maniskill_policy_rollout.py

@@ -115,7 +115,7 @@ def evaluate_maniskill_policy_rollout(
 
     trainer_config = checkpoint.get("trainer_config", {})
     dataset = CILDataset(dataset_dir)
-    group_ids = (
+    split_group_ids = (
         list(dataset.group_ids)
         if all_groups
         else _validation_group_ids(
@@ -124,6 +124,7 @@ def evaluate_maniskill_policy_rollout(
             seed=int(trainer_config.get("seed", 0)),
         )
     )
+    group_ids = list(split_group_ids)
     if max_groups is not None:
         if max_groups <= 0:
             raise ValueError("max_groups must be positive when provided")
@@ -134,10 +135,12 @@ def evaluate_maniskill_policy_rollout(
         observation_mode=model_config.observation_mode,
     )
     if selection_mode == "retrieval_lattice":
+        if all_groups:
+            raise ValueError("retrieval_lattice requires a held-out validation split")
         cases = _attach_retrieved_lattice_candidates(
             dataset,
             cases,
-            eval_group_ids=group_ids,
+            heldout_group_ids=split_group_ids,
             obs_dim=model_config.obs_dim,
             observation_mode=model_config.observation_mode,
         )
@@ -268,6 +271,67 @@ def _prepare_rollout_cases(
     return cases
 
 
+def _attach_retrieved_lattice_candidates(
+    dataset: CILDataset,
+    cases: list[_RolloutCase],
+    *,
+    heldout_group_ids: list[str],
+    obs_dim: int,
+    observation_mode: str,
+) -> list[_RolloutCase]:
+    if observation_mode != "state":
+        raise ValueError("retrieval_lattice currently supports state observations only")
+    heldout = set(heldout_group_ids)
+    bank: dict[str, list[tuple[str, np.ndarray, list[list[list[float]]], list[str]]]] = (
+        defaultdict(list)
+    )
+    for group_id in dataset.group_ids:
+        if group_id in heldout:
+            continue
+        records = dataset.get_group(group_id)
+        if not records:
+            continue
+        task_ids = {record.task_id for record in records}
+        if len(task_ids) != 1:
+            continue
+        feature = np.asarray(
+            vectorize_toy_observation(records[0].observation_inline or {}, obs_dim=obs_dim),
+            dtype=np.float32,
+        )
+        bank[next(iter(task_ids))].append(
+            (
+                group_id,
+                feature,
+                [_numeric_action_values(record) for record in records],
+                [str(record.candidate_type) for record in records],
+            )
+        )
+
+    output: list[_RolloutCase] = []
+    for case in cases:
+        candidates = bank.get(case.task_id, [])
+        if not candidates:
+            output.append(case)
+            continue
+        query = np.asarray(
+            vectorize_toy_observation(case.observation, obs_dim=obs_dim),
+            dtype=np.float32,
+        )
+        source_group_id, _feature, actions, candidate_types = min(
+            candidates,
+            key=lambda item: float(np.mean((item[1] - query) ** 2)),
+        )
+        output.append(
+            replace(
+                case,
+                candidate_action_values=actions,
+                candidate_types=candidate_types,
+                candidate_source_group_id=source_group_id,
+            )
+        )
+    return output
+
+
 def _evaluate_task_cases(
     task_id: str,
     cases: list[_RolloutCase],