forked from tangger/lerobot
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
# 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.
|
|
|
|
"""
|
|
Simple script to control a robot from teleoperation.
|
|
|
|
Example:
|
|
|
|
```shell
|
|
python -m lerobot.teleoperate \
|
|
--robot.type=so100_follower \
|
|
--robot.port=/dev/tty.usbmodem58760431541 \
|
|
--robot.cameras="{laptop: {type: opencv, camera_index: 0}}" \
|
|
--robot.id=black \
|
|
--teleop.type=so100_leader \
|
|
--teleop.port=/dev/tty.usbmodem58760431551 \
|
|
--teleop.id=blue
|
|
```
|
|
"""
|
|
|
|
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, duration: float | None = None
|
|
):
|
|
display_len = max(len(key) for key in robot.action_features)
|
|
start = time.perf_counter()
|
|
while True:
|
|
loop_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() - loop_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)")
|
|
|
|
if duration is not None and time.perf_counter() - start >= duration:
|
|
return
|
|
|
|
move_cursor_up(len(action) + 5)
|
|
|
|
|
|
@draccus.wrap()
|
|
def teleoperate(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, duration=cfg.teleop_time_s)
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
if cfg.display_data:
|
|
rr.rerun_shutdown()
|
|
teleop.disconnect()
|
|
robot.disconnect()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
teleoperate()
|