mink 0.0.1__py3-none-any.whl

mink/tasks/com_task.py

@@ -0,0 +1,82 @@
+"""Center of mass task implementation."""
+
+from typing import Optional
+
+import mujoco
+import numpy as np
+import numpy.typing as npt
+
+from ..configuration import Configuration
+from .exceptions import InvalidTarget, TargetNotSet, TaskDefinitionError
+from .task import Task
+
+
+class ComTask(Task):
+    """Regulate the center of mass (CoM) of a robot.
+
+    Attributes:
+        target_com: Target position of the CoM.
+    """
+
+    k: int = 3
+    target_com: Optional[np.ndarray]
+
+    def __init__(
+        self,
+        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.target_com = None
+
+        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(
+                f"{self.__class__.__name__} cost must be a vector of shape (1,) "
+                f"(aka identical cost for all coordinates) or ({self.k},). "
+                f"Got {cost.shape}"
+            )
+        if not np.all(cost >= 0.0):
+            raise TaskDefinitionError(f"{self.__class__.__name__} cost must be >= 0")
+        self.cost[:] = cost
+
+    def set_target(self, target_com: npt.ArrayLike) -> None:
+        """Set the target CoM 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(
+                f"Expected target CoM to have shape ({self.k},) but got "
+                f"{target_com.shape}"
+            )
+        self.target_com = target_com.copy()
+
+    def set_target_from_configuration(self, configuration: Configuration) -> None:
+        """Set the target CoM from a given robot configuration."""
+        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.
+        """
+        if self.target_com is None:
+            raise TargetNotSet(self.__class__.__name__)
+        return configuration.data.subtree_com[1] - self.target_com
+
+    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).
+        """
+        if self.target_com is None:
+            raise TargetNotSet(self.__class__.__name__)
+        jac = np.empty((self.k, configuration.nv))
+        mujoco.mj_jacSubtreeCom(configuration.model, configuration.data, jac, 1)
+        return jac

mink/tasks/exceptions.py

@@ -0,0 +1,27 @@
+"""Exceptions raised by tasks."""
+
+from ..exceptions import MinkError
+
+
+class TaskDefinitionError(MinkError):
+    """Exception raised when a task definition is ill-formed."""
+
+
+class TargetNotSet(MinkError):
+    """Exception raised when attempting to use a task with an unset target."""
+
+    def __init__(self, cls_name: str):
+        message = f"No target set for {cls_name}"
+        super().__init__(message)
+
+
+class InvalidTarget(MinkError):
+    """Exception raised when the target is invalid."""
+
+
+class InvalidGain(MinkError):
+    """Exception raised when the gain is outside the valid range."""
+
+
+class InvalidDamping(MinkError):
+    """Exception raised when the damping is outside the valid range."""

mink/tasks/frame_task.py

@@ -0,0 +1,114 @@
+"""Frame task implementation."""
+
+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 FrameTask(Task):
+    """Regulate the pose of a robot frame 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.
+    """
+
+    k: int = 6
+    transform_target_to_world: Optional[SE3]
+
+    def __init__(
+        self,
+        frame_name: str,
+        frame_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.position_cost = position_cost
+        self.orientation_cost = orientation_cost
+        self.transform_target_to_world = 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_world: SE3) -> None:
+        """Set the target pose in 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."""
+        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.
+
+        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`.
+        """
+        if self.transform_target_to_world is None:
+            raise TargetNotSet(self.__class__.__name__)
+
+        transform_frame_to_world = configuration.get_transform_frame_to_world(
+            self.frame_name, self.frame_type
+        )
+        return self.transform_target_to_world.minus(transform_frame_to_world)
+
+    def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
+        """Compute the frame task Jacobian.
+
+        The task Jacobian is the derivative of the task error with respect to the
+        current configuration. It has dimension (6, nv).
+        """
+        if self.transform_target_to_world is None:
+            raise TargetNotSet(self.__class__.__name__)
+
+        jac = configuration.get_frame_jacobian(self.frame_name, self.frame_type)
+
+        transform_frame_to_world = configuration.get_transform_frame_to_world(
+            self.frame_name, self.frame_type
+        )
+
+        T_tb = self.transform_target_to_world.inverse() @ transform_frame_to_world
+        return -T_tb.jlog() @ jac

