oopsie-data-tools 0.2.0__py3-none-any.whl

oopsie_data_tools/test/test_robot_setup.py

@@ -0,0 +1,244 @@
+"""Tests for robot profile YAML loading."""
+
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import yaml
+
+from oopsie_data_tools.utils.paths import PROFILES_DIR_NAME, repo_config_dir
+from oopsie_data_tools.utils.robot_profile.robot_profile import (
+    RobotProfile,
+    load_robot_profile,
+)
+
+# The shipped example profiles, addressed directly. This used to go through
+# openpi_example_robot_profile_path(), which resolved via the cwd-sensitive profile lookup
+# chain and so failed whenever $OOPSIE_ROBOT_PROFILES_DIR was set or pytest ran from a
+# directory containing ./robot_profiles.
+BUNDLED_PROFILES = repo_config_dir() / PROFILES_DIR_NAME if repo_config_dir() else None
+
+VALID_PROFILE = {
+    "policy_name": "test_policy",
+    "robot_name": "test_robot",
+    "is_biarm": False,
+    "uses_mobile_base": False,
+    "gripper_name": "test_gripper",
+    "control_freq": 10,
+    "camera_names": ["cam0"],
+    "robot_state_keys": ["cartesian_position", "gripper_position"],
+    "action_space": ["cartesian_position", "gripper_position"],
+}
+
+
+def _write_profile(data: dict, tmp_dir: str) -> Path:
+    p = Path(tmp_dir) / "profile.yaml"
+    p.write_text(yaml.dump(data))
+    return p
+
+
+class TestRobotProfileValid(unittest.TestCase):
+    def test_load_openpi_example_yaml(self) -> None:
+        self.assertIsNotNone(BUNDLED_PROFILES, "these tests run from a checkout")
+        path = BUNDLED_PROFILES / "openpi_example_robot_profile.yaml"
+        self.assertTrue(path.is_file(), msg=f"Missing {path}")
+        profile = load_robot_profile(path)
+        self.assertIsInstance(profile, RobotProfile)
+        self.assertIsInstance(profile.action_space, list)
+        self.assertIn("joint_velocity", profile.action_space)
+        self.assertEqual(len(profile.action_joint_names or []), 7)
+
+    def test_load_minimal_valid_profile(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(VALID_PROFILE, tmp)
+            profile = load_robot_profile(path)
+            self.assertIsInstance(profile, RobotProfile)
+            self.assertFalse(profile.is_biarm)
+
+
+class TestRobotProfileFileErrors(unittest.TestCase):
+    def test_missing_file_raises(self) -> None:
+        with self.assertRaises(FileNotFoundError):
+            load_robot_profile("/nonexistent/path/profile.yaml")
+
+    def test_non_mapping_yaml_raises(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            p = Path(tmp) / "profile.yaml"
+            p.write_text("- just\n- a\n- list\n")
+            with self.assertRaises(ValueError, msg="non-mapping YAML should raise"):
+                load_robot_profile(p)
+
+
+class TestRobotProfileMissingRequiredKeys(unittest.TestCase):
+    def _missing_key_raises(self, key: str) -> None:
+        data = {k: v for k, v in VALID_PROFILE.items() if k != key}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg=f"missing '{key}' should raise"):
+                load_robot_profile(path)
+
+    def test_missing_policy_name(self) -> None:
+        self._missing_key_raises("policy_name")
+
+    def test_missing_robot_name(self) -> None:
+        self._missing_key_raises("robot_name")
+
+    def test_missing_gripper_name(self) -> None:
+        self._missing_key_raises("gripper_name")
+
+    def test_missing_control_freq(self) -> None:
+        self._missing_key_raises("control_freq")
+
+    def test_missing_camera_names(self) -> None:
+        self._missing_key_raises("camera_names")
+
+    def test_missing_robot_state_keys(self) -> None:
+        self._missing_key_raises("robot_state_keys")
+
+    def test_missing_action_space(self) -> None:
+        self._missing_key_raises("action_space")
+
+
+class TestRobotProfileInvalidRobotStateKeys(unittest.TestCase):
+    def test_cartesian_action_without_cartesian_position_state_key(self) -> None:
+        data = {**VALID_PROFILE, "robot_state_keys": ["gripper_position"]}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaisesRegex(
+                ValueError,
+                "cartesian_position",
+                msg="Cartesian control without cartesian_position state should raise",
+            ):
+                load_robot_profile(path)
+
+    def test_missing_gripper_position_state_key(self) -> None:
+        data = {**VALID_PROFILE, "robot_state_keys": ["cartesian_position"]}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="missing gripper_position in robot_state_keys"):
+                load_robot_profile(path)
+
+    def test_joint_velocity_action_without_joint_position_state_key(self) -> None:
+        data = {
+            **VALID_PROFILE,
+            "robot_state_keys": ["gripper_position"],
+            "action_space": ["joint_velocity", "gripper_position"],
+            "action_joint_names": ["j0", "j1"],
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaisesRegex(
+                ValueError,
+                "joint_position",
+                msg="Joint velocity control without joint_position state should raise",
+            ):
+                load_robot_profile(path)
+
+    def test_joint_velocity_action_with_joint_position_state_key(self) -> None:
+        data = {
+            **VALID_PROFILE,
+            "robot_state_keys": ["joint_position", "gripper_position"],
+            "robot_state_joint_names": ["j0", "j1"],
+            "action_space": ["joint_velocity", "gripper_position"],
+            "action_joint_names": ["j0", "j1"],
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            profile = load_robot_profile(_write_profile(data, tmp))
+            self.assertEqual(
+                profile.robot_state_keys,
+                ["joint_position", "gripper_position"],
+            )
+
+    def test_joint_position_state_without_joint_names(self) -> None:
+        for joint_names in (None, []):
+            with self.subTest(robot_state_joint_names=joint_names):
+                data = {
+                    **VALID_PROFILE,
+                    "robot_state_keys": [
+                        "joint_position",
+                        "cartesian_position",
+                        "gripper_position",
+                    ],
+                }
+                if joint_names is not None:
+                    data["robot_state_joint_names"] = joint_names
+
+                with tempfile.TemporaryDirectory() as tmp:
+                    path = _write_profile(data, tmp)
+                    with self.assertRaisesRegex(
+                        ValueError,
+                        "robot_state_joint_names is required",
+                    ):
+                        load_robot_profile(path)
+
+    def test_empty_robot_state_keys(self) -> None:
+        data = {**VALID_PROFILE, "robot_state_keys": []}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError):
+                load_robot_profile(path)
+
+
+class TestRobotProfileInvalidActionSpace(unittest.TestCase):
+    def test_completely_invalid_action_key(self) -> None:
+        data = {**VALID_PROFILE, "action_space": ["invalid_key", "gripper_position"]}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="invalid action key should raise"):
+                load_robot_profile(path)
+
+    def test_no_allowed_action_key(self) -> None:
+        data = {**VALID_PROFILE, "action_space": ["gripper_position"]}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="action_space with no ALLOWED_ACTION_SPACES key"):
+                load_robot_profile(path)
+
+    def test_missing_gripper_in_action_space(self) -> None:
+        data = {**VALID_PROFILE, "action_space": ["cartesian_position"]}
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="no gripper key in action_space"):
+                load_robot_profile(path)
+
+    def test_joint_space_without_joint_names(self) -> None:
+        data = {
+            **VALID_PROFILE,
+            "robot_state_keys": ["joint_position", "gripper_position"],
+            "robot_state_joint_names": ["j0", "j1"],
+            "action_space": ["joint_position", "gripper_position"],
+            # action_joint_names intentionally omitted
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="joint_position without action_joint_names"):
+                load_robot_profile(path)
+
+    def test_joint_velocity_without_joint_names(self) -> None:
+        data = {
+            **VALID_PROFILE,
+            "robot_state_keys": ["joint_position", "gripper_position"],
+            "robot_state_joint_names": ["j0", "j1"],
+            "action_space": ["joint_velocity", "gripper_position"],
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="joint_velocity without action_joint_names"):
+                load_robot_profile(path)
+
+    def test_mobile_base_without_base_action_key(self) -> None:
+        data = {
+            **VALID_PROFILE,
+            "uses_mobile_base": True,
+            # action_space has no base_velocity or base_position
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            path = _write_profile(data, tmp)
+            with self.assertRaises(ValueError, msg="mobile base without base action key"):
+                load_robot_profile(path)
+
+
+if __name__ == "__main__":
+    unittest.main()

oopsie_data_tools/test/test_rollout_annotator.py

@@ -0,0 +1,134 @@
+"""Tests for WebRolloutAnnotator: construction and recording delegation.
+
+Server-dependent methods (start, wait_for_task, finish_rollout) require a
+running Flask process and are not covered here.
+"""
+
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+from oopsie_data_tools.annotation_tool.episode_recorder import EpisodeRecorder
+from oopsie_data_tools.annotation_tool.rollout_annotator import WebRolloutAnnotator
+from oopsie_data_tools.utils.robot_profile.robot_profile import RobotProfile
+
+
+def _profile(**overrides) -> RobotProfile:
+    defaults = dict(
+        policy_name="test_policy",
+        robot_name="test_robot",
+        is_biarm=False,
+        uses_mobile_base=False,
+        gripper_name="test_gripper",
+        control_freq=10,
+        camera_names=["left"],
+        robot_state_keys=["joint_position", "gripper_position"],
+        robot_state_joint_names=["j1", "j2", "j3", "j4", "j5", "j6", "j7"],
+        action_space=["joint_velocity", "gripper_position"],
+        action_joint_names=["j1", "j2", "j3", "j4", "j5", "j6", "j7"],
+    )
+    defaults.update(overrides)
+    return RobotProfile(**defaults)
+
+
+def _obs(profile: RobotProfile) -> dict:
+    return {
+        "robot_state": {
+            "joint_position": np.zeros(7, dtype=np.float32),
+            "gripper_position": np.zeros(1, dtype=np.float32),
+        },
+        "image_observation": {
+            cam: np.zeros((64, 64, 3), dtype=np.uint8) for cam in profile.camera_names
+        },
+    }
+
+
+def _action(profile: RobotProfile) -> dict:
+    sizes = {
+        "joint_velocity": 7,
+        "joint_position": 7,
+        "gripper_position": 1,
+        "gripper_velocity": 1,
+    }
+    return {
+        k: np.zeros(sizes.get(k, 1), dtype=np.float32) for k in profile.action_space
+    }
+
+
+class TestWebRolloutAnnotatorConstruction(unittest.TestCase):
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.profile = _profile()
+        self.annotator = WebRolloutAnnotator(
+            robot_profile=self.profile,
+            data_root_dir=Path(self._tmp.name),
+            operator_name="  tester  ",
+            annotator_name="test_annotator",
+            wait_for_annotation=False,
+            open_browser=False,
+        )
+
+    def tearDown(self):
+        self.annotator.stop()
+        self._tmp.cleanup()
+
+    def test_operator_name_stripped(self):
+        self.assertEqual(self.annotator.operator_name, "tester")
+
+    def test_wait_for_annotation_stored(self):
+        self.assertFalse(self.annotator.wait_for_annotation)
+
+    def test_internal_recorder_is_episode_recorder(self):
+        self.assertIsInstance(self.annotator._active_recorder, EpisodeRecorder)
+
+    def test_data_root_dir_resolved(self):
+        self.assertTrue(self.annotator.data_root_dir.is_absolute())
+
+    def test_stop_is_safe_without_start(self):
+        self.annotator.stop()  # _proc is None — must not raise
+
+
+class TestWebRolloutAnnotatorRecording(unittest.TestCase):
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.profile = _profile()
+        self.annotator = WebRolloutAnnotator(
+            robot_profile=self.profile,
+            data_root_dir=Path(self._tmp.name),
+            operator_name="tester",
+            annotator_name="test_annotator",
+            wait_for_annotation=False,
+            open_browser=False,
+        )
+        self.annotator.reset_episode_recorder()
+
+    def tearDown(self):
+        self.annotator.stop()
+        self._tmp.cleanup()
+
+    def test_record_step_delegates_to_recorder(self):
+        self.annotator.record_step(_obs(self.profile), _action(self.profile))
+        self.assertEqual(self.annotator._active_recorder.num_steps, 1)
+
+    def test_multiple_record_steps(self):
+        for _ in range(3):
+            self.annotator.record_step(_obs(self.profile), _action(self.profile))
+        self.assertEqual(self.annotator._active_recorder.num_steps, 3)
+
+    def test_reset_clears_recorder(self):
+        self.annotator.record_step(_obs(self.profile), _action(self.profile))
+        self.annotator.reset_episode_recorder()
+        self.assertEqual(self.annotator._active_recorder.num_steps, 0)
+
+    def test_record_step_propagates_validation_error(self):
+        bad_action = {"joint_velocity": np.zeros(7)}  # missing gripper_position
+        with self.assertRaises(ValueError):
+            self.annotator.record_step(_obs(self.profile), bad_action)
+
+
+if __name__ == "__main__":
+    unittest.main()

oopsie_data_tools/test/test_rotation_utils.py

@@ -0,0 +1,126 @@
+"""Tests for orientation-representation conversion.
+
+This module had no coverage at all, which is how the ``matrix`` branch shipped broken:
+it handed scipy a flat 9-element slice without reshaping, so every profile declaring
+``orientation_representation: matrix`` raised at the first recorded step.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from scipy.spatial.transform import Rotation as R
+
+from oopsie_data_tools.utils.robot_profile.rotation_utils import ActionQuatConversion, RotOption
+
+POSITION = np.array([0.1, -0.2, 0.3])
+# An arbitrary rotation, expressed once and then re-encoded per representation below.
+REFERENCE = R.from_euler("xyz", [0.3, -0.5, 1.1])
+
+
+def _assert_same_rotation(quat: np.ndarray, expected: R) -> None:
+    """Compare as matrices: q and -q are the same rotation, and Rotation.approx_equal
+    does not exist on the scipy that ships with our Python 3.8 floor."""
+    np.testing.assert_allclose(R.from_quat(quat).as_matrix(), expected.as_matrix(), atol=1e-8)
+
+
+def _encode(option: RotOption) -> np.ndarray:
+    if option == RotOption.QUAT:
+        return REFERENCE.as_quat()
+    if option == RotOption.MATRIX:
+        return REFERENCE.as_matrix().reshape(9)
+    if option == RotOption.ROT6D:
+        # First two columns of the rotation matrix, flattened row-wise as the converter reads it.
+        return REFERENCE.as_matrix()[:, :2].T.reshape(6)
+    if option == RotOption.ROTVEC:
+        return REFERENCE.as_rotvec()
+    return REFERENCE.as_euler(option.name)
+
+
+@pytest.mark.parametrize(
+    "option",
+    [
+        RotOption.QUAT,
+        RotOption.MATRIX,
+        RotOption.ROT6D,
+        RotOption.ROTVEC,
+        RotOption.xyz,
+        RotOption.zyx,
+        RotOption.XYZ,
+    ],
+)
+def test_every_representation_round_trips_to_the_same_rotation(option):
+    action = np.concatenate([POSITION, _encode(option)])
+
+    result = ActionQuatConversion(option).convert_position(action)
+
+    assert result.shape == (7,), "position + scalar-last quaternion"
+    np.testing.assert_allclose(result[:3], POSITION)
+    _assert_same_rotation(result[3:], REFERENCE)
+
+
+def test_matrix_is_accepted_flat():
+    """The regression: a flat 3x3 in the action vector must not reach scipy unreshaped."""
+    action = np.concatenate([POSITION, np.eye(3).reshape(9)])
+
+    result = ActionQuatConversion(RotOption.MATRIX).convert_position(action)
+
+    # 1e-9 rather than exact: the value comes out of scipy's matrix-to-quaternion routine,
+    # and pinning it to machine epsilon would make this brittle across scipy/BLAS builds
+    # for no gain — the point is that it converts at all, and to the identity rotation.
+    np.testing.assert_allclose(result, [*POSITION, 0.0, 0.0, 0.0, 1.0], atol=1e-9)
+
+
+def test_biarm_converts_each_arm_independently():
+    left = R.from_euler("xyz", [0.2, 0.0, 0.0])
+    right = R.from_euler("xyz", [0.0, 0.7, 0.0])
+    action = np.concatenate(
+        [POSITION, left.as_euler("xyz"), -POSITION, right.as_euler("xyz")]
+    )
+
+    result = ActionQuatConversion(RotOption.xyz, is_biarm=True).convert_position(action)
+
+    assert result.shape == (14,)
+    np.testing.assert_allclose(result[:3], POSITION)
+    np.testing.assert_allclose(result[7:10], -POSITION)
+    _assert_same_rotation(result[3:7], left)
+    _assert_same_rotation(result[10:14], right)
+
+
+def test_wrong_orientation_width_names_the_profile_field():
+    """A profile/policy mismatch should say so, not surface a scipy internal error."""
+    action = np.concatenate([POSITION, np.zeros(3)])  # euler-sized, but declared as quat
+
+    with pytest.raises(ValueError, match="QUAT orientation expects 4 value"):
+        ActionQuatConversion(RotOption.QUAT).convert_position(action)
+
+
+def test_biarm_rejects_an_odd_action_length():
+    with pytest.raises(ValueError, match="even length"):
+        ActionQuatConversion(RotOption.QUAT, is_biarm=True).convert_position(np.zeros(13))
+
+
+@pytest.mark.parametrize(
+    ("text", "expected"),
+    [
+        ("quat", RotOption.QUAT),
+        ("matrix", RotOption.MATRIX),
+        ("rot6d", RotOption.ROT6D),
+        ("rotvec", RotOption.ROTVEC),
+        ("euler_xyz", RotOption.xyz),
+        ("euler_XYZ", RotOption.XYZ),
+    ],
+)
+def test_from_string_accepts_every_documented_option(text, expected):
+    """Every option named in configs/robot_profiles/template.yaml must parse."""
+    assert RotOption.from_string(text) == expected
+
+
+def test_from_string_rejects_an_unknown_euler_order():
+    with pytest.raises(ValueError, match="Unsupported Euler format"):
+        RotOption.from_string("euler_abc")
+
+
+def test_from_string_rejects_an_unknown_representation():
+    with pytest.raises(ValueError, match="Unsupported rotation option"):
+        RotOption.from_string("axis_angle")