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/tdmpc/modeling_tdmpc.py
+++ b/lerobot/common/policies/tdmpc/modeling_tdmpc.py
@@ -134,7 +134,7 @@ class TDMPCPolicy(nn.Module, PyTorchModelHubMixin):
        self._prev_mean: torch.Tensor | None = None

    @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]):
+    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
        """Select a single action given environment observations."""
        batch = self.normalize_inputs(batch)
        batch["observation.image"] = batch[self.input_image_key]