Skip to content

Commit

Permalink
fix unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Cadene committed Oct 24, 2024
1 parent 5d64ba5 commit 4d03ece
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 2 deletions.
2 changes: 1 addition & 1 deletion lerobot/common/robot_devices/robots/manipulator.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def __setattr__(self, prop: str, val):
super().__setattr__(prop, val)

def __post_init__(self):
if self.robot_type not in ["koch", "aloha", "so100", "moss"]:
if self.robot_type not in ["koch", "koch_bimanual", "aloha", "so100", "moss"]:
raise ValueError(f"Provided robot type ({self.robot_type}) is not supported.")


Expand Down
35 changes: 34 additions & 1 deletion tests/test_control_robot.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
```
"""

import json
import multiprocessing
from pathlib import Path
from unittest.mock import patch
Expand All @@ -39,6 +40,31 @@
from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, TEST_ROBOT_TYPES, 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)
@require_robot
def test_teleoperate(tmpdir, request, robot_type, mock):
Expand All @@ -49,6 +75,7 @@ def test_teleoperate(tmpdir, request, robot_type, mock):
# and avoid writing calibration files in user .cache/calibration folder
tmpdir = Path(tmpdir)
calibration_dir = tmpdir / robot_type
mock_calibration_dir(calibration_dir)
overrides = [f"calibration_dir={calibration_dir}"]
else:
# Use the default .cache/calibration folder when mock=False
Expand Down Expand Up @@ -89,6 +116,7 @@ def test_record_without_cameras(tmpdir, request, robot_type, mock):
# Create an empty calibration directory to trigger manual calibration
# and avoid writing calibration files in user .cache/calibration folder
calibration_dir = Path(tmpdir) / robot_type
mock_calibration_dir(calibration_dir)
overrides.append(f"calibration_dir={calibration_dir}")

root = Path(tmpdir) / "data"
Expand Down Expand Up @@ -121,6 +149,7 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock):
# Create an empty calibration directory to trigger manual calibration
# and avoid writing calibration files in user .cache/calibration folder
calibration_dir = tmpdir / robot_type
mock_calibration_dir(calibration_dir)
overrides = [f"calibration_dir={calibration_dir}"]
else:
# Use the default .cache/calibration folder when mock=False or for aloha
Expand Down Expand Up @@ -159,7 +188,7 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock):
if robot_type == "aloha":
env_name = "aloha_real"
policy_name = "act_aloha_real"
elif robot_type in ["koch", "koch_bimanual"]:
elif robot_type in ["koch", "koch_bimanual", "so100", "moss"]:
env_name = "koch_real"
policy_name = "act_koch_real"
else:
Expand Down Expand Up @@ -248,6 +277,7 @@ def test_resume_record(tmpdir, request, robot_type, mock):
# Create an empty calibration directory to trigger manual calibration
# and avoid writing calibration files in user .cache/calibration folder
calibration_dir = tmpdir / robot_type
mock_calibration_dir(calibration_dir)
overrides = [f"calibration_dir={calibration_dir}"]
else:
# Use the default .cache/calibration folder when mock=False or for aloha
Expand Down Expand Up @@ -311,6 +341,7 @@ def test_record_with_event_rerecord_episode(tmpdir, request, robot_type, mock):
# Create an empty calibration directory to trigger manual calibration
# and avoid writing calibration files in user .cache/calibration folder
calibration_dir = tmpdir / robot_type
mock_calibration_dir(calibration_dir)
overrides = [f"calibration_dir={calibration_dir}"]
else:
# Use the default .cache/calibration folder when mock=False or for aloha
Expand Down Expand Up @@ -360,6 +391,7 @@ def test_record_with_event_exit_early(tmpdir, request, robot_type, mock):
# Create an empty calibration directory to trigger manual calibration
# and avoid writing calibration files in user .cache/calibration folder
calibration_dir = tmpdir / robot_type
mock_calibration_dir(calibration_dir)
overrides = [f"calibration_dir={calibration_dir}"]
else:
# Use the default .cache/calibration folder when mock=False or for aloha
Expand Down Expand Up @@ -410,6 +442,7 @@ def test_record_with_event_stop_recording(tmpdir, request, robot_type, mock, num
# Create an empty calibration directory to trigger manual calibration
# and avoid writing calibration files in user .cache/calibration folder
calibration_dir = tmpdir / robot_type
mock_calibration_dir(calibration_dir)
overrides = [f"calibration_dir={calibration_dir}"]
else:
# Use the default .cache/calibration folder when mock=False or for aloha
Expand Down

0 comments on commit 4d03ece

Please sign in to comment.