mink/tasks/posture_task.py

@@ -0,0 +1,103 @@
+from typing import Optional
+
+import mujoco
+import numpy as np
+import numpy.typing as npt
+
+from ..configuration import Configuration
+from ..utils import get_freejoint_dims
+from .exceptions import InvalidTarget, TargetNotSet, TaskDefinitionError
+from .task import Task
+
+
+class PostureTask(Task):
+    """Regulate joint angles to a desired posture.
+
+    A posture is a vector of actuated joint angles. Floating-base coordinates are not
+    affected by this task.
+
+    Attributes:
+        target_q: Target configuration.
+    """
+
+    target_q: Optional[np.ndarray]
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        cost: float,
+        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),
+            gain=gain,
+            lm_damping=lm_damping,
+        )
+        self.target_q = None
+
+        self._v_ids: np.ndarray | None
+        _, v_ids_or_none = get_freejoint_dims(model)
+        if v_ids_or_none:
+            self._v_ids = np.asarray(v_ids_or_none)
+        else:
+            self._v_ids = None
+
+        self.k = model.nv
+        self.nq = model.nq
+
+    def set_target(self, target_q: npt.ArrayLike) -> None:
+        """Set the target posture."""
+        target_q = np.atleast_1d(target_q)
+        if target_q.ndim != 1 or target_q.shape[0] != (self.nq):
+            raise InvalidTarget(
+                f"Expected target posture to have shape ({self.nq},) but got "
+                f"{target_q.shape}"
+            )
+        self.target_q = target_q.copy()
+
+    def set_target_from_configuration(self, configuration: Configuration) -> None:
+        """Set the target posture from the current configuration."""
+        self.set_target(configuration.q)
+
+    def compute_error(self, configuration: Configuration) -> np.ndarray:
+        """Compute the posture task error.
+
+        The error is defined as the right minus error between the target posture and
+        the current posture: `q^* ⊖ q`.
+        """
+        if self.target_q is None:
+            raise TargetNotSet(self.__class__.__name__)
+
+        # NOTE: mj_differentiatePos calculates qpos2 ⊖ qpos1.
+        qvel = np.empty(configuration.nv)
+        mujoco.mj_differentiatePos(
+            m=configuration.model,
+            qvel=qvel,
+            dt=1.0,
+            qpos1=configuration.q,
+            qpos2=self.target_q,
+        )
+
+        if self._v_ids is not None:
+            qvel[self._v_ids] = 0.0
+
+        return qvel
+
+    def compute_jacobian(self, configuration: Configuration) -> np.ndarray:
+        """Compute the posture task Jacobian.
+
+        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).
+        """
+        if self.target_q is None:
+            raise TargetNotSet(self.__class__.__name__)
+
+        jac = -np.eye(configuration.nv)
+        if self._v_ids is not None:
+            jac[:, self._v_ids] = 0.0
+
+        return jac

mink/tasks/task.py

