Docstrings.

examples/arm_ur5e.py

@@ -1,9 +1,11 @@
+from pathlib import Path
+
 import mujoco
 import mujoco.viewer
 import numpy as np
-import mink
-from pathlib import Path
 from loop_rate_limiters import RateLimiter
+
+import mink
 from mink.utils import set_mocap_pose_from_site
 
 _HERE = Path(__file__).parent

examples/dual_iiwa.py

@@ -1,14 +1,15 @@
 """Task adapted from https://github.com/stephane-caron/pink/pull/94."""
 
-from dm_control import mjcf
+from pathlib import Path
+
 import mujoco
 import mujoco.viewer
 import numpy as np
-import mink
-from pathlib import Path
+from dm_control import mjcf
 from loop_rate_limiters import RateLimiter
-from mink.utils import set_mocap_pose_from_site
 
+import mink
+from mink.utils import set_mocap_pose_from_site
 
 _HERE = Path(__file__).parent
 _XML = _HERE / "kuka_iiwa_14" / "iiwa14.xml"
@@ -95,7 +96,7 @@ def construct_model():
         mink.CollisionAvoidanceLimit(
             model=model,
             geom_pairs=collision_pairs,
-            minimum_distance_from_collisions=0.1,
+            minimum_distance_from_collisions=0.01,
             collision_detection_distance=0.2,
         ),
     ]
@@ -104,7 +105,7 @@ def construct_model():
     right_mid = model.body("r_target").mocapid[0]
     model = configuration.model
     data = configuration.data
-    solver = "quadprog"
+    solver = "osqp"
 
     l_y_des = np.array([0.392, -0.392, 0.6])
     r_y_des = np.array([0.392, 0.392, 0.6])

examples/flying_dual_arm_ur5e.py

@@ -1,12 +1,13 @@
-from dm_control import mjcf
+from pathlib import Path
+
 import mujoco
 import mujoco.viewer
 import numpy as np
-import mink
-from pathlib import Path
+from dm_control import mjcf
 from loop_rate_limiters import RateLimiter
-from mink.utils import set_mocap_pose_from_site
 
+import mink
+from mink.utils import set_mocap_pose_from_site
 
 _HERE = Path(__file__).parent
 _XML = _HERE / "universal_robots_ur5e" / "ur5e.xml"

examples/humanoid_g1.py

@@ -1,8 +1,10 @@
+from pathlib import Path
+
 import mujoco
 import mujoco.viewer
-import mink
-from pathlib import Path
 from loop_rate_limiters import RateLimiter
+
+import mink
 from mink.utils import set_mocap_pose_from_site
 
 _HERE = Path(__file__).parent

examples/main.py

@@ -0,0 +1,94 @@
+from pathlib import Path
+
+import mujoco
+import mujoco.viewer
+import numpy as np
+from loop_rate_limiters import RateLimiter
+
+import mink
+from mink.lie import SE3, SO3
+from mink.utils import set_mocap_pose_from_site
+
+_HERE = Path(__file__).parent
+_XML = _HERE / "universal_robots_ur5e" / "scene.xml"
+
+
+if __name__ == "__main__":
+    model = mujoco.MjModel.from_xml_path(_XML.as_posix())
+    data = mujoco.MjData(model)
+
+    ## =================== ##
+    ## Setup IK.
+    ## =================== ##
+
+    configuration = mink.Configuration(model)
+
+    tasks = [
+        end_effector_task := mink.FrameTask(
+            frame_name="attachment_site",
+            frame_type="site",
+            position_cost=1.0,
+            orientation_cost=1.0,
+            lm_damping=1.0,
+        ),
+    ]
+
+    # Enable collision avoidance between the following geoms:
+    collision_pairs = [
+        (["wrist_3_link"], ["floor", "wall"]),
+    ]
+
+    limits = [
+        mink.ConfigurationLimit(model=configuration.model),
+        mink.VelocityLimit(configuration.model, np.full_like(configuration.q, np.pi)),
+        mink.CollisionAvoidanceLimit(
+            model=configuration.model, geom_pairs=collision_pairs
+        ),
+    ]
+
+    ## =================== ##
+
+    mid = model.body("target").mocapid[0]
+
+    # IK settings.
+    solver = "quadprog"
+    pos_threshold = 1e-4
+    ori_threshold = 1e-4
+    max_iters = 20
+
+    with mujoco.viewer.launch_passive(
+        model=model, data=data, show_left_ui=False, show_right_ui=False
+    ) as viewer:
+        mujoco.mjv_defaultFreeCamera(model, viewer.cam)
+
+        mujoco.mj_resetDataKeyframe(model, data, model.key("home").id)
+        configuration.update(data.qpos)
+        mujoco.mj_forward(model, data)
+
+        # Initialize the mocap target at the end-effector site.
+        set_mocap_pose_from_site(model, data, "target", "attachment_site")
+
+        rate = RateLimiter(frequency=500.0)
+        while viewer.is_running():
+            target_pos = data.mocap_pos[mid]
+            target_ori = data.mocap_quat[mid]
+            target_pose = SE3.from_rotation_and_translation(SO3(target_ori), target_pos)
+            end_effector_task.set_target(target_pose)
+
+            # Compute velocity and integrate into the next configuration.
+            for i in range(max_iters):
+                vel = mink.solve_ik(configuration, tasks, limits, rate.dt, solver, 1e-3)
+                configuration.integrate_inplace(vel, rate.dt)
+                err = end_effector_task.compute_error(configuration)
+                pos_achieved = np.linalg.norm(err[:3]) <= pos_threshold
+                ori_achieved = np.linalg.norm(err[3:]) <= ori_threshold
+                if pos_achieved and ori_achieved:
+                    print(f"Exiting after {i} iterations.")
+                    break
+
+            data.ctrl = configuration.q
+            mujoco.mj_step(model, data)
+
+            # Visualize at fixed FPS.
+            viewer.sync()
+            rate.sleep()

