From 2357d4acebfe57f0cb4b651a78218188cfd29212 Mon Sep 17 00:00:00 2001
From: Simon Alibert <simon.alibert@huggingface.co>
Date: Wed, 12 Mar 2025 17:15:39 +0100
Subject: [PATCH] Update base classes typing

---
 lerobot/common/robots/robot.py               | 9 +++++----
 lerobot/common/teleoperators/teleoperator.py | 7 +++----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/lerobot/common/robots/robot.py b/lerobot/common/robots/robot.py
index 50fd9154..d08b8c32 100644
--- a/lerobot/common/robots/robot.py
+++ b/lerobot/common/robots/robot.py
@@ -1,12 +1,13 @@
 import abc
-
-import numpy as np
+from typing import Any
 
 from lerobot.common.constants import HF_LEROBOT_CALIBRATION, ROBOTS
 
 from .config import RobotConfig
 
 
+# TODO(aliberts): action/obs typing such as Generic[ObsType, ActType] similar to gym.Env ?
+# https://github.com/Farama-Foundation/Gymnasium/blob/3287c869f9a48d99454306b0d4b4ec537f0f35e3/gymnasium/core.py#L23
 class Robot(abc.ABC):
     """The main LeRobot class for implementing robots."""
 
@@ -45,12 +46,12 @@ class Robot(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def get_observation(self) -> dict[str, np.ndarray]:
+    def get_observation(self) -> dict[str, Any]:
         """Gets observation from the robot."""
         pass
 
     @abc.abstractmethod
-    def send_action(self, action: np.ndarray) -> np.ndarray:
+    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
         """Sends actions to the robot."""
         pass
 
diff --git a/lerobot/common/teleoperators/teleoperator.py b/lerobot/common/teleoperators/teleoperator.py
index f06b1983..2aa1d662 100644
--- a/lerobot/common/teleoperators/teleoperator.py
+++ b/lerobot/common/teleoperators/teleoperator.py
@@ -1,6 +1,5 @@
 import abc
-
-import numpy as np
+from typing import Any
 
 from lerobot.common.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS
 
@@ -41,12 +40,12 @@ class Teleoperator(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def get_action(self) -> np.ndarray:
+    def get_action(self) -> dict[str, Any]:
         """Gets the action to send to a teleoperator."""
         pass
 
     @abc.abstractmethod
-    def send_feedback(self, feedback: np.ndarray) -> None:
+    def send_feedback(self, feedback: dict[str, Any]) -> None:
         """Sends feedback captured from a robot to the teleoperator."""
         pass