Auto-sync: 2026-06-28 07:13:10

dovla_cil/training/trainer.py

@@ -117,11 +117,14 @@ class DoVLATrainer:
         self.output_dir = ensure_dir(config.output_dir)
         self.dataset = CILDataset(config.dataset_dir)
         self.policy_target_record_ids = _load_policy_target_map(config.policy_target_map)
+        self.policy_target_action_values = _load_policy_target_action_map(
+            config.policy_target_map
+        )
         self.train_group_ids, self.val_group_ids = _split_group_ids(
             self.dataset.group_ids, val_fraction=config.val_fraction, seed=config.seed
         )
         self.policy_target_map_coverage = _policy_target_map_coverage(
-            self.policy_target_record_ids,
+            set(self.policy_target_record_ids) | set(self.policy_target_action_values),
             train_group_ids=self.train_group_ids,
             val_group_ids=self.val_group_ids,
         )
@@ -168,7 +171,7 @@ class DoVLATrainer:
 
     def train(self) -> dict[str, Any]:
         write_json(self._resolved_config(), self.output_dir / "resolved_config.json")
-        if self.policy_target_record_ids:
+        if self.policy_target_record_ids or self.policy_target_action_values:
             coverage = self.policy_target_map_coverage
             print(
                 "policy_target_map_coverage "
@@ -292,13 +295,8 @@ class DoVLATrainer:
             pred_reward, [record.group_id for record in records]
         )
 
-        best_records = _best_records_by_group(
-            records,
-            candidate_types=self.config.policy_target_types,
-            target_record_ids=self.policy_target_record_ids,
-        )
+        best_records, best_actions = self._policy_bc_targets(records)
         best_obs = self._obs_tensor(best_records)
-        best_actions = self._action_tensor(best_records)
         best_instructions = [record.instruction for record in best_records]
         pred_best_actions = self.model.forward_policy(best_obs, best_instructions)
 
@@ -362,6 +360,50 @@ class DoVLATrainer:
         )
         return total, metrics
 
+    def _policy_bc_targets(self, records: list[CILRecord]) -> tuple[list[CILRecord], Any]:
+        assert torch is not None
+        target_records: list[CILRecord] = []
+        target_values: list[list[list[float]]] = []
+        used_group_ids: set[str] = set()
+        grouped: dict[str, list[CILRecord]] = {}
+        for record in records:
+            grouped.setdefault(record.group_id, []).append(record)
+
+        for group_id, group_records in grouped.items():
+            mapped_action = self.policy_target_action_values.get(group_id)
+            if mapped_action is None:
+                continue
+            target_records.append(group_records[0])
+            target_values.append(
+                _coerce_policy_target_action(
+                    mapped_action,
+                    action_dim=self.config.action_dim,
+                    action_horizon=self.config.action_horizon,
+                )
+            )
+            used_group_ids.add(group_id)
+
+        fallback_records = _best_records_by_group(
+            [record for record in records if record.group_id not in used_group_ids],
+            candidate_types=self.config.policy_target_types,
+            target_record_ids=self.policy_target_record_ids,
+        )
+        for record in fallback_records:
+            target_records.append(record)
+            target_values.append(
+                vectorize_toy_action(
+                    record.action_chunk,
+                    action_dim=self.config.action_dim,
+                    action_horizon=self.config.action_horizon,
+                )
+            )
+        best_actions = torch.tensor(
+            target_values,
+            dtype=torch.float32,
+            device=self.device,
+        )
+        return target_records, best_actions
+
     def _rank_loss(self, pred_reward, target_reward, pair_indices: list[tuple[int, int]]):
         assert torch is not None
         if not pair_indices:
@@ -743,18 +785,51 @@ def _load_policy_target_map(path: str | Path | None) -> dict[str, str]:
     return output
 
 
+def _load_policy_target_action_map(path: str | Path | None) -> dict[str, list[list[float]]]:
+    if path is None:
+        return {}
+    payload = json.loads(Path(path).read_text())
+    raw_targets = payload.get("targets", payload)
+    output: dict[str, list[list[float]]] = {}
+    for group_id, value in raw_targets.items():
+        if not isinstance(value, dict):
+            continue
+        raw_action = value.get("action_values", value.get("action"))
+        if raw_action is None:
+            continue
+        output[str(group_id)] = [
+            [float(item) for item in row]
+            for row in raw_action
+        ]
+    return output
+
+
+def _coerce_policy_target_action(
+    action_values: list[list[float]],
+    *,
+    action_dim: int,
+    action_horizon: int,
+) -> list[list[float]]:
+    return vectorize_toy_action(
+        action_values,
+        action_dim=action_dim,
+        action_horizon=action_horizon,
+    )
+
+
 def _policy_target_map_coverage(
-    target_record_ids: dict[str, str],
+    target_group_ids: set[str] | dict[str, Any],
     *,
     train_group_ids: list[str],
     val_group_ids: list[str],
 ) -> dict[str, float | int]:
-    train_covered = sum(1 for group_id in train_group_ids if group_id in target_record_ids)
-    val_covered = sum(1 for group_id in val_group_ids if group_id in target_record_ids)
+    target_groups = set(target_group_ids)
+    train_covered = sum(1 for group_id in train_group_ids if group_id in target_groups)
+    val_covered = sum(1 for group_id in val_group_ids if group_id in target_groups)
     train_groups = len(train_group_ids)
     val_groups = len(val_group_ids)
     return {
-        "num_targets": len(target_record_ids),
+        "num_targets": len(target_groups),
         "train_groups": train_groups,
         "train_covered": train_covered,
         "train_coverage": train_covered / train_groups if train_groups else 0.0,