Auto-sync: 2026-06-27 10:45:23

dovla_cil/training/trainer.py

@@ -65,6 +65,7 @@ class TrainerConfig:
     backbone_local_files_only: bool = True
     backbone_feature_cache: str | Path | None = None
     backbone_feature_batch_size: int = 64
+    policy_target_types: tuple[str, ...] = ()
 
     def __post_init__(self) -> None:
         if self.lr is not None:
@@ -98,6 +99,12 @@ class TrainerConfig:
                 raise ValueError("CLIP backbone requires backbone_model")
         if self.backbone_feature_batch_size <= 0:
             raise ValueError("backbone_feature_batch_size must be positive")
+        if isinstance(self.policy_target_types, str):
+            self.policy_target_types = tuple(
+                item.strip()
+                for item in self.policy_target_types.split(",")
+                if item.strip()
+            )
 
 
 class DoVLATrainer:
@@ -268,7 +275,10 @@ class DoVLATrainer:
             pred_reward, [record.group_id for record in records]
         )
 
-        best_records = _best_records_by_group(records)
+        best_records = _best_records_by_group(
+            records,
+            candidate_types=self.config.policy_target_types,
+        )
         best_obs = self._obs_tensor(best_records)
         best_actions = self._action_tensor(best_records)
         best_instructions = [record.instruction for record in best_records]
@@ -655,12 +665,28 @@ def _group_regret_from_potential(potential, group_ids: list[str]):
     return output
 
 
-def _best_records_by_group(records: list[CILRecord]) -> list[CILRecord]:
+def _best_records_by_group(
+    records: list[CILRecord], *, candidate_types: tuple[str, ...] = ()
+) -> list[CILRecord]:
+    allowed = set(candidate_types)
     best: dict[str, CILRecord] = {}
     for record in records:
+        if allowed and record.candidate_type not in allowed:
+            continue
         current = best.get(record.group_id)
         if current is None or _record_score(record) > _record_score(current):
             best[record.group_id] = record
+    if allowed:
+        missing_group_ids = {record.group_id for record in records} - set(best)
+        if missing_group_ids:
+            fallback = _best_records_by_group(records)
+            best.update(
+                {
+                    record.group_id: record
+                    for record in fallback
+                    if record.group_id in missing_group_ids
+                }
+            )
     return list(best.values())