examples/quadruped_go1.py

@@ -1,9 +1,11 @@
+from pathlib import Path
+
 import mujoco
 import mujoco.viewer
-import mink
-from pathlib import Path
 from loop_rate_limiters import RateLimiter
-from mink.utils import set_mocap_pose_from_site, set_mocap_pose_from_body
+
+import mink
+from mink.utils import set_mocap_pose_from_body, set_mocap_pose_from_site
 
 _HERE = Path(__file__).parent
 _XML = _HERE / "unitree_go1" / "scene.xml"

examples/universal_robots_ur5e/scene.xml

@@ -19,7 +19,6 @@
     <geom name="floor" size="1 1 0.01" type="plane" material="grid"/>
     <body name="target" pos="0.5 0 .5" quat="0 1 0 0" mocap="true">
       <geom type="box" size=".05 .05 .05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
-      <site type="sphere" size="0.01" rgba="0 0 1 1" group="1"/>
     </body>
     <body name="wall" pos="0.5 0 0.1">
       <geom name="wall" type="box" size=".1 .1 .1"/>

mink/__init__.py

@@ -1,16 +1,16 @@
 """mink: MuJoCo differential inverse kinematics."""
 
 from mink.configuration import Configuration
-from mink.solve_ik import build_ik, solve_ik
-from mink.tasks import FrameTask, PostureTask, Task, Objective, ComTask
+from mink.lie import SE3, SO3
 from mink.limits import (
+    CollisionAvoidanceLimit,
     ConfigurationLimit,
     Constraint,
     Limit,
     VelocityLimit,
-    CollisionAvoidanceLimit,
 )
-from mink.lie import SO3, SE3
+from mink.solve_ik import build_ik, solve_ik
+from mink.tasks import ComTask, FrameTask, Objective, PostureTask, Task
 
 __version__ = "0.0.1"
 

mink/configuration.py

@@ -1,6 +1,10 @@
-import numpy as np
+from typing import Optional
+
 import mujoco
-from mink.lie import SE3, SO3
+import numpy as np
+
+from .exceptions import FrameNotFound, KeyframeNotFound, UnsupportedFrameType
+from .lie import SE3, SO3
 
 _SUPPORTED_OBJ_TYPES = {"body", "geom", "site"}
 
@@ -28,41 +32,120 @@
 
 
 class Configuration:
+    """A struct that provides convenient access to kinematic quantities such as frame
+    transforms and frame jacobians.
+
+    The `update` function ensures the proper forward kinematics functions have been
+    called, namely:
+
+    - mujoco.mj_kinematics(model, data)
+    - mujoco.mj_comPos(model, data)
+    """
+
     def __init__(
         self,
         model: mujoco.MjModel,
         q: np.ndarray | None = None,
     ):
