File size: 2,716 Bytes
757177d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | 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:
|