Add calibrate
This commit is contained in:
60
lerobot/calibrate.py
Normal file
60
lerobot/calibrate.py
Normal file
@@ -0,0 +1,60 @@
|
||||
# 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
|
||||
from dataclasses import asdict, dataclass
|
||||
from pprint import pformat
|
||||
|
||||
import draccus
|
||||
|
||||
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 ( # noqa: F401
|
||||
Teleoperator,
|
||||
TeleoperatorConfig,
|
||||
make_teleoperator_from_config,
|
||||
)
|
||||
from lerobot.common.utils.utils import init_logging
|
||||
|
||||
from .common.teleoperators import koch_leader, so100_leader # noqa: F401
|
||||
|
||||
|
||||
@dataclass
|
||||
class CalibrateConfig:
|
||||
device: RobotConfig | TeleoperatorConfig
|
||||
|
||||
|
||||
@draccus.wrap()
|
||||
def calibrate(cfg: CalibrateConfig):
|
||||
init_logging()
|
||||
logging.info(pformat(asdict(cfg)))
|
||||
|
||||
if isinstance(cfg.device, RobotConfig):
|
||||
device = make_robot_from_config(cfg.device)
|
||||
elif isinstance(cfg.device, TeleoperatorConfig):
|
||||
device = make_teleoperator_from_config(cfg.device)
|
||||
|
||||
device.connect(calibrate=False)
|
||||
device.calibrate()
|
||||
device.disconnect()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
calibrate()
|
||||
@@ -85,7 +85,7 @@ class KochFollower(Robot):
|
||||
# TODO(aliberts): add cam.is_connected for cam in self.cameras
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
@@ -94,7 +94,7 @@ class KochFollower(Robot):
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
for cam in self.cameras.values():
|
||||
|
||||
@@ -104,12 +104,12 @@ class LeKiwi(Robot):
|
||||
# TODO(aliberts): add cam.is_connected for cam in self.cameras
|
||||
return self.bus.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.bus.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
for cam in self.cameras.values():
|
||||
|
||||
@@ -82,7 +82,7 @@ class MossRobot(Robot):
|
||||
# TODO(aliberts): add cam.is_connected for cam in self.cameras
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
@@ -91,7 +91,7 @@ class MossRobot(Robot):
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
# Connect the cameras
|
||||
|
||||
@@ -48,7 +48,7 @@ class Robot(abc.ABC):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""Connects to the robot."""
|
||||
pass
|
||||
|
||||
|
||||
@@ -82,7 +82,7 @@ class SO100Follower(Robot):
|
||||
# TODO(aliberts): add cam.is_connected for cam in self.cameras
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
@@ -91,7 +91,7 @@ class SO100Follower(Robot):
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
# Connect the cameras
|
||||
|
||||
@@ -78,7 +78,7 @@ class ViperX(Robot):
|
||||
# TODO(aliberts): add cam.is_connected for cam in self.cameras
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
@@ -87,7 +87,7 @@ class ViperX(Robot):
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
for cam in self.cameras.values():
|
||||
|
||||
@@ -69,12 +69,12 @@ class KochLeader(Teleoperator):
|
||||
def is_connected(self) -> bool:
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
self.configure()
|
||||
|
||||
@@ -66,12 +66,12 @@ class SO100Leader(Teleoperator):
|
||||
def is_connected(self) -> bool:
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
self.configure()
|
||||
|
||||
@@ -46,7 +46,7 @@ class Teleoperator(abc.ABC):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""Connects to the teleoperator."""
|
||||
pass
|
||||
|
||||
|
||||
@@ -69,12 +69,12 @@ class WidowX(Teleoperator):
|
||||
def is_connected(self) -> bool:
|
||||
return self.arm.is_connected
|
||||
|
||||
def connect(self):
|
||||
def connect(self, calibrate: bool = True):
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self.arm.connect()
|
||||
if not self.is_calibrated:
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
|
||||
self.configure()
|
||||
|
||||
@@ -67,11 +67,13 @@ class MockRobot(Robot):
|
||||
def is_connected(self) -> bool:
|
||||
return self._is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self._is_connected = True
|
||||
if calibrate:
|
||||
self.calibrate()
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
|
||||
@@ -51,11 +51,13 @@ class MockTeleop(Teleoperator):
|
||||
def is_connected(self) -> bool:
|
||||
return self._is_connected
|
||||
|
||||
def connect(self) -> None:
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self._is_connected = True
|
||||
if calibrate:
|
||||
self.calibrate()
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import time
|
||||
|
||||
from lerobot.calibrate import CalibrateConfig, calibrate
|
||||
from lerobot.record import DatasetRecordConfig, RecordConfig, record
|
||||
from lerobot.replay import DatasetReplayConfig, ReplayConfig, replay
|
||||
from lerobot.teleoperate import TeleoperateConfig, teleoperate
|
||||
@@ -8,6 +9,12 @@ from tests.mocks.mock_robot import MockRobotConfig
|
||||
from tests.mocks.mock_teleop import MockTeleopConfig
|
||||
|
||||
|
||||
def test_calibrate():
|
||||
robot_cfg = MockRobotConfig()
|
||||
cfg = CalibrateConfig(device=robot_cfg)
|
||||
calibrate(cfg)
|
||||
|
||||
|
||||
def test_teleoperate():
|
||||
robot_cfg = MockRobotConfig()
|
||||
teleop_cfg = MockTeleopConfig()
|
||||
|
||||
Reference in New Issue
Block a user