@@ -0,0 +1,105 @@
+import abc
+from typing import NamedTuple
+
+import numpy as np
+
+from ..configuration import Configuration
+from .exceptions import InvalidDamping, InvalidGain
+
+
+class Objective(NamedTuple):
+    """Quadratic objective function in the form 0.5 x^T H x + c^T x."""
+
+    H: np.ndarray  # (nv, nv)
+    c: np.ndarray  # (nv,)
+
+    def value(self, x: np.ndarray) -> float:
+        """Returns the value of the objective at the input vector."""
+        return x.T @ self.H @ x + self.c @ x
+
+
+class Task(abc.ABC):
+    """Abstract base class for kinematic tasks."""
+
+    def __init__(
+        self,
+        cost: np.ndarray,
+        gain: float = 1.0,
+        lm_damping: float = 0.0,
+    ):
+        """Constructor.
+
+        Args:
+            cost: Cost vector with the same dimension as the error of the task.
+            gain: Task gain alpha in [0, 1] for additional low-pass filtering. Defaults
+                to 1.0 (no filtering) for dead-beat control.
+            lm_damping: Unitless scale of the Levenberg-Marquardt (only when the error
+            is large) regularization term, which helps when targets are infeasible.
+            Increase this value if the task is too jerky under unfeasible targets, but
+            beware that a larger damping slows down the task.
+        """
+        if not 0.0 <= gain <= 1.0:
+            raise InvalidGain("`gain` must be in the range [0, 1]")
+
+        if not lm_damping >= 0.0:
+            raise InvalidDamping("`lm_damping` must be >= 0")
+
+        self.cost = cost
+        self.gain = gain
+        self.lm_damping = lm_damping
+
+    @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)
+
+        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.
+
+        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.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    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.
+        """
+        raise NotImplementedError
+
+    def compute_qp_objective(self, configuration: Configuration) -> Objective:
+        """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
+        of the IK is:
+
+        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.
+        """
+        jacobian = self.compute_jacobian(configuration)  # (k, nv)
+        minus_gain_error = -self.gain * self.compute_error(configuration)  # (k,)
+
+        weight = np.diag(self.cost)
+        weighted_jacobian = weight @ jacobian
+        weighted_error = weight @ minus_gain_error
+
+        mu = self.lm_damping * weighted_error @ weighted_error
+        eye_tg = np.eye(configuration.model.nv)
+
+        H = weighted_jacobian.T @ weighted_jacobian + mu * eye_tg  # (nv, nv)
+        c = -weighted_error.T @ weighted_jacobian  # (nv,)
+
+        return Objective(H, c)

mink/utils.py

@@ -0,0 +1,128 @@
+from typing import Optional
+
+import mujoco
+import numpy as np
+
+from . import constants as consts
+from .exceptions import InvalidKeyframe, InvalidMocapBody
+
+
+def move_mocap_to_frame(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    mocap_name: str,
+    frame_name: str,
+    frame_type: str,
+) -> None:
+    """Initialize mocap body pose at a desired frame.
+
+    Args:
+        model: An instance of mujoco.MjModel.
+        data: An instance of mujoco.MjData.
+        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".
+    """
+    mocap_id = model.body(mocap_name).mocapid[0]
+    if mocap_id == -1:
+        raise InvalidMocapBody(mocap_name, model)
+
+    obj_id = mujoco.mj_name2id(model, consts.FRAME_TO_ENUM[frame_type], frame_name)
+    xpos = getattr(data, consts.FRAME_TO_POS_ATTR[frame_type])[obj_id]
+    xmat = getattr(data, consts.FRAME_TO_XMAT_ATTR[frame_type])[obj_id]
+
+    data.mocap_pos[mocap_id] = xpos.copy()
+    mujoco.mju_mat2Quat(data.mocap_quat[mocap_id], xmat)
+
+
+def get_freejoint_dims(model: mujoco.MjModel) -> tuple[list[int], list[int]]:
+    """Get all floating joint configuration and tangent indices."""
+    q_ids: list[int] = []
+    v_ids: list[int] = []
+    for j in range(model.njnt):
+        if model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE:
+            qadr = model.jnt_qposadr[j]
+            vadr = model.jnt_dofadr[j]
+            q_ids.extend(range(qadr, qadr + 7))
+            v_ids.extend(range(vadr, vadr + 6))
+    return q_ids, v_ids
+
+
+def custom_configuration_vector(
+    model: mujoco.MjModel,
+    key_name: Optional[str] = None,
+    **kwargs,
+) -> np.ndarray:
+    """Generate a configuration vector where named joints have specific values.
+
+    Args:
+        model: An MjModel instance.
+        key_name: Optional keyframe name to initialize the configuration vector from.
+            Otherwise, the default pose qpos0 is used.
+        kwargs: Custom values for joint coordinates.
+
+    Returns:
+        Configuration vector where named joints have the values specified in
+            keyword arguments, and other joints have their neutral value or value
+            defined in the keyframe if provided.
+    """
+    data = mujoco.MjData(model)
+    if key_name is not None:
+        key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, key_name)
+        if key_id == -1:
+            raise InvalidKeyframe(key_name, model)
+        mujoco.mj_resetDataKeyframe(model, data, key_id)
+    else:
+        mujoco.mj_resetData(model, data)
+    q = data.qpos.copy()
+    for name, value in kwargs.items():
+        jid = model.joint(name).id
+        jnt_dim = consts.qpos_width(model.jnt_type[jid])
+        qid = model.jnt_qposadr[jid]
+        value = np.atleast_1d(value)
+        if value.shape != (jnt_dim,):
+            raise ValueError(
+                f"Joint {name} should have a qpos value of {jnt_dim,} but "
+                f"got {value.shape}"
+            )
+        q[qid : qid + jnt_dim] = value
+    return q
+
+
+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.
+        body_id: ID of body where subtree starts.
+
+    Returns:
+        A list containing all subtree geom ids.
+    """
+
+    def gather_geoms(body_id: int) -> list[int]:
+        geoms: list[int] = []
+        geom_start = model.body_geomadr[body_id]
+        geom_end = geom_start + model.body_geomnum[body_id]
+        geoms.extend(range(geom_start, geom_end))
+        children = [i for i in range(model.nbody) if model.body_parentid[i] == body_id]
+        for child_id in children:
+            geoms.extend(gather_geoms(child_id))
+        return geoms
+
+    return gather_geoms(body_id)
+
+
+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.
+        body_id: ID of body.
+
+    Returns:
+        A list containing all body geom ids.
+    """
+    geom_start = model.body_geomadr[body_id]
+    geom_end = geom_start + model.body_geomnum[body_id]
+    return list(range(geom_start, geom_end))

mink-0.0.1.dist-info/METADATA

@@ -0,0 +1,71 @@
+Metadata-Version: 2.1
+Name: mink
+Version: 0.0.1
+Summary: mink: MuJoCo inverse kinematics.
+Keywords: inverse,kinematics,mujoco
+Author-email: Kevin Zakka <zakka@berkeley.edu>
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Framework :: Robot Framework :: Library
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Requires-Dist: mujoco >= 3.1.6
+Requires-Dist: qpsolvers >= 4.3.1
+Requires-Dist: quadprog >= 0.1.11
+Requires-Dist: typing_extensions
+Requires-Dist: mypy ; extra == "dev"
+Requires-Dist: ruff ; extra == "dev"
+Requires-Dist: black ; extra == "dev"
+Requires-Dist: mink[test] ; extra == "dev"
+Requires-Dist: loop-rate-limiters >= 0.1.0 ; extra == "examples"
+Requires-Dist: dm_control >= 1.0.20 ; extra == "examples"
+Requires-Dist: absl-py ; extra == "test"
+Requires-Dist: pytest ; extra == "test"
+Requires-Dist: robot_descriptions >= 1.9.0 ; extra == "test"
+Project-URL: Changelog, https://github.com/kevinzakka/mink/blob/main/CHANGELOG.md
+Project-URL: Source, https://github.com/kevinzakka/mink
+Project-URL: Tracker, https://github.com/kevinzakka/mink/issues
+Provides-Extra: dev
+Provides-Extra: examples
+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)
+
+mink is a library for differential inverse kinematics in Python, based on the [MuJoCo](https://github.com/google-deepmind/mujoco) physics engine.
+
+## Installation
+
+```bash
+pip install -e ".[examples]"
+```
+
+## Examples
+
+* 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]()
+
+Check out the [examples](https://github.com/kevinzakka/mink/blob/main/examples/) directory for more code.
+
+## How can I help?
+
+Install the library, use it and report any bugs in the [issue tracker](https://github.com/kevinzakka/mink/issues) if you find any. If you're feeling adventurous, you can also check out the contributing [guidelines](CONTRIBUTING.md) and submit a pull request.
+
+## Acknowledgements
+
+mink is a direct port of [pink](https://github.com/stephane-caron/pink) which uses [Pinocchio](https://github.com/stack-of-tasks/pinocchio) under the hood. Stéphane Caron, the original author of pink, is a role model for open-source software in robotics. This library would not have been possible without his work and assistance throughout this project.
+
+mink also heavily adapts code from the following libraries:
+
+* The lie algebra library that powers the transforms in mink is adapted from [jaxlie](https://github.com/brentyi/jaxlie).
+* The collision avoidance constraint is adapted from [dm_robotics](https://github.com/google-deepmind/dm_robotics/tree/main/cpp/controllers)'s LSQP controller.
+