mink 0.0.2__tar.gz → 0.0.3__tar.gz

{mink-0.0.2 → mink-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mink
-Version: 0.0.2
+Version: 0.0.3
 Summary: mink: MuJoCo inverse kinematics.
 Keywords: inverse,kinematics,mujoco
 Author-email: Kevin Zakka <zakka@berkeley.edu>
@@ -43,7 +43,7 @@ Provides-Extra: test
 # mink
 
 [![Build](https://img.shields.io/github/actions/workflow/status/kevinzakka/mink/ci.yml?branch=main)](https://github.com/kevinzakka/mink/actions)
-[![Coverage Status](https://coveralls.io/repos/github/kevinzakka/mink/badge.svg)](https://coveralls.io/github/kevinzakka/mink)
+[![Coverage Status](https://coveralls.io/repos/github/kevinzakka/mink/badge.svg)](https://coveralls.io/github/kevinzakka/mink?branch=main)
 [![PyPI version](https://img.shields.io/pypi/v/mink)](https://pypi.org/project/mink/)
 ![Banner for mink](https://github.com/kevinzakka/mink/blob/assets/banner.png?raw=true)
 
@@ -79,7 +79,8 @@ mink works with a variety of robots, including:
 * Arms: [UR5e](https://github.com/kevinzakka/mink/blob/main/examples/arm_ur5e_actuators.py), [iiwa14](https://github.com/kevinzakka/mink/blob/main/examples/arm_iiwa.py), [bimanual iiwa14](https://github.com/kevinzakka/mink/blob/main/examples/dual_iiwa.py)
 * Humanoids: [Unitree G1](https://github.com/kevinzakka/mink/blob/main/examples/humanoid_g1.py)
 * Quadrupeds: [Unitree Go1](https://github.com/kevinzakka/mink/blob/main/examples/quadruped_go1.py), [Boston Dynamics Spot](https://github.com/kevinzakka/mink/blob/main/examples/quadruped_spot.py)
-* Hands: [Shadow Hand](https://github.com/kevinzakka/mink/blob/main/examples/hand_shadow.py)
+* Hands: [Shadow Hand](https://github.com/kevinzakka/mink/blob/main/examples/hand_shadow.py), [Allegro Hand](https://github.com/kevinzakka/mink/blob/main/examples/arm_hand_iiwa_allegro.py)
+* Mobile manipulators: [Stanford TidyBot](https://github.com/kevinzakka/mink/blob/main/examples/mobile_tidybot.py), [Hello Robot Stretch](https://github.com/kevinzakka/mink/blob/main/examples/mobile_stretch.py)
 
 Check out the [examples](https://github.com/kevinzakka/mink/blob/main/examples/) directory for more code.
 

{mink-0.0.2 → mink-0.0.3}/README.md

@@ -1,7 +1,7 @@
 # mink
 
 [![Build](https://img.shields.io/github/actions/workflow/status/kevinzakka/mink/ci.yml?branch=main)](https://github.com/kevinzakka/mink/actions)
-[![Coverage Status](https://coveralls.io/repos/github/kevinzakka/mink/badge.svg)](https://coveralls.io/github/kevinzakka/mink)
+[![Coverage Status](https://coveralls.io/repos/github/kevinzakka/mink/badge.svg)](https://coveralls.io/github/kevinzakka/mink?branch=main)
 [![PyPI version](https://img.shields.io/pypi/v/mink)](https://pypi.org/project/mink/)
 ![Banner for mink](https://github.com/kevinzakka/mink/blob/assets/banner.png?raw=true)
 
@@ -37,7 +37,8 @@ mink works with a variety of robots, including:
 * Arms: [UR5e](https://github.com/kevinzakka/mink/blob/main/examples/arm_ur5e_actuators.py), [iiwa14](https://github.com/kevinzakka/mink/blob/main/examples/arm_iiwa.py), [bimanual iiwa14](https://github.com/kevinzakka/mink/blob/main/examples/dual_iiwa.py)
 * Humanoids: [Unitree G1](https://github.com/kevinzakka/mink/blob/main/examples/humanoid_g1.py)
 * Quadrupeds: [Unitree Go1](https://github.com/kevinzakka/mink/blob/main/examples/quadruped_go1.py), [Boston Dynamics Spot](https://github.com/kevinzakka/mink/blob/main/examples/quadruped_spot.py)
-* Hands: [Shadow Hand](https://github.com/kevinzakka/mink/blob/main/examples/hand_shadow.py)
+* Hands: [Shadow Hand](https://github.com/kevinzakka/mink/blob/main/examples/hand_shadow.py), [Allegro Hand](https://github.com/kevinzakka/mink/blob/main/examples/arm_hand_iiwa_allegro.py)
+* Mobile manipulators: [Stanford TidyBot](https://github.com/kevinzakka/mink/blob/main/examples/mobile_tidybot.py), [Hello Robot Stretch](https://github.com/kevinzakka/mink/blob/main/examples/mobile_stretch.py)
 
 Check out the [examples](https://github.com/kevinzakka/mink/blob/main/examples/) directory for more code.
 

{mink-0.0.2 → mink-0.0.3}/mink/__init__.py

@@ -31,6 +31,7 @@ from .tasks import (
     FrameTask,
     Objective,
     PostureTask,
+    RelativeFrameTask,
     TargetNotSet,
     Task,
 )
@@ -42,8 +43,6 @@ from .utils import (
     move_mocap_to_frame,
 )
 
-__version__ = "0.0.2"
-
 __all__ = (
     "ComTask",
     "Configuration",
@@ -51,6 +50,7 @@ __all__ = (
     "solve_ik",
     "DampingTask",
     "FrameTask",
+    "RelativeFrameTask",
     "PostureTask",
     "Task",
     "Objective",

{mink-0.0.2 → mink-0.0.3}/mink/configuration.py

@@ -1,12 +1,10 @@
 """Configuration space of a robot model.
 
-The :class:`Configuration` class bundles a MuJoCo `model <https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjmodel>`__
+The :class:`Configuration` class encapsulates a MuJoCo
+`model <https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjmodel>`__
 and `data <https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjdata>`__,
-and enables easy access to kinematic quantities such as frame transforms and frame
-Jacobians.
-
-Frames are coordinate systems that can be attached to different elements of
-the robot model. mink supports frames of type `body`, `geom` and `site`.
+offering easy access to frame transforms and frame Jacobians. A frame refers to a coordinate
+system that can be attached to various parts of the robot, such as a body, geom, or site.
 """
 
 from typing import Optional
@@ -22,22 +20,17 @@ from .lie import SE3, SO3
 class Configuration:
     """Encapsulates a model and data for convenient access to kinematic quantities.
 
-    This class provides convenient methods to access and update the kinematic quantities
-    of a robot model, such as frame transforms and Jacobians. It ensures that forward
-    kinematics is computed at each time step, allowing the user to query up-to-date
-    information about the robot's state.
-
-    In this context, a frame refers to a coordinate system that can be attached to
-    different elements of the robot model. Currently supported frames include
-    `body`, `geom` and `site`.
+    This class provides methods to access and update the kinematic quantities of a robot
+    model, such as frame transforms and Jacobians. It performs forward kinematics at every
+    time step, ensuring up-to-date information about the robot's state.
 
     Key functionalities include:
 
-        - Running forward kinematics to update the state.
-        - Checking configuration limits.
-        - Computing Jacobians for different frames.
-        - Retrieving frame transforms relative to the world frame.
-        - Integrating velocities to update configurations.
+    * Running forward kinematics to update the state.
+    * Checking configuration limits.
+    * Computing Jacobians for different frames.
+    * Retrieving frame transforms relative to the world frame.
+    * Integrating velocities to update configurations.
     """
 
     def __init__(
@@ -50,7 +43,7 @@ class Configuration:
         Args:
             model: Mujoco model.
             q: Configuration to initialize from. If None, the configuration
-                is initialized to the reference configuration `qpos0`.
+                is initialized to the default configuration `qpos0`.
         """
         self.model = model
         self.data = mujoco.MjData(model)
@@ -60,7 +53,7 @@ class Configuration:
         """Run forward kinematics.
 
         Args:
-            q: Optional configuration vector to override internal data.qpos with.
+            q: Optional configuration vector to override internal `data.qpos` with.
         """
         if q is not None:
             self.data.qpos = q
@@ -74,9 +67,6 @@ class Configuration:
 
         Args:
             key_name: The name of the keyframe.
-
-        Raises:
-            ValueError: if no key named `key` was found in the model.
         """
         key_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, key_name)
         if key_id == -1:
@@ -193,6 +183,33 @@ class Configuration:
             translation=xpos,
         )
 
+    def get_transform(
+        self,
+        source_name: str,
+        source_type: str,
+        dest_name: str,
+        dest_type: str,
+    ) -> SE3:
+        """Get the pose of a frame with respect to another frame at the current
+        configuration.
+
+        Args:
+            source_name: Name of the frame in the MJCF.
+            source_type: Source type of frame. Can be a geom, a body or a site.
+            dest_name: Name of the frame to get the pose in.
+            dest_type: Dest type of frame. Can be a geom, a body or a site.
+
+        Returns:
+            The pose of `source_name` in `dest_name`.
+        """
+        transform_source_to_world = self.get_transform_frame_to_world(
+            source_name, source_type
+        )
+        transform_dest_to_world = self.get_transform_frame_to_world(
+            dest_name, dest_type
+        )
+        return transform_dest_to_world.inverse() @ transform_source_to_world
+
     def integrate(self, velocity: np.ndarray, dt: float) -> np.ndarray:
         """Integrate a velocity starting from the current configuration.
 

{mink-0.0.2 → mink-0.0.3}/mink/lie/__init__.py

@@ -1,13 +1,12 @@
 from .base import MatrixLieGroup
 from .se3 import SE3
 from .so3 import SO3
-from .utils import get_epsilon, mat2quat, skew
+from .utils import get_epsilon, skew
 
 __all__ = (
     "SE3",
     "SO3",
     "MatrixLieGroup",
     "get_epsilon",
-    "mat2quat",
     "skew",
 )

{mink-0.0.2 → mink-0.0.3}/mink/lie/utils.py

@@ -1,4 +1,3 @@
-import mujoco
 import numpy as np
 
 
@@ -17,13 +16,6 @@ def skew(x: np.ndarray) -> np.ndarray:
             [0.0, -wz, wy],
             [wz, 0.0, -wx],
             [-wy, wx, 0.0],
-        ]
+        ],
+        dtype=x.dtype,
     )
-
-
-def mat2quat(mat: np.ndarray):
-    """Convert a MuJoCo matrix (9,) to a quaternion (4,)."""
-    assert mat.shape == (9,)
-    quat = np.empty(4, dtype=np.float64)
-    mujoco.mju_mat2Quat(quat, mat)
-    return quat

{mink-0.0.2 → mink-0.0.3}/mink/limits/collision_avoidance_limit.py

@@ -18,7 +18,7 @@ CollisionPairs = Sequence[CollisionPair]
 
 
 @dataclass(frozen=True)
-class _Contact:
+class Contact:
     dist: float
     fromto: np.ndarray
     geom1: int
@@ -35,6 +35,22 @@ class _Contact:
         return self.dist == self.distmax and not self.fromto.any()
 
 
+def compute_contact_normal_jacobian(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    contact: Contact,
+) -> np.ndarray:
+    geom1_body = model.geom_bodyid[contact.geom1]
+    geom2_body = model.geom_bodyid[contact.geom2]
+    geom1_contact_pos = contact.fromto[:3]
+    geom2_contact_pos = contact.fromto[3:]
+    jac2 = np.empty((3, model.nv))
+    mujoco.mj_jac(model, data, jac2, None, geom2_contact_pos, geom2_body)
+    jac1 = np.empty((3, model.nv))
+    mujoco.mj_jac(model, data, jac1, None, geom1_contact_pos, geom1_body)
+    return contact.normal @ (jac2 - jac1)
+
+
 def _is_welded_together(model: mujoco.MjModel, geom_id1: int, geom_id2: int) -> bool:
     """Returns true if the geoms are part of the same body, or if their bodies are
     welded together."""
@@ -166,7 +182,9 @@ class CollisionAvoidanceLimit(Limit):
                 upper_bound[idx] = (self.gain * dist / dt) + self.bound_relaxation
             else:
                 upper_bound[idx] = self.bound_relaxation
-            jac = self._compute_contact_normal_jacobian(configuration.data, contact)
+            jac = compute_contact_normal_jacobian(
+                self.model, configuration.data, contact
+            )
             coefficient_matrix[idx] = -jac
         return Constraint(G=coefficient_matrix, h=upper_bound)
 
@@ -174,7 +192,7 @@ class CollisionAvoidanceLimit(Limit):
 
     def _compute_contact_with_minimum_distance(
         self, data: mujoco.MjData, geom1_id: int, geom2_id: int
-    ) -> _Contact:
+    ) -> Contact:
         """Returns the smallest signed distance between a geom pair."""
         fromto = np.empty(6)
         dist = mujoco.mj_geomDistance(
@@ -185,41 +203,10 @@ class CollisionAvoidanceLimit(Limit):
             self.collision_detection_distance,
             fromto,
         )
-        return _Contact(
+        return Contact(
             dist, fromto, geom1_id, geom2_id, self.collision_detection_distance
         )
 
-    def _compute_contact_normal_jacobian(
-        self, data: mujoco.MjData, contact: _Contact
-    ) -> np.ndarray:
-        """Computes the Jacobian mapping joint velocities to the normal component of
-        the relative Cartesian linear velocity between the geom pair.
-
-        The Jacobian-velocity relationship is given as:
-
-            J dq = n^T (v_2 - v_1)
-
-        where:
-        * J is the computed Jacobian.
-        * dq is the joint velocity vector.
-        * n^T is the transpose of the normal pointing from contact.geom1 to
-            contact.geom2.
-        *  v_1, v_2 are the linear components of the Cartesian velocity of the two
-            closest points in contact.geom1 and contact.geom2.
-
-        Note: n^T (v_2 - v_1) is a scalar that is positive if the geoms are moving away
-        from each other, and negative if they are moving towards each other.
-        """
-        geom1_body = self.model.geom_bodyid[contact.geom1]
-        geom2_body = self.model.geom_bodyid[contact.geom2]
-        geom1_contact_pos = contact.fromto[:3]
-        geom2_contact_pos = contact.fromto[3:]
-        jac2 = np.empty((3, self.model.nv))
-        mujoco.mj_jac(self.model, data, jac2, None, geom2_contact_pos, geom2_body)
-        jac1 = np.empty((3, self.model.nv))
-        mujoco.mj_jac(self.model, data, jac1, None, geom1_contact_pos, geom1_body)
-        return contact.normal @ (jac2 - jac1)
-
     def _homogenize_geom_id_list(self, geom_list: GeomSequence) -> List[int]:
         """Take a heterogeneous list of geoms (specified via ID or name) and return
         a homogenous list of IDs (int)."""

{mink-0.0.2 → mink-0.0.3}/mink/solve_ik.py

@@ -27,8 +27,8 @@ def _compute_qp_inequalities(
 ) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
     if limits is None:
         limits = [ConfigurationLimit(configuration.model)]
-    G_list = []
-    h_list = []
+    G_list: list[np.ndarray] = []
+    h_list: list[np.ndarray] = []
     for limit in limits:
         inequality = limit.compute_qp_inequalities(configuration, dt)
         if not inequality.inactive:

{mink-0.0.2 → mink-0.0.3}/mink/tasks/__init__.py

@@ -11,6 +11,7 @@ from .exceptions import (
 )
 from .frame_task import FrameTask
 from .posture_task import PostureTask
+from .relative_frame_task import RelativeFrameTask
 from .task import Objective, Task
 
 __all__ = (
@@ -19,6 +20,7 @@ __all__ = (
     "Objective",
     "DampingTask",
     "PostureTask",
+    "RelativeFrameTask",
     "Task",
     "TargetNotSet",
     "InvalidTarget",

{mink-0.0.2 → mink-0.0.3}/mink/tasks/damping_task.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import mujoco
+import numpy.typing as npt
 
 from .posture_task import PostureTask
 
@@ -14,6 +15,6 @@ class DampingTask(PostureTask):
     target configuration are set to 0 and qpos0 respectively.
     """
 
-    def __init__(self, model: mujoco.MjModel, cost: float):
+    def __init__(self, model: mujoco.MjModel, cost: npt.ArrayLike):
         super().__init__(model=model, cost=cost, gain=0.0, lm_damping=0.0)
         self.target_q = model.qpos0

{mink-0.0.2 → mink-0.0.3}/mink/tasks/posture_task.py

@@ -29,15 +29,12 @@ class PostureTask(Task):
     def __init__(
         self,
         model: mujoco.MjModel,
-        cost: float,
+        cost: npt.ArrayLike,
         gain: float = 1.0,
         lm_damping: float = 0.0,
     ):
-        if cost < 0.0:
-            raise TaskDefinitionError(f"{self.__class__.__name__} cost must be >= 0")
-
         super().__init__(
-            cost=np.asarray([cost] * model.nv),
+            cost=np.zeros((model.nv,)),
             gain=gain,
             lm_damping=lm_damping,
         )
@@ -52,6 +49,18 @@ class PostureTask(Task):
 
         self.k = model.nv
         self.nq = model.nq
+        self.set_cost(cost)
+
+    def set_cost(self, cost: npt.ArrayLike) -> None:
+        cost = np.atleast_1d(cost)
+        if cost.ndim != 1 or cost.shape[0] not in (1, self.k):
+            raise TaskDefinitionError(
+                f"{self.__class__.__name__} cost must be a vector of shape (1,) "
+                f"(aka identical cost for all dofs) or ({self.k},). Got {cost.shape}"
+            )
+        if not np.all(cost >= 0.0):
+            raise TaskDefinitionError(f"{self.__class__.__name__} cost should be >= 0")
+        self.cost[: self.k] = cost
 
     def set_target(self, target_q: npt.ArrayLike) -> None:
         """Set the target posture.

mink-0.0.3/mink/tasks/relative_frame_task.py

@@ -0,0 +1,142 @@
+"""Relative frame task implementation."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+
+from ..configuration import Configuration
+from ..lie import SE3
+from .exceptions import TargetNotSet, TaskDefinitionError
+from .task import Task
+
+
+class RelativeFrameTask(Task):
+    """Regulate the pose of a frame relative to another frame.
+
+    Attributes:
+        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`.
+        root_name: Name of the frame the task is relative to.
+        root_type: The root frame type: `body`, `geom` or `site`.
+        transform_target_to_root: Target pose in the root frame.
+    """
+
+    k: int = 6
+    transform_target_to_root: Optional[SE3]
+
+    def __init__(
+        self,
+        frame_name: str,
+        frame_type: str,
+        root_name: str,
+        root_type: str,
+        position_cost: npt.ArrayLike,
+        orientation_cost: npt.ArrayLike,
+        gain: float = 1.0,
+        lm_damping: float = 0.0,
+    ):
+        super().__init__(cost=np.zeros((self.k,)), gain=gain, lm_damping=lm_damping)
+        self.frame_name = frame_name
+        self.frame_type = frame_type
+        self.root_name = root_name
+        self.root_type = root_type
+        self.position_cost = position_cost
+        self.orientation_cost = orientation_cost
+        self.transform_target_to_root = None
+
+        self.set_position_cost(position_cost)
+        self.set_orientation_cost(orientation_cost)
+
+    def set_position_cost(self, position_cost: npt.ArrayLike) -> None:
+        position_cost = np.atleast_1d(position_cost)
+        if position_cost.ndim != 1 or position_cost.shape[0] not in (1, 3):
+            raise TaskDefinitionError(
+                f"{self.__class__.__name__} position cost should be a vector of shape "
+                "1 (aka identical cost for all coordinates) or (3,) but got "
+                f"{position_cost.shape}"
+            )
+        if not np.all(position_cost >= 0.0):
+            raise TaskDefinitionError(
+                f"{self.__class__.__name__} position cost should be >= 0"
+            )
+        self.cost[:3] = position_cost
+
+    def set_orientation_cost(self, orientation_cost: npt.ArrayLike) -> None:
+        orientation_cost = np.atleast_1d(orientation_cost)
+        if orientation_cost.ndim != 1 or orientation_cost.shape[0] not in (1, 3):
+            raise TaskDefinitionError(
+                f"{self.__class__.__name__} orientation cost should be a vector of "
+                "shape 1 (aka identical cost for all coordinates) or (3,) but got "
+                f"{orientation_cost.shape}"
+            )
+        if not np.all(orientation_cost >= 0.0):
+            raise TaskDefinitionError(
+                f"{self.__class__.__name__} position cost should be >= 0"
+            )
+        self.cost[3:] = orientation_cost
+
+    def set_target(self, transform_target_to_root: SE3) -> None:
+        """Set the target pose in the root frame.
+
+        Args:
+            transform_target_to_root: Transform from the task target frame to the
+                root frame.
+        """
+        self.transform_target_to_root = transform_target_to_root.copy()
+
+    def set_target_from_configuration(self, configuration: Configuration) -> None:
+        """Set the target pose from a given robot configuration.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+        """
+        self.set_target(
+            configuration.get_transform(
+                self.frame_name,
+                self.frame_type,
+                self.root_name,
+                self.root_type,
+            )
+        )
+
+    def compute_error(self, configuration: Configuration) -> np.ndarray:
+        if self.transform_target_to_root is None:
+            raise TargetNotSet(self.__class__.__name__)
+
+        transform_frame_to_root = configuration.get_transform(
+            self.frame_name,
+            self.frame_type,
+            self.root_name,
+            self.root_type,
+        )
+        return transform_frame_to_root.rminus(self.transform_target_to_root)
+
+    def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
+        if self.transform_target_to_root is None:
+            raise TargetNotSet(self.__class__.__name__)
+
+        jacobian_frame_in_frame = configuration.get_frame_jacobian(
+            self.frame_name, self.frame_type
+        )
+        jacobian_root_in_root = configuration.get_frame_jacobian(
+            self.root_name, self.root_type
+        )
+
+        transform_frame_to_root = configuration.get_transform(
+            self.frame_name,
+            self.frame_type,
+            self.root_name,
+            self.root_type,
+        )
+        transform_frame_to_target = (
+            self.transform_target_to_root.inverse() @ transform_frame_to_root
+        )
+
+        return transform_frame_to_target.jlog() @ (
+            jacobian_frame_in_frame
+            - transform_frame_to_root.inverse().adjoint() @ jacobian_root_in_root
+        )

{mink-0.0.2 → mink-0.0.3}/pyproject.toml

@@ -4,12 +4,13 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "mink"
+version = "0.0.3"
 readme = "README.md"
 authors = [
     {name = "Kevin Zakka", email = "zakka@berkeley.edu"},
 ]
 requires-python = ">=3.9"
-dynamic = ["version", "description"]
+dynamic = ["description"]
 classifiers = [
     "Development Status :: 5 - Production/Stable",
     "Framework :: Robot Framework :: Library",
@@ -91,3 +92,15 @@ python_version = "3.12"
 ignore_missing_imports = true
 warn_unused_configs = true
 exclude = ["scripts/"]
+
+[tool.coverage.report]
+exclude_also = [
+    "raise NotImplementedError",
+    "if __name__ == .__main__.:",
+    "class .*\\bProtocol\\):",
+    "@(abc\\.)?abstractmethod",
+    "def __repr__",
+]
+omit = [
+    "*exceptions.py",
+]

mink-0.0.2/mink/lie/tests/test_axioms.py

@@ -1,51 +0,0 @@
-"""Tests for group axioms."""
-
-from typing import Type
-
-from absl.testing import absltest, parameterized
-
-from mink import lie
-
-from ..base import MatrixLieGroup
-from .utils import assert_transforms_close
-
-
-@parameterized.named_parameters(
-    ("SO3", lie.SO3),
-    ("SE3", lie.SE3),
-)
-class TestAxioms(parameterized.TestCase):
-    def test_closure(self, group: Type[MatrixLieGroup]):
-        transform_a = group.sample_uniform()
-        transform_b = group.sample_uniform()
-        composed = transform_a @ transform_b
-        assert_transforms_close(composed, composed.normalize())
-        composed = transform_b @ transform_a
-        assert_transforms_close(composed, composed.normalize())
-        composed = transform_a @ transform_b
-        assert_transforms_close(composed, composed.normalize())
-
-    def test_identity(self, group: Type[MatrixLieGroup]):
-        transform = group.sample_uniform()
-        identity = group.identity()
-        assert_transforms_close(transform, identity @ transform)
-        assert_transforms_close(transform, transform @ identity)
-
-    def test_inverse(self, group: Type[MatrixLieGroup]):
-        transform = group.sample_uniform()
-        identity = group.identity()
-        assert_transforms_close(identity, transform.inverse() @ transform)
-        assert_transforms_close(identity, transform @ transform.inverse())
-
-    def test_associative(self, group: Type[MatrixLieGroup]):
-        transform_a = group.sample_uniform()
-        transform_b = group.sample_uniform()
-        transform_c = group.sample_uniform()
-        assert_transforms_close(
-            (transform_a @ transform_b) @ transform_c,
-            transform_a @ (transform_b @ transform_c),
-        )
-
-
-if __name__ == "__main__":
-    absltest.main()