move mock_calibration_dir in utils

This commit is contained in:
Remi Cadene
2024-10-24 11:45:02 +02:00
parent 8d624643aa
commit 3368e8c4d7
3 changed files with 31 additions and 28 deletions

View File

@@ -23,7 +23,6 @@ pytest -sx 'tests/test_control_robot.py::test_teleoperate[aloha-True]'
``` ```
""" """
import json
import multiprocessing import multiprocessing
from pathlib import Path from pathlib import Path
from unittest.mock import patch from unittest.mock import patch
@@ -37,32 +36,7 @@ from lerobot.common.utils.utils import init_hydra_config
from lerobot.scripts.control_robot import calibrate, record, replay, teleoperate from lerobot.scripts.control_robot import calibrate, record, replay, teleoperate
from lerobot.scripts.train import make_optimizer_and_scheduler from lerobot.scripts.train import make_optimizer_and_scheduler
from tests.test_robots import make_robot from tests.test_robots import make_robot
from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, TEST_ROBOT_TYPES, require_robot from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, TEST_ROBOT_TYPES, mock_calibration_dir, require_robot
def mock_calibration_dir(calibration_dir):
# calibration file produced with Moss v1, but works with Koch, Koch bimanual and SO-100
example_calib = {
"homing_offset": [-1416, -845, 2130, 2872, 1950, -2211],
"drive_mode": [0, 0, 1, 1, 1, 0],
"start_pos": [1442, 843, 2166, 2849, 1988, 1835],
"end_pos": [2440, 1869, -1106, -1848, -926, 3235],
"calib_mode": ["DEGREE", "DEGREE", "DEGREE", "DEGREE", "DEGREE", "LINEAR"],
"motor_names": ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"],
}
calibration_dir.mkdir(parents=True, exist_ok=True)
with open(calibration_dir / "main_follower.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "main_leader.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "left_follower.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "left_leader.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "right_follower.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "right_leader.json", "w") as f:
json.dump(example_calib, f)
@pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES) @pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES)

View File

@@ -30,7 +30,7 @@ import torch
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
from tests.utils import TEST_ROBOT_TYPES, make_robot, require_robot from tests.utils import TEST_ROBOT_TYPES, make_robot, mock_calibration_dir, require_robot
@pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES) @pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES)
@@ -54,6 +54,7 @@ def test_robot(tmpdir, request, robot_type, mock):
tmpdir = Path(tmpdir) tmpdir = Path(tmpdir)
calibration_dir = tmpdir / robot_type calibration_dir = tmpdir / robot_type
overrides_calibration_dir = [f"calibration_dir={calibration_dir}"] overrides_calibration_dir = [f"calibration_dir={calibration_dir}"]
mock_calibration_dir(calibration_dir)
robot_kwargs["calibration_dir"] = calibration_dir robot_kwargs["calibration_dir"] = calibration_dir
# Test connecting without devices raises an error # Test connecting without devices raises an error

View File

@@ -13,10 +13,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import json
import os import os
import platform import platform
from copy import copy from copy import copy
from functools import wraps from functools import wraps
from pathlib import Path
import pytest import pytest
import torch import torch
@@ -281,6 +283,32 @@ def require_motor(func):
return wrapper return wrapper
def mock_calibration_dir(calibration_dir):
# TODO(rcadene): remove this hack
# calibration file produced with Moss v1, but works with Koch, Koch bimanual and SO-100
example_calib = {
"homing_offset": [-1416, -845, 2130, 2872, 1950, -2211],
"drive_mode": [0, 0, 1, 1, 1, 0],
"start_pos": [1442, 843, 2166, 2849, 1988, 1835],
"end_pos": [2440, 1869, -1106, -1848, -926, 3235],
"calib_mode": ["DEGREE", "DEGREE", "DEGREE", "DEGREE", "DEGREE", "LINEAR"],
"motor_names": ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"],
}
Path(str(calibration_dir)).mkdir(parents=True, exist_ok=True)
with open(calibration_dir / "main_follower.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "main_leader.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "left_follower.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "left_leader.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "right_follower.json", "w") as f:
json.dump(example_calib, f)
with open(calibration_dir / "right_leader.json", "w") as f:
json.dump(example_calib, f)
def make_robot(robot_type: str, overrides: list[str] | None = None, mock=False) -> Robot: def make_robot(robot_type: str, overrides: list[str] | None = None, mock=False) -> Robot:
if mock: if mock:
overrides = [] if overrides is None else copy(overrides) overrides = [] if overrides is None else copy(overrides)