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

mink/tasks/com_task.py

@@ -1,4 +1,6 @@
-"""Center of mass task implementation."""
+"""Center-of-mass task implementation."""
+
+from __future__ import annotations
 
 from typing import Optional
 
@@ -12,7 +14,7 @@ from .task import Task
 
 
 class ComTask(Task):
-    """Regulate the center of mass (CoM) of a robot.
+    """Regulate the center-of-mass (CoM) of a robot.
 
     Attributes:
         target_com: Target position of the CoM.
@@ -33,7 +35,6 @@ class ComTask(Task):
         self.set_cost(cost)
 
     def set_cost(self, cost: npt.ArrayLike) -> None:
-        """Set a new cost for all CoM coordinates."""
         cost = np.atleast_1d(cost)
         if cost.ndim != 1 or cost.shape[0] not in (1, self.k):
             raise TaskDefinitionError(
@@ -46,7 +47,11 @@ class ComTask(Task):
         self.cost[:] = cost
 
     def set_target(self, target_com: npt.ArrayLike) -> None:
-        """Set the target CoM position in the world frame."""
+        """Set the target CoM position in the world frame.
+
+        Args:
+            target_com: Desired center-of-mass position in the world frame.
+        """
         target_com = np.atleast_1d(target_com)
         if target_com.ndim != 1 or target_com.shape[0] != (self.k):
             raise InvalidTarget(
@@ -56,14 +61,21 @@ class ComTask(Task):
         self.target_com = target_com.copy()
 
     def set_target_from_configuration(self, configuration: Configuration) -> None:
-        """Set the target CoM from a given robot configuration."""
+        """Set the target CoM from a given robot configuration.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+        """
         self.set_target(configuration.data.subtree_com[1])
 
     def compute_error(self, configuration: Configuration) -> np.ndarray:
         """Compute the CoM task error.
 
-        The error is the difference between the target CoM and current CoM positions,
-        expressed in the world frame.
+        Args:
+            configuration: Robot configuration :math:`q`.
+
+        Returns:
+            Center-of-mass task error vector :math:`e(q)`.
         """
         if self.target_com is None:
             raise TargetNotSet(self.__class__.__name__)
@@ -72,8 +84,11 @@ class ComTask(Task):
     def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
         """Compute the CoM task Jacobian.
 
-        The task Jacobian is the derivative of the task error with respect to the
-        current configuration. It has dimension (3, nv).
+        Args:
+            configuration: Robot configuration :math:`q`.
+
+        Returns:
+            Center-of-mass task jacobian :math:`J(q)`.
         """
         if self.target_com is None:
             raise TargetNotSet(self.__class__.__name__)

mink/tasks/damping_task.py

@@ -0,0 +1,20 @@
+"""Damping task implementation."""
+
+from __future__ import annotations
+
+import mujoco
+import numpy.typing as npt
+
+from .posture_task import PostureTask
+
+
+class DampingTask(PostureTask):
+    """Minimize joint velocities.
+
+    This task is implemented as a special case of the PostureTask where the gain and
+    target configuration are set to 0 and qpos0 respectively.
+    """
+
+    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/tasks/frame_task.py

@@ -1,5 +1,7 @@
 """Frame task implementation."""
 
+from __future__ import annotations
+
 from typing import Optional
 
 import numpy as np
@@ -12,12 +14,13 @@ from .task import Task
 
 
 class FrameTask(Task):
-    """Regulate the pose of a robot frame in the world frame.
+    """Regulate the pose of a frame expressed in the world frame.
 
     Attributes:
-        frame_name: Name of the frame to regulate.
-        frame_type: The frame type: body, geom or site.
-        transform_frame_to_world: Target pose of the frame in the world frame.
+        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.
     """
 
     k: int = 6
@@ -71,21 +74,44 @@ class FrameTask(Task):
         self.cost[3:] = orientation_cost
 
     def set_target(self, transform_target_to_world: SE3) -> None:
-        """Set the target pose in the world frame."""
+        """Set the target pose.
+
+        Args:
+            transform_target_to_world: Transform from the task target frame to the
+                world frame.
+        """
         self.transform_target_to_world = transform_target_to_world.copy()
 
     def set_target_from_configuration(self, configuration: Configuration) -> None:
-        """Set the target pose from a given robot configuration."""
+        """Set the target pose from a given robot configuration.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+        """
         self.set_target(
             configuration.get_transform_frame_to_world(self.frame_name, self.frame_type)
         )
 
     def compute_error(self, configuration: Configuration) -> np.ndarray:
-        """Compute the frame task error.
+        r"""Compute the frame task error.
+
+        This error is a twist :math:`e(q) \in se(3)` expressed in the local frame,
+        i.e., it is a body twist. It is computed by taking the right-minus difference
+        between the target pose :math:`T_{0t}` and current frame pose :math:`T_{0b}`:
 
-        This error is a twist expressed in the local frame, i.e., it is a body twist.
-        It is computed by taking the right minus difference between the target pose
-        and the current frame pose: `e = T_wt ⊖ T_wf`.
+        .. math::
+
+            e(q) := {}_b \xi_{0b} = -(T_{t0} \ominus T_{b0})
+            = -\log(T_{t0} \cdot T_{0b}) = -\log(T_{tb}) = \log(T_{bt})
+
+        where :math:`b` denotes our frame, :math:`t` the target frame and
+        :math:`0` the inertial frame.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+
+        Returns:
+            Frame task error vector :math:`e(q)`.
         """
         if self.transform_target_to_world is None:
             raise TargetNotSet(self.__class__.__name__)
@@ -96,10 +122,16 @@ class FrameTask(Task):
         return self.transform_target_to_world.minus(transform_frame_to_world)
 
     def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
-        """Compute the frame task Jacobian.
+        r"""Compute the frame task Jacobian.
+
+        The derivation of the formula for this Jacobian is detailed in
+        [FrameTaskJacobian]_.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
 
-        The task Jacobian is the derivative of the task error with respect to the
-        current configuration. It has dimension (6, nv).
+        Returns:
+            Frame task jacobian :math:`J(q)`.
         """
         if self.transform_target_to_world is None:
             raise TargetNotSet(self.__class__.__name__)

mink/tasks/posture_task.py

@@ -1,3 +1,7 @@
+"""Posture task implementation."""
+
+from __future__ import annotations
+
 from typing import Optional
 
 import mujoco
@@ -11,7 +15,7 @@ from .task import Task
 
 
 class PostureTask(Task):
-    """Regulate joint angles to a desired posture.
+    """Regulate the joint angles of the robot towards a desired posture.
 
     A posture is a vector of actuated joint angles. Floating-base coordinates are not
     affected by this task.
@@ -25,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,
         )
@@ -48,9 +49,25 @@ 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."""
+        """Set the target posture.
+
+        Args:
+            target_q: Desired joint configuration.
+        """
         target_q = np.atleast_1d(target_q)
         if target_q.ndim != 1 or target_q.shape[0] != (self.nq):
             raise InvalidTarget(
@@ -60,14 +77,27 @@ class PostureTask(Task):
         self.target_q = target_q.copy()
 
     def set_target_from_configuration(self, configuration: Configuration) -> None:
-        """Set the target posture from the current configuration."""
+        """Set the target posture from the current configuration.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+        """
         self.set_target(configuration.q)
 
     def compute_error(self, configuration: Configuration) -> np.ndarray:
-        """Compute the posture task error.
+        r"""Compute the posture task error.
 
-        The error is defined as the right minus error between the target posture and
-        the current posture: `q^* ⊖ q`.
+        The error is defined as:
+
+        .. math::
+
+            e(q) = q^* \ominus q
+
+        Args:
+            configuration: Robot configuration :math:`q`.
+
+        Returns:
+            Posture task error vector :math:`e(q)`.
         """
         if self.target_q is None:
             raise TargetNotSet(self.__class__.__name__)
@@ -88,10 +118,19 @@ class PostureTask(Task):
         return qvel
 
     def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
-        """Compute the posture task Jacobian.
+        r"""Compute the posture task Jacobian.
+
+        The task Jacobian is defined as:
+
+        .. math::
+
+            J(q) = I_{n_v}
+
+        Args:
+            configuration: Robot configuration :math:`q`.
 
-        The task Jacobian is the derivative of the task error with respect to the
-        current configuration. It is equal to -I and has dimension (nv, nv).
+        Returns:
+            Posture task jacobian :math:`J(q)`.
         """
         if self.target_q is None:
             raise TargetNotSet(self.__class__.__name__)

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

@@ -1,3 +1,5 @@
+"""All kinematic tasks derive from the :class:`Task` base class."""
+
 import abc
 from typing import NamedTuple
 
@@ -8,10 +10,12 @@ from .exceptions import InvalidDamping, InvalidGain
 
 
 class Objective(NamedTuple):
-    """Quadratic objective function in the form 0.5 x^T H x + c^T x."""
+    r"""Quadratic objective of the form :math:`\frac{1}{2} x^T H x + c^T x`."""
 
-    H: np.ndarray  # (nv, nv)
-    c: np.ndarray  # (nv,)
+    H: np.ndarray
+    """Hessian matrix, of shape (n_v, n_v)"""
+    c: np.ndarray
+    """Linear vector, of shape (n_v,)."""
 
     def value(self, x: np.ndarray) -> float:
         """Returns the value of the objective at the input vector."""
@@ -19,7 +23,33 @@ class Objective(NamedTuple):
 
 
 class Task(abc.ABC):
-    """Abstract base class for kinematic tasks."""
+    r"""Abstract base class for kinematic tasks.
+
+    Subclasses must implement the configuration-dependent task error
+    :py:meth:`~Task.compute_error` and Jacobian :py:meth:`~Task.compute_jacobian`
+    functions.
+
+    The error function :math:`e(q) \in \mathbb{R}^{k}` is the quantity that
+    the task aims to drive to zero (:math:`k` is the dimension of the
+    task). It appears in the first-order task dynamics:
+
+    .. math::
+
+        J(q) \Delta q = -\alpha e(q)
+
+    The Jacobian matrix :math:`J(q) \in \mathbb{R}^{k \times n_v}`, with
+    :math:`n_v` the dimension of the robot's tangent space, is the
+    derivative of the task error :math:`e(q)` with respect to the
+    configuration :math:`q \in \mathbb{R}^{n_q}`. The configuration displacement
+    :math:`\Delta q` is the output of inverse kinematics; we divide it by dt to get a
+    commanded velocity.
+
+    In the first-order task dynamics, the error :math:`e(q)` is multiplied
+    by the task gain :math:`\alpha \in [0, 1]`. This gain can be 1.0 for
+    dead-beat control (*i.e.* converge as fast as possible), but might be
+    unstable as it neglects our first-order approximation. Lower values
+    cause slow down the task, similar to low-pass filtering.
+    """
 
     def __init__(
         self,
@@ -41,7 +71,7 @@ class Task(abc.ABC):
         if not 0.0 <= gain <= 1.0:
             raise InvalidGain("`gain` must be in the range [0, 1]")
 
-        if not lm_damping >= 0.0:
+        if lm_damping < 0.0:
             raise InvalidDamping("`lm_damping` must be >= 0")
 
         self.cost = cost
@@ -50,23 +80,13 @@ class Task(abc.ABC):
 
     @abc.abstractmethod
     def compute_error(self, configuration: Configuration) -> np.ndarray:
-        """Compute the task error function at the current configuration.
-
-        The error function e(q), of shape (k,), is the quantity that the task aims to
-        drive to zero. It appears in the first-order task dynamics:
-
-            J(q) Δq = -alpha e(q)
+        """Compute the task error at the current configuration.
 
-        The Jacobian matrix J(q), of shape (k, nv), is the derivative of the task error
-        with respect to the configuration. This Jacobian is implemented in
-        `compute_jacobian`. Finally, the configuration displacement Δq is the output
-        of inverse kinematics.
+        Args:
+            configuration: Robot configuration :math:`q`.
 
-        In the first-order task dynamics, the error e(q) is multiplied by the task gain
-        alpha which is between [0, 1]. This gain can be 1.0 for dead-beat control
-        (i.e., converge as fast as possible), but might be unstable as it neglects the
-        first-order approximation. Lower values make the task slower and have a similar
-        effect to low-pass filtering.
+        Returns:
+            Task error vector :math:`e(q)`.
         """
         raise NotImplementedError
 
@@ -74,20 +94,33 @@ class Task(abc.ABC):
     def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
         """Compute the task Jacobian at the current configuration.
 
-        The task Jacobian J(q), of shape (k, nv) is the first-order derivative of the
-        error e(q) that defines the task, with k the dimension of the task and
-        nv the dimension of the robot's tangent space.
+        Args:
+            configuration: Robot configuration :math:`q`.
+
+        Returns:
+            Task jacobian :math:`J(q)`.
         """
         raise NotImplementedError
 
     def compute_qp_objective(self, configuration: Configuration) -> Objective:
-        """Compute the matrix-vector pair :math:`(H, c)` of the QP objective.
+        r"""Compute the matrix-vector pair :math:`(H, c)` of the QP objective.
+
+        This pair is such that the contribution of the task to the QP objective is:
 
-        This pair is such that the contribution of the task to the QP objective
-        of the IK is:
+        .. math::
+
+            \| J \Delta q + \alpha e \|_{W}^2 = \frac{1}{2} \Delta q^T H
+            \Delta q + c^T q
+
+        The weight matrix :math:`W \in \mathbb{R}^{k \times k}` weights and
+        normalizes task coordinates to the same unit. The unit of the overall
+        contribution is [cost]^2.
+
+        Args:
+            configuration: Robot configuration :math:`q`.
 
-        The weight matrix W, of shape (k, k), weighs and normalizes task coordinates
-        to the same unit. The unit of the overall contribution is [cost]^2.
+        Returns:
+            Pair :math:`(H(q), c(q))`.
         """
         jacobian = self.compute_jacobian(configuration)  # (k, nv)
         minus_gain_error = -self.gain * self.compute_error(configuration)  # (k,)

mink/utils.py

@@ -17,8 +17,8 @@ def move_mocap_to_frame(
     """Initialize mocap body pose at a desired frame.
 
     Args:
-        model: An instance of mujoco.MjModel.
-        data: An instance of mujoco.MjData.
+        model: Mujoco model.
+        data: Mujoco data.
         mocap_name: The name of the mocap body.
         frame_name: The desired frame name.
         frame_type: The desired frame type. Can be "body", "geom" or "site".
@@ -36,7 +36,15 @@ def move_mocap_to_frame(
 
 
 def get_freejoint_dims(model: mujoco.MjModel) -> tuple[list[int], list[int]]:
-    """Get all floating joint configuration and tangent indices."""
+    """Get all floating joint configuration and tangent indices.
+
+    Args:
+        model: Mujoco model.
+
+    Returns:
+        A (q_ids, v_ids) pair containing all floating joint indices in the
+        configuration and tangent spaces respectively.
+    """
     q_ids: list[int] = []
     v_ids: list[int] = []
     for j in range(model.njnt):
@@ -56,9 +64,9 @@ def custom_configuration_vector(
     """Generate a configuration vector where named joints have specific values.
 
     Args:
-        model: An MjModel instance.
+        model: Mujoco model.
         key_name: Optional keyframe name to initialize the configuration vector from.
-            Otherwise, the default pose qpos0 is used.
+            Otherwise, the default pose `qpos0` is used.
         kwargs: Custom values for joint coordinates.
 
     Returns:
@@ -93,7 +101,7 @@ def get_subtree_geom_ids(model: mujoco.MjModel, body_id: int) -> list[int]:
     """Get all geoms belonging to subtree starting at a given body.
 
     Args:
-        model: An MjModel instance.
+        model: Mujoco model.
         body_id: ID of body where subtree starts.
 
     Returns:
@@ -117,7 +125,7 @@ def get_body_geom_ids(model: mujoco.MjModel, body_id: int) -> list[int]:
     """Get all geoms belonging to a given body.
 
     Args:
-        model: An MjModel instance.
+        model: Mujoco model.
         body_id: ID of body.
 
     Returns: