mink 0.0.1__py3-none-any.whl

mink/limits/collision_avoidance_limit.py

@@ -0,0 +1,274 @@
+"""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
+"""
+
+import itertools
+from dataclasses import dataclass
+from typing import List, Sequence, Union
+
+import mujoco
+import numpy as np
+
+from ..configuration import Configuration
+from .limit import Constraint, Limit
+
+# Type aliases.
+Geom = Union[int, str]
+GeomSequence = Sequence[Geom]
+CollisionPair = tuple[GeomSequence, GeomSequence]
+CollisionPairs = Sequence[CollisionPair]
+
+
+@dataclass(frozen=True)
+class Contact:
+    dist: float
+    fromto: np.ndarray
+    geom1: int
+    geom2: int
+    distmax: float
+
+    @property
+    def normal(self) -> np.ndarray:
+        normal = self.fromto[3:] - self.fromto[:3]
+        return normal / (np.linalg.norm(normal) + 1e-9)
+
+    @property
+    def inactive(self) -> bool:
+        return self.dist == self.distmax and not self.fromto.any()
+
+
+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."""
+    body1 = model.geom_bodyid[geom_id1]
+    body2 = model.geom_bodyid[geom_id2]
+    weld1 = model.body_weldid[body1]
+    weld2 = model.body_weldid[body2]
+    return weld1 == weld2
+
+
+def _are_geom_bodies_parent_child(
+    model: mujoco.MjModel, geom_id1: int, geom_id2: int
+) -> bool:
+    """Returns true if the geom bodies have a parent-child relationship."""
+    body_id1 = model.geom_bodyid[geom_id1]
+    body_id2 = model.geom_bodyid[geom_id2]
+
+    # body_weldid is the ID of the body's weld.
+    body_weldid1 = model.body_weldid[body_id1]
+    body_weldid2 = model.body_weldid[body_id2]
+
+    # weld_parent_id is the ID of the parent of the body's weld.
+    weld_parent_id1 = model.body_parentid[body_weldid1]
+    weld_parent_id2 = model.body_parentid[body_weldid2]
+
+    # weld_parent_weldid is the weld ID of the parent of the body's weld.
+    weld_parent_weldid1 = model.body_weldid[weld_parent_id1]
+    weld_parent_weldid2 = model.body_weldid[weld_parent_id2]
+
+    cond1 = body_weldid1 == weld_parent_weldid2
+    cond2 = body_weldid2 == weld_parent_weldid1
+    return cond1 or cond2
+
+
+def _is_pass_contype_conaffinity_check(
+    model: mujoco.MjModel, geom_id1: int, geom_id2: int
+) -> bool:
+    """Returns true if the geoms pass the contype/conaffinity check."""
+    cond1 = bool(model.geom_contype[geom_id1] & model.geom_conaffinity[geom_id2])
+    cond2 = bool(model.geom_contype[geom_id2] & model.geom_conaffinity[geom_id1])
+    return cond1 or cond2
+
+
+class CollisionAvoidanceLimit(Limit):
+    """Normal velocity limit between geom pairs."""
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        geom_pairs: CollisionPairs,
+        gain: float = 0.85,
+        minimum_distance_from_collisions: float = 0.005,
+        collision_detection_distance: float = 0.01,
+        bound_relaxation: float = 0.0,
+    ):
+        """Initialize collision avoidance limit.
+
+        Args:
+            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 collision pair, the mapper 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 to tighten or relax
+        """
+        self.model = model
+        self.gain = gain
+        self.minimum_distance_from_collisions = minimum_distance_from_collisions
+        self.collision_detection_distance = collision_detection_distance
+        self.bound_relaxation = bound_relaxation
+        self.geom_id_pairs = self._construct_geom_id_pairs(geom_pairs)
+        self.max_num_contacts = len(self.geom_id_pairs)
+
+    def compute_qp_inequalities(
+        self,
+        configuration: Configuration,
+        dt: float,
+    ) -> Constraint:
+        upper_bound = np.full((self.max_num_contacts,), np.inf)
+        coefficient_matrix = np.zeros((self.max_num_contacts, self.model.nv))
+        for idx, (geom1_id, geom2_id) in enumerate(self.geom_id_pairs):
+            contact = self._compute_contact_with_minimum_distance(
+                configuration.data, geom1_id, geom2_id
+            )
+            if contact.inactive:
+                continue
+            hi_bound_dist = contact.dist
+            if hi_bound_dist > self.minimum_distance_from_collisions:
+                dist = hi_bound_dist - self.minimum_distance_from_collisions
+                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)
+            coefficient_matrix[idx] = -jac
+        return Constraint(G=coefficient_matrix, h=upper_bound)
+
+    # Private methods.
+
+    def _compute_contact_with_minimum_distance(
+        self, data: mujoco.MjData, geom1_id: int, geom2_id: int
+    ) -> Contact:
+        """Returns the smallest signed distance between a geom pair."""
+        fromto = np.empty(6)
+        dist = mujoco.mj_geomDistance(
+            self.model,
+            data,
+            geom1_id,
+            geom2_id,
+            self.collision_detection_distance,
+            fromto,
+        )
+        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)."""
+        list_of_int: list[int] = []
+        for g in geom_list:
+            if isinstance(g, int):
+                list_of_int.append(g)
+            else:
+                assert isinstance(g, str)
+                list_of_int.append(self.model.geom(g).id)
+        return list_of_int
+
+    def _collision_pairs_to_geom_id_pairs(self, collision_pairs: CollisionPairs):
+        geom_id_pairs = []
+        for collision_pair in collision_pairs:
+            id_pair_A = self._homogenize_geom_id_list(collision_pair[0])
+            id_pair_B = self._homogenize_geom_id_list(collision_pair[1])
+            id_pair_A = list(set(id_pair_A))
+            id_pair_B = list(set(id_pair_B))
+            geom_id_pairs.append((id_pair_A, id_pair_B))
+        return geom_id_pairs
+
+    def _construct_geom_id_pairs(self, geom_pairs):
+        """Returns a set of geom ID pairs for all possible geom-geom collisions.
+
+        The contacts are added based on the following heuristics:
+            1) Geoms that are part of the same body or weld are not included.
+            2) Geoms where the body of one geom is a parent of the body of the other
+                geom are not included.
+            3) Geoms that fail the contype-conaffinity check are ignored.
+
+        Note:
+            1) If two bodies are kinematically welded together (no joints between them)
+                they are considered to be the same body within this function.
+        """
+        geom_id_pairs = []
+        for id_pair in self._collision_pairs_to_geom_id_pairs(geom_pairs):
+            for geom_a, geom_b in itertools.product(*id_pair):
+                weld_body_cond = not _is_welded_together(self.model, geom_a, geom_b)
+                parent_child_cond = not _are_geom_bodies_parent_child(
+                    self.model, geom_a, geom_b
+                )
+                contype_conaffinity_cond = _is_pass_contype_conaffinity_check(
+                    self.model, geom_a, geom_b
+                )
+                if weld_body_cond and parent_child_cond and contype_conaffinity_cond:
+                    geom_id_pairs.append((min(geom_a, geom_b), max(geom_a, geom_b)))
+        return geom_id_pairs

