Add stretch boilerplate

lerobot/common/robot_devices/robots/factory.py

@@ -1,7 +1,9 @@
 import hydra
 from omegaconf import DictConfig
 
+from lerobot.common.robot_devices.robots.utils import Robot
 
-def make_robot(cfg: DictConfig):
+
+def make_robot(cfg: DictConfig) -> Robot:
     robot = hydra.utils.instantiate(cfg)
     return robot

lerobot/common/robot_devices/robots/stretch.py

@@ -0,0 +1,43 @@
+from stretch_body.gamepad_teleop import GamePadTeleop
+from stretch_body.robot import Robot as Stretch
+
+
+class LeRobotStretchTeleop(GamePadTeleop):
+    """Wrapper of stretch_body.gamepad_teleop.GamePadTeleop"""
+
+    def __init__(self):
+        super().__init__()
+
+
+class LeRobotStretch(Stretch):
+    """Wrapper of stretch_body.robot.Robot"""
+
+    def __init__(self, teleoperate: bool = False, **kwargs):
+        super().__init__()
+        self.robot_type = "stretch"
+        self.is_connected = False
+        self.teleop = GamePadTeleop(robot_instance=False) if teleoperate else None
+
+    def connect(self):
+        self.is_connected = self.startup()
+
+    def run_calibration(self):
+        if not self.is_homed():
+            self.home()
+
+    def teleop_step(self, record_data=False):
+        if self.teleop is None:
+            raise ValueError
+        ...
+
+    def capture_observation(self): ...
+
+    def send_action(self, action): ...
+
+    def disconnect(self):
+        self.stop()
+        if self.teleop is not None:
+            self.teleop.stop()
+
+    def __del__(self):
+        self.disconnect()

lerobot/common/robot_devices/robots/utils.py

@@ -9,8 +9,9 @@ def get_arm_id(name, arm_type):
 
 
 class Robot(Protocol):
-    def init_teleop(self): ...
+    def connect(self): ...
     def run_calibration(self): ...
     def teleop_step(self, record_data=False): ...
     def capture_observation(self): ...
     def send_action(self, action): ...
+    def disconnect(self): ...

lerobot/configs/robot/stretch.yaml

@@ -0,0 +1,3 @@
+_target_: lerobot.common.robot_devices.robots.stretch.LeRobotStretch
+robot_type: stretch
+teleoperate: false

lerobot/scripts/control_robot.py

@@ -127,6 +127,7 @@
 from lerobot.common.datasets.video_utils import encode_video_frames
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.robot_devices.robots.factory import make_robot
+from lerobot.common.robot_devices.robots.stretch import LeRobotStretch
 from lerobot.common.robot_devices.robots.utils import Robot, get_arm_id
 from lerobot.common.robot_devices.utils import busy_wait
 from lerobot.common.utils.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
@@ -242,6 +243,13 @@ def is_headless():
 
 
 def calibrate(robot: Robot, arms: list[str] | None):
+    if isinstance(robot, LeRobotStretch):
+        if not robot.is_connected:
+            robot.connect()
+        if not robot.is_homed():
+            robot.home()
+        return
+
     available_arms = []
     for name in robot.follower_arms:
         arm_id = get_arm_id(name, "follower")