| |
| |
| |
| |
| @@ -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: |
|
|