mink/limits/configuration_limit.py

@@ -0,0 +1,126 @@
+"""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]
+"""
+
+import mujoco
+import numpy as np
+
+from ..configuration import Configuration
+from ..constants import qpos_width
+from .exceptions import LimitDefinitionError
+from .limit import Constraint, Limit
+
+
+class ConfigurationLimit(Limit):
+    """Limit for joint positions in a model.
+
+    Floating base joints (joint type="free") are ignored.
+
+    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.
+    """
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        gain: float = 0.95,
+        min_distance_from_limits: float = 0.0,
+    ):
+        """Initialize configuration limits.
+
+        Args:
+            model: MuJoCo model.
+            gain: Gain factor in (0, 1] that determines how fast each joint is
+                allowed to move towards the joint limits at each timestep. Values lower
+                ttan 1 are safer but may make the joints move slowly.
+            min_distance_from_limits: Offset in meters (slide joints) or radians
+                (hinge joints) to be added to the limits. Positive values decrease the
+                range of motion, negative values increase it (i.e. negative values
+                allow penetration).
+        """
+        if not 0.0 < gain <= 1.0:
+            raise LimitDefinitionError(
+                f"{self.__class__.__name__} gain must be in the range (0, 1]"
+            )
+
+        index_list: list[int] = []
+        lower = np.full(model.nq, -np.inf)
+        upper = np.full(model.nq, np.inf)
+        for jnt in range(model.njnt):
+            jnt_type = model.jnt_type[jnt]
+            qpos_dim = qpos_width(jnt_type)
+            jnt_range = model.jnt_range[jnt]
+            padr = model.jnt_qposadr[jnt]
+            if jnt_type == mujoco.mjtJoint.mjJNT_FREE or not model.jnt_limited[jnt]:
+                continue
+            lower[padr : padr + qpos_dim] = jnt_range[0] + min_distance_from_limits
+            upper[padr : padr + qpos_dim] = jnt_range[1] - min_distance_from_limits
+            index_list.append(jnt)
+
+        self.indices = np.array(index_list)
+        self.indices.setflags(write=False)
+
+        dim = len(self.indices)
+        self.projection_matrix = np.eye(model.nv)[self.indices] if dim > 0 else None
+
+        self.lower = lower
+        self.upper = upper
+        self.model = model
+        self.gain = gain
+
+    def compute_qp_inequalities(
+        self,
+        configuration: Configuration,
+        dt: float,
+    ) -> Constraint:
+        del dt  # Unused.
+
+        # Upper.
+        delta_q_max = np.zeros(self.model.nv)
+        mujoco.mj_differentiatePos(
+            m=self.model,
+            qvel=delta_q_max,
+            dt=1.0,
+            qpos1=configuration.q,
+            qpos2=self.upper,
+        )
+
+        # Lower.
+        delta_q_min = np.zeros(self.model.nv)
+        mujoco.mj_differentiatePos(
+            m=self.model,
+            qvel=delta_q_min,
+            dt=1.0,
+            # NOTE: mujoco.mj_differentiatePos does `qpos2 - qpos1` so notice the order
+            # swap here compared to above.
+            qpos1=self.lower,
+            qpos2=configuration.q,
+        )
+
+        p_min = self.gain * delta_q_min[self.indices]
+        p_max = self.gain * delta_q_max[self.indices]
+        G = np.vstack([self.projection_matrix, -self.projection_matrix])
+        h = np.hstack([p_max, p_min])
+        return Constraint(G=G, h=h)

