Add clamping for SO3 and SE3 (#96)

* PoseConstraintTask implementation

* make clamping a function for SO3 and SE3, remove PoseConstraintTask

* docstring format fix, typo

Author: adlarkin

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Added
+
+- `clamp` method for SO3 and SE3.
+
 ### Changed
 
 - `solve_ik` now raises a new exception, `NoSolutionFound`, when the QP solver fails to find a solution.

examples/arm_hand_xarm_leap.py

@@ -99,6 +99,21 @@ def construct_model() -> mujoco.MjModel:
         configuration.update(data.qpos)
         posture_task.set_target_from_configuration(configuration)
 
+        # Define a rectangular region centered about the origin of the initial
+        # end-effector pose. This can be used to clamp the end_effector_task target
+        # to ensure that the end-effector stays within this region.
+        # Roll and pitch limits are also defined to keep the hand pointed downward.
+        T_eef_init = configuration.get_transform_frame_to_world(
+            "attachment_site", "site"
+        )
+        translation_init = T_eef_init.translation()
+        x_limits = (translation_init[0] - 0.1, translation_init[0] + 0.1)
+        y_limits = (translation_init[1] - 0.1, translation_init[1] + 0.1)
+        z_limits = (translation_init[2] - 0.1, translation_init[2] + 0.1)
+        rpy_init = T_eef_init.rotation().as_rpy_radians()
+        roll_limits = (rpy_init.roll - 0.2, rpy_init.roll + 0.2)
+        pitch_limits = (rpy_init.pitch - 0.2, rpy_init.pitch + 0.2)
+
         # Initialize the mocap target at the end-effector site.
         mink.move_mocap_to_frame(model, data, "target", "attachment_site", "site")
         for finger in fingers:
@@ -112,9 +127,16 @@ def construct_model() -> mujoco.MjModel:
 
         rate = RateLimiter(frequency=200.0, warn=False)
         while viewer.is_running():
-            # Update kuka end-effector task.
+            # Clamp the target within the limits and then update the end-effector task.
             T_wt = mink.SE3.from_mocap_name(model, data, "target")
-            end_effector_task.set_target(T_wt)
+            T_wt_clamped = T_wt.clamp(
+                x_translation=x_limits,
+                y_translation=y_limits,
+                z_translation=z_limits,
+                roll_radians=roll_limits,
+                pitch_radians=pitch_limits,
+            )
+            end_effector_task.set_target(T_wt_clamped)
 
             # Update finger tasks.
             for finger, task in zip(fingers, finger_tasks):

mink/lie/se3.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Tuple
 
 import mujoco
 import numpy as np
@@ -221,6 +222,36 @@ def ljacinv(cls, other: np.ndarray) -> np.ndarray:
         ljacinv_se3[3:, 3:] = ljacinv_so3
         return ljacinv_se3
 
+    def clamp(
+        self,
+        x_translation: Tuple[float, float] = (-np.inf, np.inf),
+        y_translation: Tuple[float, float] = (-np.inf, np.inf),
+        z_translation: Tuple[float, float] = (-np.inf, np.inf),
+        roll_radians: Tuple[float, float] = (-np.inf, np.inf),
+        pitch_radians: Tuple[float, float] = (-np.inf, np.inf),
+        yaw_radians: Tuple[float, float] = (-np.inf, np.inf),
+    ) -> SE3:
+        """Clamp a SE3 within translation and RPY limits.
+
+        Args:
+            x_translation: The (lower, upper) limits for translation along the x axis.
+            y_translation: The (lower, upper) limits for translation along the y axis.
+            z_translation: The (lower, upper) limits for translation along the z axis.
+            roll_radians: The (lower, upper) limits for the roll.
+            pitch_radians: The (lower, upper) limits for the pitch.
+            yaw_radians: The (lower, upper) limits for the yaw.
+
+        Returns:
+            A SE3 within the translation and RPY limits.
+        """
+        translation_limits = np.array([x_translation, y_translation, z_translation])
+        return SE3.from_rotation_and_translation(
+            rotation=self.rotation().clamp(roll_radians, pitch_radians, yaw_radians),
+            translation=np.clip(
+                self.translation(), translation_limits[:, 0], translation_limits[:, 1]
+            ),
+        )
+
 
 # Eqn 180.
 def _getQ(c) -> np.ndarray:

mink/lie/so3.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import NamedTuple
+from typing import NamedTuple, Tuple
 
 import mujoco
 import numpy as np
@@ -260,3 +260,25 @@ def ljacinv(cls, other: np.ndarray) -> np.ndarray:
         ljacinv[1, 1] += 1.0
         ljacinv[2, 2] += 1.0
         return ljacinv
+
+    def clamp(
+        self,
+        roll_radians: Tuple[float, float] = (-np.inf, np.inf),
+        pitch_radians: Tuple[float, float] = (-np.inf, np.inf),
+        yaw_radians: Tuple[float, float] = (-np.inf, np.inf),
+    ) -> SO3:
+        """Clamp a SO3 within RPY limits.
+
+        Args:
+            roll_radians: The (lower, upper) limits for the roll.
+            pitch_radians: The (lower, upper) limits for the pitch.
+            yaw_radians: The (lower, upper) limits for the yaw.
+
+        Returns:
+            A SO3 within the RPY limits.
+        """
+        return SO3.from_rpy_radians(
+            roll=np.clip(self.compute_roll_radians(), *roll_radians),
+            pitch=np.clip(self.compute_pitch_radians(), *pitch_radians),
+            yaw=np.clip(self.compute_yaw_radians(), *yaw_radians),
+        )

mink/tasks/frame_task.py

@@ -20,7 +20,7 @@ class FrameTask(Task):
         frame_name: Name of the frame to regulate, typically the name of body, geom
             or site in the robot model.
         frame_type: The frame type: `body`, `geom` or `site`.
-        transform_frame_to_world: Target pose of the frame in the world frame.
+        transform_target_to_world: Target pose of the frame in the world frame.
 
     Example:
 

tests/test_lie_operations.py

@@ -148,6 +148,34 @@ def test_so3_apply_throws_assertion_error_if_wrong_shape(self):
         with self.assertRaises(AssertionError):
             rot.apply(vec)
 
+    def test_so3_clamp(self):
+        # Clamping with the default RPY limits (+- infinity) means the SO3 should remain unchanged.
+        rot = SO3.from_rpy_radians(roll=np.pi, pitch=0, yaw=-np.pi)
+        self.assertEqual(rot, rot.clamp())
+
+        original_rpy = rot.as_rpy_radians()
+
+        # Test clamping the roll.
+        clamped_rot = rot.clamp(roll_radians=(0, 1))
+        clamped_rpy = clamped_rot.as_rpy_radians()
+        self.assertAlmostEqual(clamped_rpy.roll, 1)
+        self.assertAlmostEqual(clamped_rpy.pitch, original_rpy.pitch)
+        self.assertAlmostEqual(clamped_rpy.yaw, original_rpy.yaw)
+
+        # Test clamping the pitch.
+        clamped_rot = rot.clamp(pitch_radians=(1, 2))
+        clamped_rpy = clamped_rot.as_rpy_radians()
+        self.assertAlmostEqual(clamped_rpy.roll, original_rpy.roll)
+        self.assertAlmostEqual(clamped_rpy.pitch, 1)
+        self.assertAlmostEqual(clamped_rpy.yaw, original_rpy.yaw)
+
+        # Test clamping the yaw.
+        clamped_rot = rot.clamp(yaw_radians=(np.pi, 2 * np.pi))
+        clamped_rpy = clamped_rot.as_rpy_radians()
+        self.assertAlmostEqual(clamped_rpy.roll, original_rpy.roll)
+        self.assertAlmostEqual(clamped_rpy.pitch, original_rpy.pitch)
+        self.assertAlmostEqual(clamped_rpy.yaw, np.pi)
+
     # SE3.
 
     def test_se3_equality(self):
@@ -257,6 +285,55 @@ def test_se3_interpolate(self):
             start.interpolate(end, alpha=2.0)
         self.assertIn(expected_error_message, str(cm.exception))
 
+    def test_se3_clamp(self):
+        T = SE3.from_rotation_and_translation(
+            rotation=SO3.from_rpy_radians(roll=np.pi, pitch=0, yaw=-np.pi),
+            translation=np.array([0, 1, 2]),
+        )
+        original_rpy = T.rotation().as_rpy_radians()
+
+        # Clamping with the default limits (+- infinity) means the SE3 should remain unchanged.
+        self.assertEqual(T, T.clamp())
+
+        # Test clamping the x translation.
+        clamped_T = T.clamp(x_translation=(1, 2))
+        np.testing.assert_allclose(clamped_T.translation(), np.array([1, 1, 2]))
+        self.assertEqual(clamped_T.rotation(), T.rotation())
+
+        # Test clamping the y translation.
+        clamped_T = T.clamp(y_translation=(-1, 0))
+        np.testing.assert_allclose(clamped_T.translation(), np.array([0, 0, 2]))
+        self.assertEqual(clamped_T.rotation(), T.rotation())
+
+        # Test clamping the z translation.
+        clamped_T = T.clamp(z_translation=(5, 10))
+        np.testing.assert_allclose(clamped_T.translation(), np.array([0, 1, 5]))
+        self.assertEqual(clamped_T.rotation(), T.rotation())
+
+        # Test clamping the roll.
+        clamped_T = T.clamp(roll_radians=(0, 1))
+        clamped_rpy = clamped_T.rotation().as_rpy_radians()
+        np.testing.assert_equal(clamped_T.translation(), T.translation())
+        self.assertAlmostEqual(clamped_rpy.roll, 1)
+        self.assertAlmostEqual(clamped_rpy.pitch, original_rpy.pitch)
+        self.assertAlmostEqual(clamped_rpy.yaw, original_rpy.yaw)
+
+        # Test clamping the pitch.
+        clamped_T = T.clamp(pitch_radians=(1, 2))
+        clamped_rpy = clamped_T.rotation().as_rpy_radians()
+        np.testing.assert_equal(clamped_T.translation(), T.translation())
+        self.assertAlmostEqual(clamped_rpy.roll, original_rpy.roll)
+        self.assertAlmostEqual(clamped_rpy.pitch, 1)
+        self.assertAlmostEqual(clamped_rpy.yaw, original_rpy.yaw)
+
+        # Test clamping the yaw.
+        clamped_T = T.clamp(yaw_radians=(np.pi, 2 * np.pi))
+        clamped_rpy = clamped_T.rotation().as_rpy_radians()
+        np.testing.assert_equal(clamped_T.translation(), T.translation())
+        self.assertAlmostEqual(clamped_rpy.roll, original_rpy.roll)
+        self.assertAlmostEqual(clamped_rpy.pitch, original_rpy.pitch)
+        self.assertAlmostEqual(clamped_rpy.yaw, np.pi)
+
 
 if __name__ == "__main__":
     absltest.main()