mink 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

mink/__init__.py

@@ -25,7 +25,16 @@ from .limits import (
     VelocityLimit,
 )
 from .solve_ik import build_ik, solve_ik
-from .tasks import ComTask, FrameTask, Objective, PostureTask, TargetNotSet, Task
+from .tasks import (
+    ComTask,
+    DampingTask,
+    FrameTask,
+    Objective,
+    PostureTask,
+    RelativeFrameTask,
+    TargetNotSet,
+    Task,
+)
 from .utils import (
     custom_configuration_vector,
     get_body_geom_ids,
@@ -34,14 +43,14 @@ from .utils import (
     move_mocap_to_frame,
 )
 
-__version__ = "0.0.1"
-
 __all__ = (
     "ComTask",
     "Configuration",
     "build_ik",
     "solve_ik",
+    "DampingTask",
     "FrameTask",
+    "RelativeFrameTask",
     "PostureTask",
     "Task",
     "Objective",

mink/configuration.py

@@ -1,3 +1,12 @@
+"""Configuration space of a robot model.
+
+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>`__,
+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
 
 import mujoco
@@ -9,14 +18,19 @@ from .lie import SE3, SO3
 
 
 class Configuration:
-    """A struct that provides convenient access to kinematic quantities such as frame
-    transforms and frame jacobians. Frames can be defined at bodies, geoms or sites.
+    """Encapsulates a model and data for convenient access to kinematic quantities.
+
+    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.
 
-    The `update` function ensures the proper forward kinematics functions have been
-    called, namely:
+    Key functionalities include:
 
-    - mujoco.mj_kinematics(model, data)
-    - mujoco.mj_comPos(model, data)
+    * 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__(
@@ -27,9 +41,9 @@ class Configuration:
         """Constructor.
 
         Args:
-            model: An instance of MjModel.
-            q: Optional configuration to initialize from. If None, the configuration
-                is initialized to the reference configuration `qpos0`.
+            model: Mujoco model.
+            q: Configuration to initialize from. If None, the configuration
+                is initialized to the default configuration `qpos0`.
         """
         self.model = model
         self.data = mujoco.MjData(model)
@@ -38,15 +52,13 @@ class Configuration:
     def update(self, q: Optional[np.ndarray] = None) -> None:
         """Run forward kinematics.
 
-        The minimal function call required to get updated frame transforms (aka forward
-        kinematics) is `mj_kinematics`. An extra call to `mj_comPos` is needed for
-        updated Jacobians.
-
-        Args:`
-            q: Optional configuration vector to override internal data.qpos with.
+        Args:
+            q: Optional configuration vector to override internal `data.qpos` with.
         """
         if q is not None:
             self.data.qpos = q
+        # The minimal function call required to get updated frame transforms is
+        # mj_kinematics. An extra call to mj_comPos is required for updated Jacobians.
         mujoco.mj_kinematics(self.model, self.data)
         mujoco.mj_comPos(self.model, self.data)
 
@@ -55,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:
@@ -95,24 +104,27 @@ class Configuration:
                     )
                 else:
                     print(
-                        f"Value {qval} at index {jnt} is out of limits: "
-                        f"[{qmin}, {qmax}]"
+                        f"Value {qval:.2f} at index {jnt} is outside of its limits: "
+                        f"[{qmin:.2f}, {qmax:.2f}]"
                     )
 
     def get_frame_jacobian(self, frame_name: str, frame_type: str) -> np.ndarray:
-        """Compute the Jacobian matrix of a frame velocity.
+        r"""Compute the Jacobian matrix of a frame velocity.
+
+        Denoting our frame by :math:`B` and the world frame by :math:`W`, the
+        Jacobian matrix :math:`{}_B J_{WB}` is related to the body velocity
+        :math:`{}_B v_{WB}` by:
 
-        Denoting our frame by B and the world frame by W, the Jacobian matrix B_J_WB
-        is related to the body velocity B_v_WB by:
+        .. math::
 
-            B_v_WB = B_J_WB q_dot
+            {}_B v_{WB} = {}_B J_{WB} \dot{q}
 
         Args:
             frame_name: Name of the frame in the MJCF.
             frame_type: Type of frame. Can be a geom, a body or a site.
 
         Returns:
-            Jacobian B_J_WB of the frame.
+            Jacobian :math:`{}_B J_{WB}` of the frame.
         """
         if frame_type not in consts.SUPPORTED_FRAMES:
             raise exceptions.UnsupportedFrame(frame_type, consts.SUPPORTED_FRAMES)
@@ -142,9 +154,7 @@ class Configuration:
         return jac
 
     def get_transform_frame_to_world(self, frame_name: str, frame_type: str) -> SE3:
-        """Get the pose of a frame in the current configuration.
-
-        Denoting our frame by B and the world frame by W, this function returns T_WB.
+        """Get the pose of a frame at the current configuration.
 
         Args:
             frame_name: Name of the frame in the MJCF.
@@ -173,11 +183,38 @@ 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.
 
         Args:
-            velocity: The velocity.
+            velocity: The velocity in tangent space.
             dt: Integration duration in [s].
 
         Returns:
@@ -191,7 +228,7 @@ class Configuration:
         """Integrate a velocity and update the current configuration inplace.
 
         Args:
-            velocity: The velocity.
+            velocity: The velocity in tangent space.
             dt: Integration duration in [s].
         """
         mujoco.mj_integratePos(self.model, self.data.qpos, velocity, dt)

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/lie/base.py

@@ -10,8 +10,8 @@ class MatrixLieGroup(abc.ABC):
 
     Attributes:
         matrix_dim: Dimension of square matrix output.
-        parameters_dim:: Dimension of underlying parameters.
-        tangent_dim:: Dimension of tangent space.
+        parameters_dim: Dimension of underlying parameters.
+        tangent_dim: Dimension of tangent space.
         space_dim: Dimension of coordinates that can be transformed.
     """
 

mink/lie/so3.py

@@ -14,6 +14,8 @@ _INVERT_QUAT_SIGN = np.array([1.0, -1.0, -1.0, -1.0], dtype=np.float64)
 
 
 class RollPitchYaw(NamedTuple):
+    """Struct containing roll, pitch, and yaw Euler angles."""
+
     roll: float
     pitch: float
     yaw: float

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/limits/collision_avoidance_limit.py

@@ -1,32 +1,4 @@
-"""Collision avoidance limit.
-
-Derivation
-==========
-
-p1, p2: closest points between g1 and g2
-d: distance between g1 and g2 (d = ||p1 - p2||)
-n: normal vector from g1 to g2 (n = (p2 - p1) / ||p2 - p1||)
-
-The relative velocity constraint between g1 and g2 is given by:
-
-    n^T [J2(p2) - J1(p1)] dq <= -k * (d - d_min) / dt
-    V_n <= -k * (d - d_min)
-
-where [J2(p2) - J1(p1)] dq denotes the relative velocity between g1 and g2, and
-n^T [J2(p2) - J1(p1)] dq projects this velocity onto the direction connecting the
-closest points p1 and p2.
-
-We have three cases for the distance d:
-
-1. V_n <= R_b                                if dist_c < dist_m
-2. V_n <= -k * (dist_c - dist_m) / dt + R_b  if dist_m <= d <= dist_d
-3. V_n <= inf                                otherwise
-
-where:
-- dist_c: current normal distance between g1 and g2
-- dist_m: minimum allowed distance between g1 and g2
-- dist_d: collision detection distance
-"""
+"""Collision avoidance limit."""
 
 import itertools
 from dataclasses import dataclass
@@ -63,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."""
@@ -107,7 +95,31 @@ def _is_pass_contype_conaffinity_check(
 
 
 class CollisionAvoidanceLimit(Limit):
-    """Normal velocity limit between geom pairs."""
+    """Normal velocity limit between geom pairs.
+
+    Attributes:
+        model: MuJoCo model.
+        geom_pairs: Set of collision pairs in which to perform active collision
+            avoidance. A collision pair is defined as a pair of geom groups. A geom
+            group is a set of geom names. For each geom pair, the solver will
+            attempt to compute joint velocities that avoid collisions between every
+            geom in the first geom group with every geom in the second geom group.
+            Self collision is achieved by adding a collision pair with the same
+            geom group in both pair fields.
+        gain: Gain factor in (0, 1] that determines how fast the geoms are
+            allowed to move towards each other at each iteration. Smaller values
+            are safer but may make the geoms move slower towards each other.
+        minimum_distance_from_collisions: The minimum distance to leave between
+            any two geoms. A negative distance allows the geoms to penetrate by
+            the specified amount.
+        collision_detection_distance: The distance between two geoms at which the
+            active collision avoidance limit will be active. A large value will
+            cause collisions to be detected early, but may incur high computational
+            cost. A negative value will cause the geoms to be detected only after
+            they penetrate by the specified amount.
+        bound_relaxation: An offset on the upper bound of each collision avoidance
+            constraint.
+    """
 
     def __init__(
         self,
@@ -141,7 +153,7 @@ class CollisionAvoidanceLimit(Limit):
                 cost. A negative value will cause the geoms to be detected only after
                 they penetrate by the specified amount.
             bound_relaxation: An offset on the upper bound of each collision avoidance
-                constraint to tighten or relax
+                constraint.
         """
         self.model = model
         self.gain = gain
@@ -170,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)
 
@@ -193,37 +207,6 @@ class CollisionAvoidanceLimit(Limit):
             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/limits/configuration_limit.py

@@ -1,24 +1,4 @@
-"""Limit on joint positions.
-
-Derivation
-==========
-
-Using a first order Taylor expansion on the configuration, we can write the limit as:
-
-    q_min     <= q + v * dt <= q_max
-    q_min     <= q + dq     <= q_max
-    q_min - q <= dq         <= q_max - q
-
-Rewriting as G dq <= h:
-
-    +I * dq <= q_max - q
-    -I * dq <= q - q_min
-
-Stacking them together, we get:
-
-    G = [+I, -I]
-    h = [q_max - q, q - q_min]
-"""
+"""Joint position limit."""
 
 import mujoco
 import numpy as np
@@ -30,16 +10,9 @@ from .limit import Constraint, Limit
 
 
 class ConfigurationLimit(Limit):
-    """Limit for joint positions in a model.
-
-    Floating base joints (joint type="free") are ignored.
+    """Inequality constraint on joint positions in a robot model.
 
-    Attributes:
-        indices: Tangent indices corresponding to configuration-limited joints.
-        projection_matrix: Projection from tangent space to subspace with
-            configuration-limited joints.
-        lower: Lower configuration limit.
-        upper: Upper configuration limit.
+    Floating base joints are ignored.
     """
 
     def __init__(
@@ -95,6 +68,26 @@ class ConfigurationLimit(Limit):
         configuration: Configuration,
         dt: float,
     ) -> Constraint:
+        r"""Compute the configuration-dependent joint position limits.
+
+        The limits are defined as:
+
+        .. math::
+
+            {q \ominus q_{min}} \leq \Delta q \leq {q_{max} \ominus q}
+
+        where :math:`q \in {\cal C}` is the robot's configuration and
+        :math:`\Delta q \in T_q({\cal C})` is the displacement in the tangent
+        space at :math:`q`. See the :ref:`derivations` section for more information.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+            dt: Integration timestep in [s].
+
+        Returns:
+            Pair :math:`(G, h)` representing the inequality constraint as
+            :math:`G \Delta q \leq h`, or ``None`` if there is no limit.
+        """
         del dt  # Unused.
 
         # Upper.

mink/limits/limit.py

@@ -1,3 +1,5 @@
+"""All kinematic limits derive from the :class:`Limit` base class."""
+
 import abc
 from typing import NamedTuple, Optional
 
@@ -7,21 +9,29 @@ from ..configuration import Configuration
 
 
 class Constraint(NamedTuple):
-    """Linear inequalities in the form G(q) dq <= h(q).
+    r"""Linear inequality constraint of the form :math:`G(q) \Delta q \leq h(q)`.
 
-    The limit is considered inactive when G or h are None.
+    Inactive if G and h are None.
     """
 
-    G: Optional[np.ndarray] = None  # (nv, nv)
-    h: Optional[np.ndarray] = None  # (nv,)
+    G: Optional[np.ndarray] = None
+    """Shape (nv, nv)."""
+    h: Optional[np.ndarray] = None
+    """Shape (nv,)."""
 
     @property
     def inactive(self) -> bool:
-        return self.G is None or self.h is None
+        """Returns True if the constraint is inactive."""
+        return self.G is None and self.h is None
 
 
 class Limit(abc.ABC):
-    """Abstract base class for kinematic limits."""
+    """Abstract base class for kinematic limits.
+
+    Subclasses must implement the :py:meth:`~Limit.compute_qp_inequalities` method
+    which takes in the current robot configuration and integration time step and
+    returns an instance of :class:`Constraint`.
+    """
 
     @abc.abstractmethod
     def compute_qp_inequalities(
@@ -29,13 +39,21 @@ class Limit(abc.ABC):
         configuration: Configuration,
         dt: float,
     ) -> Constraint:
-        """Compute limit as linearized QP inequalities.
+        r"""Compute limit as linearized QP inequalities of the form:
+
+        .. math::
+
+            G(q) \Delta q \leq h(q)
+
+        where :math:`q \in {\cal C}` is the robot's configuration and
+        :math:`\Delta q \in T_q({\cal C})` is the displacement in the tangent
+        space at :math:`q`.
 
         Args:
-            configuration: Current configuration.
+            configuration: Robot configuration :math:`q`.
             dt: Integration time step in [s].
 
         Returns:
-            Pair (G, h) representing the inequality constraint.
+            Pair :math:`(G, h)`.
         """
         raise NotImplementedError

mink/limits/velocity_limit.py

@@ -1,24 +1,4 @@
-"""Limit on joint velocities.
-
-Derivation
-==========
-
-Given maximum joint velocity magnitudes v_max, we can express joint velocity limits as:
-
-    -v_max      <= v       <= v_max
-    -v_max      <= dq / dt <= v_max
-    -v_max * dt <= dq      <= v_max * dt
-
-Rewriting as G dq <= h:
-
-    +I * dq <= v_max * dt
-    -I * dq <= v_max * dt
-
-Stacking them together, we get:
-
-    G = [+I, -I]
-    h = [v_max * dt, v_max * dt]
-"""
+"""Joint velocity limit."""
 
 from typing import Mapping
 
@@ -33,17 +13,22 @@ from .limit import Constraint, Limit
 
 
 class VelocityLimit(Limit):
-    """Limit for joint velocities in a model.
+    """Inequality constraint on joint velocities in a robot model.
 
     Floating base joints are ignored.
 
     Attributes:
         indices: Tangent indices corresponding to velocity-limited joints.
-        limit: Maximum allowed velocity magnitude for velocity-limited joints.
+        limit: Maximum allowed velocity magnitude for velocity-limited joints, in
+            [m]/[s] for slide joints and [rad]/[s] for hinge joints.
         projection_matrix: Projection from tangent space to subspace with
             velocity-limited joints.
     """
 
+    indices: np.ndarray
+    limit: np.ndarray
+    projection_matrix: np.ndarray
+
     def __init__(
         self,
         model: mujoco.MjModel,
@@ -52,6 +37,7 @@ class VelocityLimit(Limit):
         """Initialize velocity limits.
 
         Args:
+            model: MuJoCo model.
             velocities: Dictionary mapping joint name to maximum allowed magnitude in
                 [m]/[s] for slide joints and [rad]/[s] for hinge joints.
         """
@@ -84,6 +70,27 @@ class VelocityLimit(Limit):
     def compute_qp_inequalities(
         self, configuration: Configuration, dt: float
     ) -> Constraint:
+        r"""Compute the configuration-dependent joint velocity limits.
+
+        The limits are defined as:
+
+        .. math::
+
+            -v_{\text{max}} \cdot dt \leq \Delta q \leq v_{\text{max}} \cdot dt
+
+        where :math:`v_{max} \in {\cal T}` is the robot's velocity limit
+        vector and :math:`\Delta q \in T_q({\cal C})` is the displacement in the
+        tangent space at :math:`q`. See the :ref:`derivations` section for
+        more information.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+            dt: Integration timestep in [s].
+
+        Returns:
+            Pair :math:`(G, h)` representing the inequality constraint as
+            :math:`G \Delta q \leq h`, or ``None`` if there is no limit.
+        """
         del configuration  # Unused.
         if self.projection_matrix is None:
             return Constraint()

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/tasks/__init__.py

@@ -1,6 +1,7 @@
 """Kinematic tasks."""
 
 from .com_task import ComTask
+from .damping_task import DampingTask
 from .exceptions import (
     InvalidDamping,
     InvalidGain,
@@ -10,13 +11,16 @@ 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__ = (
     "ComTask",
     "FrameTask",
     "Objective",
+    "DampingTask",
     "PostureTask",
+    "RelativeFrameTask",
     "Task",
     "TargetNotSet",
     "InvalidTarget",