From 069c39ee5d4fe3104ba2de66adf99183a570f71c Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Fri, 12 Jul 2024 14:18:24 -0400 Subject: [PATCH] Docstrings. --- examples/arm_ur5e.py | 6 +- examples/dual_iiwa.py | 13 +-- examples/flying_dual_arm_ur5e.py | 9 +- examples/humanoid_g1.py | 6 +- examples/main.py | 94 ++++++++++++++++++ examples/quadruped_go1.py | 8 +- examples/universal_robots_ur5e/scene.xml | 1 - mink/__init__.py | 8 +- mink/configuration.py | 115 +++++++++++++++++++++-- mink/exceptions.py | 61 ++++++++++++ mink/lie/__init__.py | 4 +- mink/lie/interpolate.py | 11 --- mink/lie/se3.py | 57 ++++++++++- mink/lie/so3.py | 47 ++++++++- mink/lie/utils.py | 2 +- mink/limits/__init__.py | 4 +- mink/limits/collision_avoidance_limit.py | 5 +- mink/limits/configuration_limit.py | 4 +- mink/limits/limit.py | 4 +- mink/limits/velocity_limit.py | 5 +- mink/solve_ik.py | 8 +- mink/tasks/__init__.py | 5 +- mink/tasks/com_task.py | 4 +- mink/tasks/frame_task.py | 46 +++------ mink/tasks/posture_task.py | 5 +- pyproject.toml | 26 +++++ tests/compare.py | 96 +++++++++++++++++++ tests/test_configuration.py | 12 ++- tests/test_jacobians.py | 5 +- 29 files changed, 556 insertions(+), 115 deletions(-) create mode 100644 examples/main.py create mode 100644 mink/exceptions.py delete mode 100644 mink/lie/interpolate.py create mode 100644 tests/compare.py diff --git a/examples/arm_ur5e.py b/examples/arm_ur5e.py index 83db3fe..1fa7670 100644 --- a/examples/arm_ur5e.py +++ b/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 diff --git a/examples/dual_iiwa.py b/examples/dual_iiwa.py index d3d0767..fbf759b 100644 --- a/examples/dual_iiwa.py +++ b/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]) diff --git a/examples/flying_dual_arm_ur5e.py b/examples/flying_dual_arm_ur5e.py index 76004cf..aeb1bae 100644 --- a/examples/flying_dual_arm_ur5e.py +++ b/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" diff --git a/examples/humanoid_g1.py b/examples/humanoid_g1.py index b1c9c31..463bbcd 100644 --- a/examples/humanoid_g1.py +++ b/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 diff --git a/examples/main.py b/examples/main.py new file mode 100644 index 0000000..558c67d --- /dev/null +++ b/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() diff --git a/examples/quadruped_go1.py b/examples/quadruped_go1.py index af7eece..f120d5a 100644 --- a/examples/quadruped_go1.py +++ b/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" diff --git a/examples/universal_robots_ur5e/scene.xml b/examples/universal_robots_ur5e/scene.xml index aab5099..26758b7 100644 --- a/examples/universal_robots_ur5e/scene.xml +++ b/examples/universal_robots_ur5e/scene.xml @@ -19,7 +19,6 @@ - diff --git a/mink/__init__.py b/mink/__init__.py index 7b15d31..ef461be 100644 --- a/mink/__init__.py +++ b/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" diff --git a/mink/configuration.py b/mink/configuration.py index e00b101..67ae131 100644 --- a/mink/configuration.py +++ b/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,11 +32,28 @@ 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: @@ -40,29 +61,91 @@ def __init__( 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() diff --git a/mink/exceptions.py b/mink/exceptions.py new file mode 100644 index 0000000..0bea4df --- /dev/null +++ b/mink/exceptions.py @@ -0,0 +1,61 @@ +"""Exceptions specific to mink.""" + +import mujoco + + +class MinkError(Exception): + """Base class for Mink exceptions.""" + + +class UnsupportedFrameType(MinkError): + """Exception raised when a frame type is unsupported.""" + + def __init__(self, frame_type: str, supported_types: list[str]): + message = ( + f"{frame_type} is not supported." + f"Supported frame types are: {supported_types}" + ) + super().__init__(message) + + +class FrameNotFound(MinkError): + """Exception raised when a frame is not found in the robot model.""" + + def __init__( + self, + frame_name: str, + frame_type: str, + model: mujoco.MjModel, + ): + if frame_type == "body": + available_names_of_type_frame_type = [ + model.body(i).name for i in range(model.nbody) + ] + elif frame_type == "site": + available_names_of_type_frame_type = [ + model.site(i).name for i in range(model.nsite) + ] + else: + assert frame_type == "geom" + available_names_of_type_frame_type = [ + model.geom(i).name for i in range(model.ngeom) + ] + + message = ( + f"Frame '{frame_name}' does not exist in the model. " + f"Available {frame_type} names: {available_names_of_type_frame_type}" + ) + + super().__init__(message) + + +class KeyframeNotFound(MinkError): + """Exception raised when a keyframe is not found in the robot model.""" + + def __init__(self, keyframe_name: str, model: mujoco.MjModel): + available_keyframes = [model.key(i).name for i in range(model.nkey)] + message = ( + f"Keyframe {keyframe_name} does not exist in the model. " + f"Available keyframe names: {available_keyframes}" + ) + super().__init__(message) diff --git a/mink/lie/__init__.py b/mink/lie/__init__.py index d39104c..5cf3484 100644 --- a/mink/lie/__init__.py +++ b/mink/lie/__init__.py @@ -1,13 +1,11 @@ from mink.lie.se3 import SE3 from mink.lie.so3 import SO3 -from mink.lie.utils import get_epsilon, skew, mat2quat -from mink.lie.interpolate import interpolate +from mink.lie.utils import get_epsilon, mat2quat, skew __all__ = ( "SE3", "SO3", "get_epsilon", - "interpolate", "mat2quat", "skew", ) diff --git a/mink/lie/interpolate.py b/mink/lie/interpolate.py deleted file mode 100644 index 89df5c5..0000000 --- a/mink/lie/interpolate.py +++ /dev/null @@ -1,11 +0,0 @@ -from mink.lie import SE3, SO3 - - -def interpolate( - p0: SE3 | SO3, - p1: SE3 | SO3, - alpha: float = 0.5, -) -> SE3 | SO3: - assert 0.0 <= alpha <= 1.0 - exp_func = getattr(type(p0), "exp") - return p0 @ exp_func(alpha * (p0.inverse() @ p1).log()) diff --git a/mink/lie/se3.py b/mink/lie/se3.py index 39969f2..659b3e2 100644 --- a/mink/lie/se3.py +++ b/mink/lie/se3.py @@ -1,8 +1,9 @@ from __future__ import annotations -import numpy as np from dataclasses import dataclass +import numpy as np + from mink.lie.so3 import SO3 from mink.lie.utils import get_epsilon, skew @@ -121,7 +122,18 @@ def log(self) -> np.ndarray: + (skew_omega @ skew_omega) / 12.0 ) else: - V_inv = np.eye(3) - 0.5 * skew_omega + (1 - theta_safe * np.cos(half_theta_safe) / (2 * np.sin(half_theta_safe))) / (theta_safe * theta_safe) * (skew_omega @ skew_omega) + V_inv = ( + np.eye(3) + - 0.5 * skew_omega + + ( + 1 + - theta_safe + * np.cos(half_theta_safe) + / (2 * np.sin(half_theta_safe)) + ) + / (theta_safe * theta_safe) + * (skew_omega @ skew_omega) + ) V_inv_2 = ( np.eye(3, dtype=np.float64) - 0.5 * skew_omega @@ -142,8 +154,8 @@ def adjoint(self) -> np.ndarray: R = self.rotation().as_matrix() return np.block( [ - [R, np.zeros((3, 3), dtype=np.float64)], - [skew(self.translation()) @ R, R], + [R, skew(self.translation()) @ R], + [np.zeros((3, 3), dtype=np.float64), R], ] ) @@ -180,3 +192,40 @@ def __matmul__(self, other: SE3 | np.ndarray) -> SE3 | np.ndarray: def copy(self) -> SE3: return SE3(wxyz_xyz=self.wxyz_xyz.copy()) + + def jlog(self): + """Derivatve of log(this.inv() * x) by x at x=this.""" + rotation = self.rotation() + translation = self.translation() + jlog_so3 = rotation.jlog() + w = rotation.log() + theta = np.linalg.norm(w) + use_taylor = theta < get_epsilon(theta.dtype) + t2 = theta * theta + tinv = 1 / theta + t2inv = tinv * tinv + st, ct = np.sin(theta), np.cos(theta) + inv_2_2ct = 1 / (2 * (1 - ct)) + beta = np.where(use_taylor, 1 / 12 + t2 / 720, t2inv - st * tinv * inv_2_2ct) + + beta_dot_over_theta = np.where( + use_taylor, + 1 / 360, + -2 * t2inv * t2inv + (1 + st * tinv) * t2inv * inv_2_2ct, + ) + wTp = w @ translation + v3_tmp = (beta_dot_over_theta * wTp) * w - ( + theta**2 * beta_dot_over_theta + 2 * beta + ) * translation + C = ( + np.outer(v3_tmp, w) + + beta * np.outer(w, translation) + + wTp * beta * np.eye(3) + ) + C = C + 0.5 * skew(translation) + B = C @ jlog_so3 + jlog = np.zeros((6, 6)) + jlog[:3, :3] = jlog_so3 + jlog[3:, 3:] = jlog_so3 + jlog[:3, 3:] = B + return jlog diff --git a/mink/lie/so3.py b/mink/lie/so3.py index 9a9334e..757c9fd 100644 --- a/mink/lie/so3.py +++ b/mink/lie/so3.py @@ -1,10 +1,11 @@ from __future__ import annotations -import numpy as np from dataclasses import dataclass + import mujoco +import numpy as np -from mink.lie.utils import get_epsilon +from mink.lie.utils import get_epsilon, skew _IDENTITIY_WXYZ = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64) _INVERT_QUAT_SIGN = np.array([1.0, -1.0, -1.0, -1.0], dtype=np.float64) @@ -192,3 +193,45 @@ def __matmul__(self, other: SO3 | np.ndarray) -> SO3 | np.ndarray: return self.multiply(other=other) else: raise ValueError(f"Unsupported argument type for @ operator: {type(other)}") + + def jlog(self): + """Derivatve of log(this.inv() * x) by x at x=this.""" + theta = self.log() + + theta_squared = np.sum(np.square(theta), axis=-1) + use_taylor = theta_squared < get_epsilon(theta_squared.dtype) + + # Shim to avoid NaNs in np.where branches, which cause failures for + # reverse-mode AD. + theta_squared_safe = np.where( + use_taylor, + np.ones_like(theta_squared), # Any non-zero value should do here. + theta_squared, + ) + del theta_squared + theta_safe = np.sqrt(theta_squared_safe) + half_theta_safe = theta_safe / 2.0 + + skew_omega = skew(theta) + V_inv = np.where( + use_taylor[None, None], + np.eye(3) + - 0.5 * skew_omega + + np.einsum("ij,jk->ik", skew_omega, skew_omega) / 12.0, + ( + np.eye(3) + - 0.5 * skew_omega + + ( + ( + 1.0 + - theta_safe + * np.cos(half_theta_safe) + / (2.0 * np.sin(half_theta_safe)) + ) + / theta_squared_safe + )[None, None] + * np.einsum("ij,jk->ik", skew_omega, skew_omega) + ), + ) + + return V_inv.T diff --git a/mink/lie/utils.py b/mink/lie/utils.py index 4b12159..5faadf1 100644 --- a/mink/lie/utils.py +++ b/mink/lie/utils.py @@ -1,5 +1,5 @@ -import numpy as np import mujoco +import numpy as np def get_epsilon(dtype: np.dtype) -> float: diff --git a/mink/limits/__init__.py b/mink/limits/__init__.py index 2713942..aaf9c26 100644 --- a/mink/limits/__init__.py +++ b/mink/limits/__init__.py @@ -1,9 +1,9 @@ """Kinematic limits.""" -from mink.limits.limit import Limit, Constraint +from mink.limits.collision_avoidance_limit import CollisionAvoidanceLimit from mink.limits.configuration_limit import ConfigurationLimit +from mink.limits.limit import Constraint, Limit from mink.limits.velocity_limit import VelocityLimit -from mink.limits.collision_avoidance_limit import CollisionAvoidanceLimit __all__ = ( "ConfigurationLimit", diff --git a/mink/limits/collision_avoidance_limit.py b/mink/limits/collision_avoidance_limit.py index 5dc8a2b..dabb105 100644 --- a/mink/limits/collision_avoidance_limit.py +++ b/mink/limits/collision_avoidance_limit.py @@ -1,10 +1,11 @@ from dataclasses import dataclass from typing import Sequence -import numpy as np + import mujoco +import numpy as np -from mink.limits import Limit, Constraint from mink.configuration import Configuration +from mink.limits import Constraint, Limit CollisionPairs = Sequence[tuple[Sequence[str], Sequence[str]]] diff --git a/mink/limits/configuration_limit.py b/mink/limits/configuration_limit.py index 5d28b6e..45cbbe5 100644 --- a/mink/limits/configuration_limit.py +++ b/mink/limits/configuration_limit.py @@ -1,8 +1,8 @@ -import numpy as np import mujoco +import numpy as np -from mink.limits import Limit, Constraint from mink.configuration import Configuration +from mink.limits import Constraint, Limit _SUPPORTED_JOINT_TYPES = {mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE} diff --git a/mink/limits/limit.py b/mink/limits/limit.py index 8039c54..ddc6e03 100644 --- a/mink/limits/limit.py +++ b/mink/limits/limit.py @@ -1,8 +1,8 @@ -from typing import NamedTuple - import abc +from typing import NamedTuple import numpy as np + from mink.configuration import Configuration diff --git a/mink/limits/velocity_limit.py b/mink/limits/velocity_limit.py index 9418cef..98933e5 100644 --- a/mink/limits/velocity_limit.py +++ b/mink/limits/velocity_limit.py @@ -1,7 +1,8 @@ -import numpy as np import mujoco -from mink.limits import Limit, Constraint +import numpy as np + from mink.configuration import Configuration +from mink.limits import Constraint, Limit class VelocityLimit(Limit): diff --git a/mink/solve_ik.py b/mink/solve_ik.py index 3709efb..55fc67f 100644 --- a/mink/solve_ik.py +++ b/mink/solve_ik.py @@ -1,13 +1,13 @@ """Build and solve the inverse kinematics problem.""" from typing import Sequence + import numpy as np +import qpsolvers from mink.configuration import Configuration -from mink.tasks import Task, Objective -from mink.limits import Limit, Constraint - -import qpsolvers +from mink.limits import Constraint, Limit +from mink.tasks import Objective, Task def _compute_qp_objective( diff --git a/mink/tasks/__init__.py b/mink/tasks/__init__.py index 90e4085..7a102e4 100644 --- a/mink/tasks/__init__.py +++ b/mink/tasks/__init__.py @@ -1,10 +1,9 @@ """Kinematic tasks.""" -from mink.tasks.task import Task, Objective +from mink.tasks.com_task import ComTask from mink.tasks.frame_task import FrameTask from mink.tasks.posture_task import PostureTask -from mink.tasks.com_task import ComTask - +from mink.tasks.task import Objective, Task __all__ = ( "ComTask", diff --git a/mink/tasks/com_task.py b/mink/tasks/com_task.py index c8b7471..232b021 100644 --- a/mink/tasks/com_task.py +++ b/mink/tasks/com_task.py @@ -1,8 +1,8 @@ -import numpy as np import mujoco +import numpy as np -from mink.tasks import Task from mink.configuration import Configuration +from mink.tasks import Task class ComTask(Task): diff --git a/mink/tasks/frame_task.py b/mink/tasks/frame_task.py index 56c8e56..8c42bae 100644 --- a/mink/tasks/frame_task.py +++ b/mink/tasks/frame_task.py @@ -1,9 +1,10 @@ import mujoco import numpy as np +import pinocchio as pin +from mink.configuration import Configuration from mink.lie import SE3, SO3 from mink.tasks import Task -from mink.configuration import Configuration class FrameTask(Task): @@ -50,50 +51,25 @@ def set_target_from_configuration(self, configuration: Configuration) -> None: def compute_error(self, configuration: Configuration) -> np.ndarray: if self.world_T_target is None: raise ValueError("Target transform in world frame not set.") - - ## OLD WAY OF DOING IT - # world_T_frame = configuration.get_transform_frame_to_world( - # frame_name=self.frame_name, - # frame_type=self.frame_type, - # ) - # frame_T_target = world_T_frame.inverse() @ self.world_T_target - # return frame_T_target.log() - - world_T_frame = configuration.get_transform_frame_to_world( + T_0b = configuration.get_transform_frame_to_world( frame_name=self.frame_name, frame_type=self.frame_type, ) - error = np.zeros(6) - error[:3] = world_T_frame.translation() - self.world_T_target.translation() - site_quat = np.zeros(4) - mujoco.mju_mat2Quat(site_quat, world_T_frame.rotation().as_matrix().ravel()) - mujoco.mju_subQuat(error[3:], self.world_T_target.rotation().wxyz, site_quat) - - return error + T_0t = self.world_T_target + T_bt = T_0b.inverse() @ T_0t + return T_bt.log() def compute_jacobian(self, configuration: Configuration) -> np.ndarray: if self.world_T_target is None: raise ValueError("Target transform in world frame not set.") - - jac = configuration.get_frame_jacobian( + T_0b = configuration.get_transform_frame_to_world( frame_name=self.frame_name, frame_type=self.frame_type, ) - world_T_frame = configuration.get_transform_frame_to_world( + T_0t = self.world_T_target + jac = configuration.get_frame_jacobian( frame_name=self.frame_name, frame_type=self.frame_type, ) - - ## OLD WAY OF DOING IT - # return -world_T_frame.inverse().adjoint() @ jac - - effector_quat = np.empty(4) - mujoco.mju_mat2Quat(effector_quat, world_T_frame.rotation().as_matrix().ravel()) - target_quat = self.world_T_target.rotation().wxyz - Deffector = np.empty((3, 3)) - mujoco.mjd_subQuat(target_quat, effector_quat, None, Deffector) - target_mat = self.world_T_target.rotation().as_matrix() - mat = Deffector.T @ target_mat.T - jac[3:] = mat @ jac[3:] - - return jac + T_tb = T_0t.inverse() @ T_0b + return -pin.Jlog6(pin.SE3(T_tb.as_matrix())) @ jac diff --git a/mink/tasks/posture_task.py b/mink/tasks/posture_task.py index 47e0d82..f652ce3 100644 --- a/mink/tasks/posture_task.py +++ b/mink/tasks/posture_task.py @@ -1,7 +1,8 @@ +import mujoco import numpy as np -from mink.tasks import Task + from mink.configuration import Configuration -import mujoco +from mink.tasks import Task class PostureTask(Task): diff --git a/pyproject.toml b/pyproject.toml index 2f27207..1bc3467 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,3 +45,29 @@ dev = [ "black", "mink[test]" ] + +[tool.ruff] +lint.select = [ + "E", # pycodestyle errors. + "F", # Pyflakes rules. + "PLC", # Pylint convention warnings. + "PLE", # Pylint errors. + "PLR", # Pylint refactor recommendations. + "PLW", # Pylint warnings. + "I" # Import sorting. +] +lint.ignore = [ + "E741", # Ambiguous variable name. (l, O, or I) + "E501", # Line too long. + "PLR2004", # Magic value used in comparison. + "PLR0915", # Too many statements. + "PLR0913", # Too many arguments. + "PLC0414", # Import alias does not rename variable. (this is used for exporting names) + "PLC1901", # Use falsey strings. + "PLR5501", # Use `elif` instead of `else if`. + "PLR0911", # Too many return statements. + "PLR0912", # Too many branches. + "PLW0603", # Global statement updates are discouraged. + "PLW2901" # For loop variable overwritten. +] +extend-exclude = ["**/_argparse.py"] diff --git a/tests/compare.py b/tests/compare.py new file mode 100644 index 0000000..bf50ad9 --- /dev/null +++ b/tests/compare.py @@ -0,0 +1,96 @@ +import mujoco +import numpy as np +import pinocchio as pin +from robot_descriptions.loaders.mujoco import ( + load_robot_description as mj_load_description, +) +from robot_descriptions.loaders.pinocchio import ( + load_robot_description as pin_load_description, +) + +""" +WORLD: refers to the origin of the world as if it was attached to the desired frame. +LOCAL_WORLD_ALIGNED: refers to the frame centered at the desired frame but with axes aligned with the WORLD axes. +""" + + +def skew(x): + wx, wy, wz = x + return np.array( + [ + [0.0, -wz, wy], + [wz, 0.0, -wx], + [-wy, wx, 0.0], + ] + ) + + +def adjoint(R, t): + assert R.shape == (3, 3) + assert t.shape == (3,) + A = np.block( + [ + [R, np.zeros((3, 3))], + [skew(t) @ R, R], + ] + ) + assert A.shape == (6, 6) + return A + + +def inverse(R, t): + assert R.shape == (3, 3) + assert t.shape == (3,) + R_inv = R.T + t_inv = -R_inv @ t + return R_inv, t_inv + + +np.set_printoptions(precision=2, threshold=1e-5, suppress=True) + +# PINNOCHIO +print("Loading pinnochio model...") +robot = pin_load_description("panda_description") +q = 0.1 * np.pi * (2.0 * np.random.random((9,)) - 1.0) +model_pin = robot.model +data_pin = robot.data +pin.computeJointJacobians(model_pin, data_pin, q) +pin.updateFramePlacements(model_pin, data_pin) +frame_id = model_pin.getFrameId("panda_link7") + +# MUJOCO +print("Loading mujoco model...") +model_mj = mj_load_description("panda_mj_description") +data_mj = mujoco.MjData(model_mj) +data_mj.qpos = q +mujoco.mj_kinematics(model_mj, data_mj) +mujoco.mj_comPos(model_mj, data_mj) + + +# Check frame transforms are identical. +T_sb_pin = data_pin.oMf[frame_id].copy() +t_sb_pin = T_sb_pin.np[:3, 3] +R_sb_pin = T_sb_pin.np[:3, :3] +t_sb = data_mj.body("link7").xpos.copy() +R_sb = data_mj.body("link7").xmat.copy().reshape(3, 3) +np.testing.assert_allclose(t_sb, t_sb_pin, atol=1e-7) +np.testing.assert_allclose(R_sb, R_sb_pin, atol=1e-7) + +# Check that mujoco's jac matches pinnochio's jac in frame: LOCAL_WORLD_ALIGNED. +jac_pin_local_world_aligned: np.ndarray = pin.getFrameJacobian( + model_pin, data_pin, frame_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED +) +jac_muj = np.zeros((6, model_mj.nv)) +mujoco.mj_jacBody( + model_mj, data_mj, jac_muj[:3], jac_muj[3:], model_mj.body("link7").id +) +np.testing.assert_allclose(jac_pin_local_world_aligned, jac_muj, atol=1e-7) + +# Compute Jacobian in the local frame. +jac_pin_local: np.ndarray = pin.getFrameJacobian( + model_pin, data_pin, frame_id, pin.ReferenceFrame.LOCAL +) +R_bs, t_bs = inverse(R_sb, t_sb) +A_bs = adjoint(R_bs, np.zeros(3)) +jac_muj_local = A_bs @ jac_muj +np.testing.assert_allclose(jac_muj_local, jac_pin_local, atol=1e-7) diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 0397ed2..793e91e 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -1,9 +1,8 @@ +import numpy as np from absl.testing import absltest +from robot_descriptions.loaders.mujoco import load_robot_description import mink -import numpy as np - -from robot_descriptions.loaders.mujoco import load_robot_description class TestConfiguration(absltest.TestCase): @@ -23,19 +22,22 @@ def test_site_transform_world_frame(self): site_name = "attachment_site" configuration = mink.Configuration(self.model) - # From a keyframe. + # Randomly sample a joint configuration. np.random.seed(12345) configuration.data.qpos = np.random.uniform(*configuration.model.jnt_range.T) configuration.update() + world_T_site = configuration.get_transform_frame_to_world(site_name, "site") + expected_translation = configuration.data.site(site_name).xpos np.testing.assert_array_equal(world_T_site.translation(), expected_translation) + expected_xmat = configuration.data.site(site_name).xmat.reshape(3, 3) np.testing.assert_almost_equal( world_T_site.rotation().as_matrix(), expected_xmat ) - def test_site_transform_raises_error_if_site_not_exists(self): + def test_site_transform_raises_error_if_site_name_is_invalid(self): configuration = mink.Configuration(self.model) with self.assertRaises(ValueError): configuration.get_transform_frame_to_world("invalid_name", "site") diff --git a/tests/test_jacobians.py b/tests/test_jacobians.py index 499eff5..065a2cf 100644 --- a/tests/test_jacobians.py +++ b/tests/test_jacobians.py @@ -1,10 +1,9 @@ +import numpy as np from absl.testing import absltest +from robot_descriptions.loaders.mujoco import load_robot_description import mink from mink import lie -import numpy as np - -from robot_descriptions.loaders.mujoco import load_robot_description class TestJacobians(absltest.TestCase):