Improves Type Annotations (#252)

2024-06-10 19:09:48 +01:00
parent a06598678c
commit 54c9776bde
7 changed files with 54 additions and 23 deletions
--- a/lerobot/common/policies/policy_protocol.py
+++ b/lerobot/common/policies/policy_protocol.py
@@ -57,7 +57,7 @@ class Policy(Protocol):
        other items should be logging-friendly, native Python types.
        """

-    def select_action(self, batch: dict[str, Tensor]):
+    def select_action(self, batch: dict[str, Tensor]) -> 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