import dataclasses as dc from typing import Any import unittest import numpy as np from pydrake.math import RigidTransform, RollPitchYaw, RotationMatrix from robot_gym.multiarm_spaces import ( CURRENT_VERSION, CameraDepthImage, CameraImageSet, CameraLabelImage, CameraRgbImage, MultiarmObservation, PosesAndGrippers, PosesAndGrippersActualAndDesired, RestorePosesAndGrippersConfig, ) from robot_gym.multiarm_spaces_conversions import ( dp_vector_to_multiarm_action, matrix_to_rotation_6d, multiarm_action_to_dp_vector, multiarm_observation_to_dp_dict, rotation_6d_to_matrix, ) def rpy_deg(rpy_deg): """Converts RPY in degress to a RotationMatrix.""" return RotationMatrix(RollPitchYaw(np.deg2rad(rpy_deg))) def xyz_rpy(xyz, rpy): """Shorthand to create an isometry from XYZ and RPY.""" return RigidTransform(R=RotationMatrix(rpy=RollPitchYaw(rpy)), p=xyz) def xyz_rpy_deg(xyz, rpy_deg): return xyz_rpy(xyz, np.deg2rad(rpy_deg)) def assert_equal_recursive( test: unittest.TestCase, a: Any, b: Any, np_tol=0.0 ): """Recursively asserts equality between two objects, a and b.""" def recurse(a, b, path): T = type(a) test.assertEqual(type(b), T, path) if T == np.ndarray: test.assertEqual(a.shape, b.shape, path) test.assertEqual(a.dtype, b.dtype, path) np.testing.assert_allclose( a, b, err_msg=path, atol=np_tol, rtol=0.0 ) elif isinstance(a, (list, tuple)): test.assertEqual(len(a), len(b), path) pair_iter = zip(a, b, strict=True) for i, (ai, bi) in enumerate(pair_iter): recurse(ai, bi, f"{path}[{i}]") elif T == dict: test.assertEqual(len(a), len(b), path) test.assertEqual(a.keys(), b.keys(), path) for k, ai in a.items(): bi = b[k] recurse(ai, bi, f"{path}[{repr(k)}]") elif dc.is_dataclass(T): fields = dc.fields(a) test.assertEqual(dc.fields(b), fields) for field in fields: ai = getattr(a, field.name) bi = getattr(b, field.name) recurse(ai, bi, f"{path}.{field.name}") elif T == RigidTransform: recurse(a.GetAsMatrix4(), b.GetAsMatrix4(), path) else: # Assume a type that has a builtin equality operator. test.assertEqual(a, b, path) recurse(a, b, "") def make_example_obs_and_act(): h, w, c = (48, 64, 3) fake_rgb = np.zeros((h, w, c), dtype=np.uint8) fake_depth = np.zeros((h, w), dtype=np.uint16) fake_label = np.zeros((h, w), dtype=np.uint16) fake_K = np.eye(3) fake_X_TC = RigidTransform() camera_name = "my_camera" model_1 = "right::panda" gripper_1 = "right::panda_hand" model_2 = "left::panda" gripper_2 = "left::panda_hand" obs = MultiarmObservation( robot=PosesAndGrippersActualAndDesired( actual=PosesAndGrippers( poses={ model_1: xyz_rpy_deg([0.1, 0.2, 0.3], [15, 30, 45]), model_2: xyz_rpy_deg([-0.4, 0.5, -0.6], [-45, 60, -15]), }, grippers={ gripper_1: 0.02, gripper_2: 0.04, }, joint_position={ model_1: np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]), model_2: np.array([-0.1, 0.2, -0.3, 0.4, -0.5, 0.6, -0.7]), }, joint_velocity={ model_1: np.array([1, 2, 3, 4, 5, 6, 7]), model_2: np.array([-1, 2, -3, 4, -5, 6, -7]), }, joint_torque={ model_1: np.array([10, 20, 30, 40, 50, 60, 70]), model_2: np.array([-10, 20, -30, 40, -50, 60, -70]), }, joint_torque_external={ model_1: np.array([20, 40, 60, 80, 100, 120, 140]), model_2: np.array([-20, 40, -60, 80, -100, 120, -140]), }, wrench={ model_1: np.array([0.2, 0.6, 0.8, 2, 6, 8]), model_2: np.array([-0.2, 0.6, -0.8, 2, -6, 8]), }, external_wrench={ model_1: np.array([0.02, 0.06, 0.08, 0.2, 0.6, 0.8]), model_2: np.array([-0.02, 0.06, -0.08, 0.2, -0.6, 0.8]), }, timestamp_data=0.123, timestamp_sent=None, timestamp_received=0.456, ), desired=PosesAndGrippers( poses={ model_1: xyz_rpy_deg( [0.11, 0.21, 0.31], [15.1, 30.1, 45.1] ), model_2: xyz_rpy_deg( [-0.41, 0.51, -0.61], [-45.1, 60.1, -15.1] ), }, grippers={ gripper_1: 0.021, gripper_2: 0.041, }, joint_position={ model_1: np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]), model_2: np.array([-0.1, 0.2, -0.3, 0.4, -0.5, 0.6, -0.7]), }, joint_velocity=None, joint_torque=None, joint_torque_external=None, timestamp_sent=0.789, ), version=CURRENT_VERSION, ), visuo={ camera_name: CameraImageSet( rgb=CameraRgbImage(fake_rgb, fake_K, fake_X_TC), depth=CameraDepthImage(fake_depth, fake_K, fake_X_TC), label=CameraLabelImage(fake_label, fake_K, fake_X_TC), ), }, timestamp_packaged=1.234, ) act = PosesAndGrippers( poses={ "right::panda": xyz_rpy_deg([0.1, 0.2, 0.3], [15, 30, 45]), "left::panda": xyz_rpy_deg([-0.4, 0.5, -0.6], [-45, 60, -15]), }, grippers={ "right::panda_hand": 0.02, "left::panda_hand": 0.04, }, debugging_output={ "skill": ["something"], }, ) return obs, act def make_example_obs_and_obs_dict(): obs, _ = make_example_obs_and_act() rgb = obs.visuo["my_camera"].rgb.array depth = obs.visuo["my_camera"].depth.array label = obs.visuo["my_camera"].label.array obs_dict = { "my_camera": rgb, "my_camera_depth": depth, "my_camera_label": label, "robot__actual__poses__right::panda__xyz": np.array([0.1, 0.2, 0.3]), "robot__actual__poses__right::panda__rot_6d": np.array( [ 0.61237244, -0.59150635, 0.52451905, 0.61237244, 0.77451905, 0.15849365, ] ), "robot__actual__poses__left::panda__xyz": np.array([-0.4, 0.5, -0.6]), "robot__actual__poses__left::panda__rot_6d": np.array( [ 0.48296291, -0.40849365, 0.77451905, -0.12940952, 0.84150635, 0.52451905, ] ), "robot__actual__grippers__right::panda_hand": np.array([0.02]), "robot__actual__grippers__left::panda_hand": np.array([0.04]), "robot__actual__joint_position__right::panda": np.array( [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7] ), "robot__actual__joint_position__left::panda": np.array( [-0.1, 0.2, -0.3, 0.4, -0.5, 0.6, -0.7] ), "robot__actual__joint_velocity__right::panda": np.array( [1, 2, 3, 4, 5, 6, 7] ), "robot__actual__joint_velocity__left::panda": np.array( [-1, 2, -3, 4, -5, 6, -7] ), "robot__actual__joint_torque__right::panda": np.array( [10, 20, 30, 40, 50, 60, 70] ), "robot__actual__joint_torque__left::panda": np.array( [-10, 20, -30, 40, -50, 60, -70] ), "robot__actual__joint_torque_external__right::panda": np.array( [20, 40, 60, 80, 100, 120, 140] ), "robot__actual__joint_torque_external__left::panda": np.array( [-20, 40, -60, 80, -100, 120, -140] ), "robot__actual__wrench__right::panda": np.array( [0.2, 0.6, 0.8, 2.0, 6.0, 8.0] ), "robot__actual__wrench__left::panda": np.array( [-0.2, 0.6, -0.8, 2.0, -6.0, 8.0] ), "robot__actual__external_wrench__right::panda": np.array( [0.02, 0.06, 0.08, 0.2, 0.6, 0.8] ), "robot__actual__external_wrench__left::panda": np.array( [-0.02, 0.06, -0.08, 0.2, -0.6, 0.8] ), "robot__actual__timestamp_data": np.array([0.123]), "robot__actual__timestamp_received": np.array([0.456]), "robot__desired__poses__right::panda__xyz": np.array( [0.11, 0.21, 0.31] ), "robot__desired__poses__right::panda__rot_6d": np.array( [ 0.61068579, -0.59166356, 0.52630513, 0.61282122, 0.77404131, 0.1590918, ] ), "robot__desired__poses__left::panda__xyz": np.array( [-0.41, 0.51, -0.61] ), "robot__desired__poses__left::panda__rot_6d": np.array( [ 0.48127627, -0.40897299, 0.77531558, -0.1298583, 0.84146443, 0.52447539, ] ), "robot__desired__grippers__right::panda_hand": np.array([0.021]), "robot__desired__grippers__left::panda_hand": np.array([0.041]), "robot__desired__joint_position__right::panda": np.array( [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7] ), "robot__desired__joint_position__left::panda": np.array( [-0.1, 0.2, -0.3, 0.4, -0.5, 0.6, -0.7] ), "robot__desired__timestamp_sent": np.array([0.789]), "robot__version": np.array([CURRENT_VERSION]), "timestamp_packaged": np.array([1.234]), } return obs, obs_dict class Test(unittest.TestCase): def setUp(self): self.maxDiff = None # Ensure we get good assertion errors here. def test_rotation_6d(self): R = rpy_deg([15, 30, 45]).matrix() d6 = matrix_to_rotation_6d(R) self.assertEqual(d6.shape, (6,)) R_again = rotation_6d_to_matrix(d6) np.testing.assert_allclose(R, R_again, atol=1e-15, rtol=0) def test_rotation_6d_orthonormality_regression_test(self): # This caused orthonormality errors (#13969). d6 = np.array( [ 0.6298105716705322, 0.035225264728069305, -0.7759491205215454, 0.22841385006904602, -0.9632018208503723, 0.14166943728923798, ], dtype=np.float32, ) R = rotation_6d_to_matrix(d6) # This is the tolerance used by Drake's # RotationMatrix::kInternalToleranceForOrthonormality. tolerance = 128 * np.finfo(np.float64).eps orthonormality_measure = np.max(np.abs(R @ R.T - np.eye(3))) self.assertLess(orthonormality_measure, tolerance) def test_observation_to_dp_dict(self): obs, obs_dict_expected = make_example_obs_and_obs_dict() obs_dict = multiarm_observation_to_dp_dict(obs) assert_equal_recursive(self, obs_dict, obs_dict_expected, np_tol=1e-8) def test_action_to_and_from_dp_vector(self): _, act = make_example_obs_and_act() vector = multiarm_action_to_dp_vector(act) config = RestorePosesAndGrippersConfig.make_default() act_again = dp_vector_to_multiarm_action(config, vector) # N.B. `DeepDiff` does not seem to respect more precise tolerances; # DeepDiff(*, significant_digits=16, number_format_notation="e") does # not trigger an expected failure. act.debugging_output = None assert_equal_recursive(self, act, act_again, np_tol=2e-16) if __name__ == "__main__": unittest.main()