mink/limits/exceptions.py

@@ -0,0 +1,7 @@
+"""Exceptions raised by limits."""
+
+from ..exceptions import MinkError
+
+
+class LimitDefinitionError(MinkError):
+    """Exception raised when a limit definition is ill-formed."""

mink/limits/limit.py

@@ -0,0 +1,41 @@
+import abc
+from typing import NamedTuple, Optional
+
+import numpy as np
+
+from ..configuration import Configuration
+
+
+class Constraint(NamedTuple):
+    """Linear inequalities in the form G(q) dq <= h(q).
+
+    The limit is considered inactive when G or h are None.
+    """
+
+    G: Optional[np.ndarray] = None  # (nv, nv)
+    h: Optional[np.ndarray] = None  # (nv,)
+
+    @property
+    def inactive(self) -> bool:
+        return self.G is None or self.h is None
+
+
+class Limit(abc.ABC):
+    """Abstract base class for kinematic limits."""
+
+    @abc.abstractmethod
+    def compute_qp_inequalities(
+        self,
+        configuration: Configuration,
+        dt: float,
+    ) -> Constraint:
+        """Compute limit as linearized QP inequalities.
+
+        Args:
+            configuration: Current configuration.
+            dt: Integration time step in [s].
+
+        Returns:
+            Pair (G, h) representing the inequality constraint.
+        """
+        raise NotImplementedError

mink/limits/velocity_limit.py

@@ -0,0 +1,92 @@
+"""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]
+"""
+
+from typing import Mapping
+
+import mujoco
+import numpy as np
+import numpy.typing as npt
+
+from ..configuration import Configuration
+from ..constants import dof_width
+from .exceptions import LimitDefinitionError
+from .limit import Constraint, Limit
+
+
+class VelocityLimit(Limit):
+    """Limit for joint velocities in a model.
+
+    Floating base joints are ignored.
+
+    Attributes:
+        indices: Tangent indices corresponding to velocity-limited joints.
+        limit: Maximum allowed velocity magnitude for velocity-limited joints.
+        projection_matrix: Projection from tangent space to subspace with
+            velocity-limited joints.
+    """
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        velocities: Mapping[str, npt.ArrayLike] = {},
+    ):
+        """Initialize velocity limits.
+
+        Args:
+            velocities: Dictionary mapping joint name to maximum allowed magnitude in
+                [m]/[s] for slide joints and [rad]/[s] for hinge joints.
+        """
+        limit_list: list[float] = []
+        index_list: list[int] = []
+        for joint_name, max_vel in velocities.items():
+            jid = model.joint(joint_name).id
+            jnt_type = model.jnt_type[jid]
+            jnt_dim = dof_width(jnt_type)
+            jnt_id = model.jnt_dofadr[jid]
+            if jnt_type == mujoco.mjtJoint.mjJNT_FREE:
+                raise LimitDefinitionError(f"Free joint {joint_name} is not supported")
+            max_vel = np.atleast_1d(max_vel)
+            if max_vel.shape != (jnt_dim,):
+                raise LimitDefinitionError(
+                    f"Joint {joint_name} must have a limit of shape ({jnt_dim},). "
+                    f"Got: {max_vel.shape}"
+                )
+            index_list.extend(range(jnt_id, jnt_id + jnt_dim))
+            limit_list.extend(max_vel.tolist())
+
+        self.indices = np.array(index_list)
+        self.indices.setflags(write=False)
+        self.limit = np.array(limit_list)
+        self.limit.setflags(write=False)
+
+        dim = len(self.indices)
+        self.projection_matrix = np.eye(model.nv)[self.indices] if dim > 0 else None
+
+    def compute_qp_inequalities(
+        self, configuration: Configuration, dt: float
+    ) -> Constraint:
+        del configuration  # Unused.
+        if self.projection_matrix is None:
+            return Constraint()
+        G = np.vstack([self.projection_matrix, -self.projection_matrix])
+        h = np.hstack([dt * self.limit, dt * self.limit])
+        return Constraint(G=G, h=h)

