diff --git a/cyclo_brain/policy/lerobot/lerobot_engine/prediction.py b/cyclo_brain/policy/lerobot/lerobot_engine/prediction.py index 6e298f3..9db6547 100644 --- a/cyclo_brain/policy/lerobot/lerobot_engine/prediction.py +++ b/cyclo_brain/policy/lerobot/lerobot_engine/prediction.py @@ -21,22 +21,43 @@ logger = logging.getLogger("lerobot_engine") class PredictionMixin: """Policy input batch -> action chunk.""" + # Policies whose predict_action_chunk cannot be called standalone: it + # stacks internal observation queues (self._queues) that ONLY select_action + # populates, so calling it directly raises "stack expects a non-empty + # TensorList" and the robot never receives an action. + # + # DiffusionPolicy.predict_action_chunk: + # batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch ...} + # VQBeTPolicy.predict_action_chunk: same, plus a combined OBS_IMAGES key + # that only select_action builds. + # + # ACT does not use queues (n_obs_steps=1), which is why it works with the + # direct call. Route the queue-based ones through select_action instead; + # they manage their own action-chunk queue internally. + _QUEUE_BASED_POLICIES = {"VQBeTPolicy", "DiffusionPolicy"} + def _predict_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: """Return a chunk tensor of shape (1, T, A).""" assert self._policy is not None + if type(self._policy).__name__ in self._QUEUE_BASED_POLICIES: + return self._select_action_chunk(batch) try: action = self._policy.predict_action_chunk(batch) if action.dim() == 2: action = action.unsqueeze(1) return action - except (NotImplementedError, AttributeError): + except (NotImplementedError, AttributeError, RuntimeError, AssertionError): logger.debug( - "predict_action_chunk unavailable; falling back to select_action" + "predict_action_chunk unavailable/failed; falling back to select_action" ) - action = self._policy.select_action(batch) - if action.dim() == 1: - action = action.unsqueeze(0) - return action.unsqueeze(1) + return self._select_action_chunk(batch) + + def _select_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + """select_action -> (1, 1, A) chunk.""" + action = self._policy.select_action(batch) + if action.dim() == 1: + action = action.unsqueeze(0) + return action.unsqueeze(1) @staticmethod def _to_numpy_chunk(action: torch.Tensor) -> np.ndarray: