diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 724f920f3..61a856f0b 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -40,6 +40,7 @@ T = TypeVar("T", bound="PreTrainedPolicy") class ActionSelectKwargs(TypedDict, total=False): noise: Tensor | None + return_extra: bool class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): @@ -187,20 +188,32 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): raise NotImplementedError @abc.abstractmethod - def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor: + def predict_action_chunk( + self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs] + ) -> Tensor | tuple[Tensor, dict[str, Tensor]]: """Returns the action chunk (for action chunking policies) for a given observation, potentially in batch mode. Child classes using action chunking should use this method within `select_action` to form the action chunk cached for selection. + + By default returns just the action `Tensor`. If `return_extra=True`, returns `(action, extra)` + where `extra` is a (possibly empty) `dict[str, Tensor]` of auxiliary outputs a policy may + expose (e.g. world-model predictions). Policies that produce nothing extra may ignore the kwarg. """ raise NotImplementedError @abc.abstractmethod - def select_action(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor: + def select_action( + self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs] + ) -> Tensor | tuple[Tensor, dict[str, Tensor]]: """Return one action to run in the environment (potentially in batch mode). When the model uses a history of observations, or outputs a sequence of actions, this method deals with caching. + + By default returns just the action `Tensor`. If `return_extra=True`, returns `(action, extra)` + where `extra` is a (possibly empty) `dict[str, Tensor]` of auxiliary outputs a policy may + expose (e.g. world-model predictions). Policies that produce nothing extra may ignore the kwarg. """ raise NotImplementedError