Add teleoperate

lerobot/common/utils/visualization_utils.py

@@ -0,0 +1,12 @@
+import os
+
+import rerun as rr
+
+
+def _init_rerun(session_name: str = "lerobot_control_loop") -> None:
+    """Initializes the Rerun SDK for visualizing the control loop."""
+    batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000")
+    os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size
+    rr.init(session_name)
+    memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
+    rr.spawn(memory_limit=memory_limit)

lerobot/teleoperate.py

@@ -0,0 +1,106 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import time
+from dataclasses import asdict, dataclass
+from pprint import pformat
+
+import draccus
+import numpy as np
+import rerun as rr
+
+from lerobot.common.cameras import intel, opencv  # noqa: F401
+from lerobot.common.robots import (  # noqa: F401
+    Robot,
+    RobotConfig,
+    koch_follower,
+    make_robot_from_config,
+    so100_follower,
+)
+from lerobot.common.teleoperators import (
+    Teleoperator,
+    TeleoperatorConfig,
+    make_teleoperator_from_config,
+)
+from lerobot.common.utils.utils import init_logging, move_cursor_up
+from lerobot.common.utils.visualization_utils import _init_rerun
+
+from .common.teleoperators import koch_leader, so100_leader  # noqa: F401
+
+
+@dataclass
+class TeleoperateConfig:
+    teleop: TeleoperatorConfig
+    robot: RobotConfig
+    # Limit the maximum frames per second. By default, no limit.
+    fps: int | None = None
+    teleop_time_s: float | None = None
+    # Display all cameras on screen
+    display_data: bool = False
+
+
+def teleop_loop(teleop: Teleoperator, robot: Robot, display_data: bool = False):
+    display_len = max(len(key) for key in robot.action_features)
+    while True:
+        start = time.perf_counter()
+        action = teleop.get_action()
+        if display_data:
+            observation = robot.get_observation()
+            for obs, val in observation.items():
+                if isinstance(val, float):
+                    rr.log(f"observation_{obs}", rr.Scalar(val))
+                elif isinstance(val, np.ndarray):
+                    rr.log(f"observation_{obs}", rr.Image(val), static=True)
+            for act, val in action.items():
+                if isinstance(val, float):
+                    rr.log(f"action_{act}", rr.Scalar(val))
+
+        robot.send_action(action)
+        loop_s = time.perf_counter() - start
+
+        print("\n" + "-" * (display_len + 10))
+        print(f"{'NAME':<{display_len}} | {'NORM':>7}")
+        for motor, value in action.items():
+            print(f"{motor:<{display_len}} | {value:>7.2f}")
+        print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
+        move_cursor_up(len(action) + 5)
+
+
+@draccus.wrap()
+def control_robot(cfg: TeleoperateConfig):
+    init_logging()
+    logging.info(pformat(asdict(cfg)))
+    if cfg.display_data:
+        _init_rerun(session_name="teleoperation")
+
+    teleop = make_teleoperator_from_config(cfg.teleop)
+    robot = make_robot_from_config(cfg.robot)
+
+    teleop.connect()
+    robot.connect()
+
+    try:
+        teleop_loop(teleop, robot, display_data=cfg.display_data)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        if cfg.display_data:
+            rr.rerun_shutdown()
+        teleop.disconnect()
+        robot.disconnect()
+
+
+if __name__ == "__main__":
+    control_robot()