+        """Constructor.
+
+        Args:
+            model: An instance of MjModel.
+            q: Optional configuration to initialize from. If None, the configuration
+            is initialized to the reference configuration `qpos0`.
+        """
         self.model = model
         self.data = mujoco.MjData(model)
         if q is not None:
             self.data.qpos = q
         self.update()
 
     def update_from_keyframe(self, key: str) -> None:
-        mujoco.mj_resetDataKeyframe(self.model, self.data, self.model.key(key).id)
+        """Update the configuration from a keyframe and run forward kinematics.
+
+        Args:
+            key: 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)
+        if key_id == -1:
+            raise KeyframeNotFound(key, self.model)
+        mujoco.mj_resetDataKeyframe(self.model, self.data, key_id)
         self.update()
 
-    def update(self) -> None:
+    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.
+        """
+        if q is not None:
+            self.data.qpos = q
         mujoco.mj_kinematics(self.model, self.data)
         mujoco.mj_comPos(self.model, self.data)
 
     def get_frame_jacobian(self, frame_name: str, frame_type: str) -> np.ndarray:
-        assert frame_type in _SUPPORTED_OBJ_TYPES
+        """Compute the Jacobian matrix of a frame velocity.
+
+        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:
+
+            B_v_WB = B_J_WB q_dot
+
+        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.
+        """
+        if frame_type not in _SUPPORTED_OBJ_TYPES:
+            raise UnsupportedFrameType(frame_type, _SUPPORTED_OBJ_TYPES)
+
         frame_id = mujoco.mj_name2id(self.model, _TYPE_TO_ENUM[frame_type], frame_name)
         if frame_id == -1:
-            raise ValueError(f"Frame '{frame_name}' not found")
+            raise FrameNotFound(
+                frame_name=frame_name,
+                frame_type=frame_type,
+                model=self.model,
+            )
+
         jac = np.zeros((6, self.model.nv))
         _TYPE_TO_JAC_FUNCTION[frame_type](
             self.model, self.data, jac[:3], jac[3:], frame_id
         )
+        R_sb = getattr(self.data, _TYPE_TO_XMAT_ATTRIBUTE[frame_type])[frame_id]
+        R_sb = R_sb.reshape(3, 3).T
+        jac[3:] = R_sb @ jac[3:]
+        jac[:3] = R_sb @ jac[:3]
         return jac
 
     def get_transform_frame_to_world(self, frame_name: str, frame_type: str) -> SE3:
-        assert frame_type in _SUPPORTED_OBJ_TYPES
+        """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.
+
+        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:
+            The pose of the frame in the world frame.
+        """
+        if frame_type not in _SUPPORTED_OBJ_TYPES:
+            raise UnsupportedFrameType(frame_type, _SUPPORTED_OBJ_TYPES)
+
         frame_id = mujoco.mj_name2id(self.model, _TYPE_TO_ENUM[frame_type], frame_name)
         if frame_id == -1:
-            raise ValueError(f"Frame '{frame_name}' not found")
+            raise FrameNotFound(
+                frame_name=frame_name,
+                frame_type=frame_type,
+                model=self.model,
+            )
+
         xpos = getattr(self.data, _TYPE_TO_POS_ATTRIBUTE[frame_type])[frame_id]
         xmat = getattr(self.data, _TYPE_TO_XMAT_ATTRIBUTE[frame_type])[frame_id]
         return SE3.from_rotation_and_translation(
@@ -71,14 +154,30 @@ def get_transform_frame_to_world(self, frame_name: str, frame_type: str) -> SE3:
         )
 
     def integrate(self, velocity: np.ndarray, dt: float) -> np.ndarray:
+        """Integrate a velocity starting from the current configuration.
+
+        Args:
+            velocity: The velocity.
+            dt: Integration duration in [s].
+
+        Returns:
+            The new configuration after integration.
+        """
         q = self.data.qpos.copy()
         mujoco.mj_integratePos(self.model, q, velocity, dt)
         return q
 
     def integrate_inplace(self, velocity: np.ndarray, dt: float) -> None:
+        """Integrate a velocity and update the current configuration inplace.
+
+        Args:
+            velocity: The velocity.
+            dt: Integration duration in [s].
+        """
         mujoco.mj_integratePos(self.model, self.data.qpos, velocity, dt)
         self.update()
 
     @property
     def q(self) -> np.ndarray:
+        """The current configuration vector."""
         return self.data.qpos.copy()