Auto-sync: 2026-06-27 09:54:24

dovla_cil/eval/maniskill_policy_rollout.py

@@ -292,6 +292,13 @@ def _evaluate_task_cases(
                     dtype=torch.float32,
                     device=device,
                 )
+            env_dim = _env_action_dim(env)
+            action_low, action_high = _action_bounds_for_tensor(
+                env,
+                torch=torch,
+                device=device,
+                action_dim=model_config.action_dim,
+            )
             predicted, candidate_index = _select_action_chunk(
                 model,
                 observations,
@@ -301,9 +308,10 @@ def _evaluate_task_cases(
                 num_candidates=num_candidates,
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed + start,
+                action_low=action_low,
+                action_high=action_high,
             )
             predicted_np = predicted.detach().cpu().numpy().astype(np.float32)
-            env_dim = _env_action_dim(env)
             adapted = _adapt_action_dim(predicted_np, env_dim)
             adapted = _clip_to_action_space(adapted, env)
             after_state, rewards, successes, restore_error = execute_grouped_action_lattice_batch(
@@ -365,6 +373,8 @@ def _select_action_chunk(
     num_candidates: int,
     candidate_sigma: float,
     selection_seed: int,
+    action_low: Any | None = None,
+    action_high: Any | None = None,
 ) -> tuple[Any, np.ndarray]:
     """Return the executed action chunk per state and the chosen candidate index.
 
@@ -381,7 +391,7 @@ def _select_action_chunk(
 
     generator = torch.Generator(device=policy_mean.device)
     generator.manual_seed(int(selection_seed))
-    candidates = [policy_mean]
+    candidates = [_clamp_action_tensor(policy_mean, action_low=action_low, action_high=action_high)]
     for _ in range(num_candidates - 1):
         noise = torch.randn(
             policy_mean.shape,
@@ -389,10 +399,16 @@ def _select_action_chunk(
             device=policy_mean.device,
             dtype=policy_mean.dtype,
         )
-        candidates.append(policy_mean + candidate_sigma * noise)
+        candidates.append(
+            _clamp_action_tensor(
+                policy_mean + candidate_sigma * noise,
+                action_low=action_low,
+                action_high=action_high,
+            )
+        )
 
     best_potential = None
-    best_actions = policy_mean.clone()
+    best_actions = candidates[0].clone()
     best_index = torch.zeros(batch_size, dtype=torch.long, device=policy_mean.device)
     for cand_idx, candidate in enumerate(candidates):
         field = model.forward_field(observations, instructions, candidate)
@@ -413,6 +429,39 @@ def _select_action_chunk(
     return best_actions, best_index.detach().cpu().numpy()
 
 
+def _action_bounds_for_tensor(
+    env: Any,
+    *,
+    torch: Any,
+    device: str,
+    action_dim: int,
+) -> tuple[Any | None, Any | None]:
+    space = getattr(env, "single_action_space", None) or getattr(env, "action_space", None)
+    low = getattr(space, "low", None)
+    high = getattr(space, "high", None)
+    if low is None or high is None:
+        return None, None
+    low_arr = np.asarray(low, dtype=np.float32).reshape(-1)
+    high_arr = np.asarray(high, dtype=np.float32).reshape(-1)
+    if low_arr.size < action_dim or high_arr.size < action_dim:
+        return None, None
+    low_tensor = torch.as_tensor(low_arr[-action_dim:], dtype=torch.float32, device=device).reshape(
+        1, 1, action_dim
+    )
+    high_tensor = torch.as_tensor(
+        high_arr[-action_dim:], dtype=torch.float32, device=device
+    ).reshape(1, 1, action_dim)
+    return low_tensor, high_tensor
+
+
+def _clamp_action_tensor(action: Any, *, action_low: Any | None, action_high: Any | None) -> Any:
+    if action_low is not None:
+        action = action.clamp_min(action_low)
+    if action_high is not None:
+        action = action.clamp_max(action_high)
+    return action
+
+
 def _source_dataset(record: CILRecord, fallback: Path) -> Path:
     raw = record.metadata.get("source_dataset")
     return Path(str(raw)).resolve() if raw else fallback.resolve()

dovla_cil/training/trainer.py

@@ -148,6 +148,7 @@ class DoVLATrainer:
             )
             self.wandb_run = self._maybe_init_wandb()
         self.best_metric: float | None = None
+        self.best_policy_metric: float | None = None
 
     def train(self) -> dict[str, Any]:
         write_json(self._resolved_config(), self.output_dir / "resolved_config.json")
@@ -156,6 +157,7 @@ class DoVLATrainer:
 
         history: list[dict[str, Any]] = []
         best_metrics: dict[str, float] | None = None
+        best_policy_metrics: dict[str, float] | None = None
         for epoch in range(1, self.config.epochs + 1):
             train_metrics = self._run_epoch(self.train_group_ids, train=True, epoch=epoch)
             val_metrics = self._run_epoch(self.val_group_ids, train=False, epoch=epoch)
@@ -170,11 +172,23 @@ class DoVLATrainer:
             if self._is_better(val_metrics):
                 best_metrics = val_metrics
                 self._save_checkpoint("best.pt", epoch=epoch, metrics=epoch_summary)
+            if self._is_better_policy(val_metrics):
+                best_policy_metrics = val_metrics
+                self._save_checkpoint("best_policy.pt", epoch=epoch, metrics=epoch_summary)
 
         if best_metrics is None and history:
             self._save_checkpoint("best.pt", epoch=history[-1]["epoch"], metrics=history[-1])
             best_metrics = history[-1]["val"]
-        result = {"history": history, "best": best_metrics or {}}
+        if best_policy_metrics is None and history:
+            self._save_checkpoint(
+                "best_policy.pt", epoch=history[-1]["epoch"], metrics=history[-1]
+            )
+            best_policy_metrics = history[-1]["val"]
+        result = {
+            "history": history,
+            "best": best_metrics or {},
+            "best_policy": best_policy_metrics or {},
+        }
         write_json(result, self.output_dir / "metrics.json")
         if self.wandb_run is not None:
             self.wandb_run.finish()
@@ -451,6 +465,13 @@ class DoVLATrainer:
             return True
         return False
 
+    def _is_better_policy(self, val_metrics: dict[str, float]) -> bool:
+        bc_loss = float(val_metrics.get("bc_loss", float("inf")))
+        if self.best_policy_metric is None or bc_loss < self.best_policy_metric:
+            self.best_policy_metric = bc_loss
+            return True
+        return False
+
     def _resolve_device(self, device: str) -> str:
         if torch is None:
             return "cpu"
@@ -496,9 +517,15 @@ class DoVLATrainer:
         train_metrics = _heuristic_metrics(self.dataset, self.train_group_ids)
         val_metrics = _heuristic_metrics(self.dataset, self.val_group_ids)
         history = [{"epoch": 1, "train": train_metrics, "val": val_metrics}]
-        summary = {"history": history, "best": val_metrics, "torch_available": False}
+        summary = {
+            "history": history,
+            "best": val_metrics,
+            "best_policy": val_metrics,
+            "torch_available": False,
+        }
         self._save_checkpoint("latest.pt", epoch=1, metrics=history[0])
         self._save_checkpoint("best.pt", epoch=1, metrics=history[0])
+        self._save_checkpoint("best_policy.pt", epoch=1, metrics=history[0])
         write_json(summary, self.output_dir / "metrics.json")
         print(
             "torch is not installed; wrote deterministic heuristic smoke checkpoints "

logs/auto_sync_hf.log

@@ -130,3 +130,4 @@ No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
+No files have been modified since last commit. Skipping to prevent empty commit.