mink/solve_ik.py

@@ -0,0 +1,105 @@
+"""Build and solve the inverse kinematics problem."""
+
+from typing import Optional, Sequence
+
+import numpy as np
+import qpsolvers
+
+from .configuration import Configuration
+from .limits import ConfigurationLimit, Limit
+from .tasks import Objective, Task
+
+
+def _compute_qp_objective(
+    configuration: Configuration, tasks: Sequence[Task], damping: float
+) -> Objective:
+    H = np.eye(configuration.model.nv) * damping
+    c = np.zeros(configuration.model.nv)
+    for task in tasks:
+        H_task, c_task = task.compute_qp_objective(configuration)
+        H += H_task
+        c += c_task
+    return Objective(H, c)
+
+
+def _compute_qp_inequalities(
+    configuration: Configuration, limits: Optional[Sequence[Limit]], dt: float
+) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
+    if limits is None:
+        limits = [ConfigurationLimit(configuration.model)]
+    G_list = []
+    h_list = []
+    for limit in limits:
+        inequality = limit.compute_qp_inequalities(configuration, dt)
+        if not inequality.inactive:
+            assert inequality.G is not None and inequality.h is not None  # mypy.
+            G_list.append(inequality.G)
+            h_list.append(inequality.h)
+    if not G_list:
+        return None, None
+    return np.vstack(G_list), np.hstack(h_list)
+
+
+def build_ik(
+    configuration: Configuration,
+    tasks: Sequence[Task],
+    dt: float,
+    damping: float = 1e-12,
+    limits: Optional[Sequence[Limit]] = None,
+) -> qpsolvers.Problem:
+    """Build quadratic program from current configuration and tasks.
+
+    Args:
+        configuration: Robot configuration.
+        tasks: List of kinematic tasks.
+        dt: Integration timestep in [s].
+        damping: Levenberg-Marquardt damping.
+        limits: List of limits to enforce. Set to empty list to disable. If None,
+            defaults to a configuration limit.
+
+    Returns:
+        Quadratic program of the inverse kinematics problem.
+    """
+    P, q = _compute_qp_objective(configuration, tasks, damping)
+    G, h = _compute_qp_inequalities(configuration, limits, dt)
+    return qpsolvers.Problem(P, q, G, h)
+
+
+def solve_ik(
+    configuration: Configuration,
+    tasks: Sequence[Task],
+    dt: float,
+    solver: str,
+    damping: float = 1e-12,
+    safety_break: bool = False,
+    limits: Optional[Sequence[Limit]] = None,
+    **kwargs,
+) -> np.ndarray:
+    """Solve the differential inverse kinematics problem.
+
+    Computes a velocity tangent to the current robot configuration. The computed
+    velocity satisfies at (weighted) best the set of provided kinematic tasks.
+
+    Args:
+        configuration: Robot configuration.
+        tasks: List of kinematic tasks.
+        dt: Integration timestep in [s].
+        solver: Backend quadratic programming (QP) solver.
+        damping: Levenberg-Marquardt damping.
+        safety_break: If True, stop execution and raise an exception if
+            the current configuration is outside limits. If False, print a
+            warning and continue execution.
+        limits: List of limits to enforce. Set to empty list to disable. If None,
+            defaults to a configuration limit.
+        kwargs: Keyword arguments to forward to the backend QP solver.
+
+    Returns:
+        Velocity `v` in tangent space.
+    """
+    configuration.check_limits(safety_break=safety_break)
+    problem = build_ik(configuration, tasks, dt, damping, limits)
+    result = qpsolvers.solve_problem(problem, solver=solver, **kwargs)
+    dq = result.x
+    assert dq is not None
+    v: np.ndarray = dq / dt
+    return v

mink/tasks/__init__.py

@@ -0,0 +1,26 @@
+"""Kinematic tasks."""
+
+from .com_task import ComTask
+from .exceptions import (
+    InvalidDamping,
+    InvalidGain,
+    InvalidTarget,
+    TargetNotSet,
+    TaskDefinitionError,
+)
+from .frame_task import FrameTask
+from .posture_task import PostureTask
+from .task import Objective, Task
+
+__all__ = (
+    "ComTask",
+    "FrameTask",
+    "Objective",
+    "PostureTask",
+    "Task",
+    "TargetNotSet",
+    "InvalidTarget",
+    "TaskDefinitionError",
+    "InvalidGain",
+    "InvalidDamping",
+)