diff --git a/README.md b/README.md index 0664dbac..60a2cdcf 100644 --- a/README.md +++ b/README.md @@ -31,9 +31,9 @@ - [CasADi interface](#casadi-interface) - [PyTorch interface](#pytorch-interface) - [PyTorch Batched interface](#pytorch-batched-interface) + - [MuJoCo interface](#mujoco-interface) - [Inverse Kinematics](#inverse-kinematics) - [πŸ¦Έβ€β™‚οΈ Contributing](#️-contributing) - - [Todo](#todo) ## 🐍 Dependencies @@ -352,6 +352,59 @@ M = kinDyn.mass_matrix(w_H_b_batch, joints_batch) w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b_batch, joints_batch) ``` +### MuJoCo interface + +adam supports loading models directly from [MuJoCo](https://mujoco.org/) `MjModel` objects. This is useful when working with MuJoCo simulations or models from [robot_descriptions](https://github.com/robot-descriptions/robot_descriptions.py). + +```python +import mujoco +import numpy as np +from adam import Representations +from adam.numpy import KinDynComputations + +# Load a MuJoCo model (e.g., from robot_descriptions) +from robot_descriptions.loaders.mujoco import load_robot_description +mj_model = load_robot_description("g1_mj_description") + +# Create KinDynComputations directly from MuJoCo model +kinDyn = KinDynComputations.from_mujoco_model(mj_model) + +# Set velocity representation (default is mixed) +kinDyn.set_frame_velocity_representation(Representations.MIXED_REPRESENTATION) + +# Set gravity to match MuJoCo settings +kinDyn.g = np.concatenate([mj_model.opt.gravity, np.zeros(3)]) + +# Create MuJoCo data and set state +mj_data = mujoco.MjData(mj_model) +mj_data.qpos[:] = your_qpos # Your configuration +mj_data.qvel[:] = your_qvel # Your velocities +mujoco.mj_forward(mj_model, mj_data) + +# Extract base transform from MuJoCo state (for floating-base robots) +from scipy.spatial.transform import Rotation as R +base_rot = R.from_quat(mj_data.qpos[3:7], scalar_first=True).as_matrix() +base_pos = mj_data.qpos[0:3] +w_H_b = np.eye(4) +w_H_b[:3, :3] = base_rot +w_H_b[:3, 3] = base_pos + +# Joint positions (excluding free joint). +# Be sure the serialization between mujoco and adam is the same +joints = mj_data.qpos[7:] + +# Compute dynamics quantities +M = kinDyn.mass_matrix(w_H_b, joints) +com_pos = kinDyn.CoM_position(w_H_b, joints) +J = kinDyn.jacobian('frame_name', w_H_b, joints) +``` + +> [!NOTE] +> When using `from_mujoco_model`, adam automatically extracts the joint names from the MuJoCo model. You can also specify `use_mujoco_actuators=True` to use actuator names instead of joint names. + +> [!WARNING] +> MuJoCo uses a different velocity representation for the floating base. The free joint velocity in MuJoCo is `[I \dot{p}_B, B \omega_B]`, while mixed representation uses `[I \dot{p}_B, I \omega_B]`. Make sure to handle this transformation when comparing with MuJoCo computations. + ### Inverse Kinematics adam provides an interface for solving inverse kinematics problems using CasADi. The solver supports @@ -400,13 +453,3 @@ Open an issue with your feature request or if you spot a bug. Then, you can also > [!WARNING] > REPOSITORY UNDER DEVELOPMENT! We cannot guarantee stable API - -## Todo - -- [x] Center of Mass position -- [x] Jacobians -- [x] Forward kinematics -- [x] Mass Matrix via CRBA -- [x] Centroidal Momentum Matrix via CRBA -- [x] Recursive Newton-Euler algorithm (still no acceleration in the algorithm, since it is used only for the computation of the bias force) -- [ ] Articulated Body algorithm diff --git a/ci_env.yml b/ci_env.yml index 5955a9e1..ee02a0e0 100644 --- a/ci_env.yml +++ b/ci_env.yml @@ -21,3 +21,5 @@ dependencies: - requests - array-api-compat - liecasadi + - mujoco + - robot_descriptions diff --git a/ci_env_win.yml b/ci_env_win.yml index 5c06164a..b2a6cf13 100644 --- a/ci_env_win.yml +++ b/ci_env_win.yml @@ -19,3 +19,5 @@ dependencies: - requests - liecasadi - array-api-compat + - mujoco + - robot_descriptions diff --git a/pyproject.toml b/pyproject.toml index 5387def0..c4d1a1c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,8 @@ test = [ "jaxlib", "casadi", "torch", + "mujoco", + "robot_descriptions", "pytest", "idyntree", "icub-models", diff --git a/src/adam/casadi/computations.py b/src/adam/casadi/computations.py index 3fba0daf..3191e96a 100644 --- a/src/adam/casadi/computations.py +++ b/src/adam/casadi/computations.py @@ -8,10 +8,11 @@ from adam.casadi.casadi_like import SpatialMath from adam.core import RBDAlgorithms from adam.core.constants import Representations -from adam.model import Model, URDFModelFactory +from adam.model import Model, build_model_factory +from adam.model.kindyn_mixin import KinDynFactoryMixin -class KinDynComputations: +class KinDynComputations(KinDynFactoryMixin): """Class that retrieves robot quantities using CasADi for Floating Base systems.""" def __init__( @@ -24,12 +25,13 @@ def __init__( ) -> None: """ Args: - urdfstring (str): either path or string of the urdf + urdfstring (str): path/string of a URDF or a MuJoCo MjModel. + NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release. joints_name_list (list): list of the actuated joints root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None. """ math = SpatialMath() - factory = URDFModelFactory(path=urdfstring, math=math) + factory = build_model_factory(description=urdfstring, math=math) model = Model.build(factory=factory, joints_name_list=joints_name_list) self.rbdalgos = RBDAlgorithms(model=model, math=math) self.NDoF = self.rbdalgos.NDoF diff --git a/src/adam/jax/computations.py b/src/adam/jax/computations.py index fa12a8b8..341d5f6f 100644 --- a/src/adam/jax/computations.py +++ b/src/adam/jax/computations.py @@ -8,11 +8,12 @@ from adam.core.constants import Representations from adam.core.rbd_algorithms import RBDAlgorithms from adam.jax.jax_like import SpatialMath -from adam.model import Model, URDFModelFactory +from adam.model import Model, build_model_factory +from adam.model.kindyn_mixin import KinDynFactoryMixin from adam.core.array_api_math import spec_from_reference -class KinDynComputations: +class KinDynComputations(KinDynFactoryMixin): """This is a small class that retrieves robot quantities using Jax for Floating Base systems.""" def __init__( @@ -25,13 +26,14 @@ def __init__( ) -> None: """ Args: - urdfstring (str): either path or string of the urdf + urdfstring (str): path/string of a URDF or a MuJoCo MjModel. + NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release. joints_name_list (list): list of the actuated joints root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None. """ ref = jnp.array(0.0, dtype=dtype) math = SpatialMath(spec_from_reference(ref)) - factory = URDFModelFactory(path=urdfstring, math=math) + factory = build_model_factory(description=urdfstring, math=math) model = Model.build(factory=factory, joints_name_list=joints_name_list) self.rbdalgos = RBDAlgorithms(model=model, math=math) self.NDoF = self.rbdalgos.NDoF diff --git a/src/adam/model/__init__.py b/src/adam/model/__init__.py index a619183f..30fcc18f 100644 --- a/src/adam/model/__init__.py +++ b/src/adam/model/__init__.py @@ -1,6 +1,8 @@ from .abc_factories import Inertia, Inertial, Joint, Limits, Link, ModelFactory, Pose +from .factory import build_model_factory from .model import Model from .std_factories.std_joint import StdJoint from .std_factories.std_link import StdLink +from .mj_factory.mujoco_model import MujocoModelFactory from .std_factories.std_model import URDFModelFactory from .tree import Node, Tree diff --git a/src/adam/model/factory.py b/src/adam/model/factory.py new file mode 100644 index 00000000..ea25176b --- /dev/null +++ b/src/adam/model/factory.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import pathlib +from typing import Any + +from adam.model.abc_factories import ModelFactory +from adam.model.std_factories.std_model import URDFModelFactory + + +def _is_mujoco_model(obj: Any) -> bool: + """Check if obj is a MuJoCo MjModel without importing mujoco.""" + cls = obj.__class__ + cls_name = getattr(cls, "__name__", "") + cls_module = getattr(cls, "__module__", "") + if cls_name != "MjModel" or "mujoco" not in cls_module: + return False + return all(hasattr(obj, attr) for attr in ("nq", "nv", "nu", "nbody")) + + +def _is_urdf(obj: Any) -> bool: + """Check if obj is a URDF.""" + if isinstance(obj, pathlib.Path): + return obj.suffix.lower() == ".urdf" + + if isinstance(obj, str): + s = obj.lstrip() + if s.startswith("<") and " ModelFactory: + """Return a ModelFactory from a URDF string/path or a MuJoCo model.""" + + if _is_mujoco_model(description): + + from adam.model.mj_factory.mujoco_model import MujocoModelFactory + + return MujocoModelFactory(mj_model=description, math=math) + + if _is_urdf(description): + return URDFModelFactory(path=description, math=math) + + raise ValueError( + f"Unsupported model description. Expected a URDF path/string or a mujoco.MjModel. " + f"Got: {type(description)!r}" + ) diff --git a/src/adam/model/kindyn_mixin.py b/src/adam/model/kindyn_mixin.py new file mode 100644 index 00000000..98e0bba4 --- /dev/null +++ b/src/adam/model/kindyn_mixin.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from typing import Any + + +class KinDynFactoryMixin: + """Shared helpers to instantiate KinDyn* classes from different model sources.""" + + @classmethod + def from_urdf(cls: type[KinDynFactoryMixin], urdfstring: Any, *args, **kwargs): + """Instantiate using a URDF path/string. + + Args: + urdfstring (str): path/string of a URDF + + Returns: + KinDynFactoryMixin: An instance of the class initialized with the provided URDF and arguments. + """ + return cls(urdfstring, *args, **kwargs) + + @classmethod + def from_mujoco_model( + cls: type[KinDynFactoryMixin], + mujoco_model: Any, + use_mujoco_actuators: bool = False, + *args, + **kwargs, + ): + """Instantiate using a MuJoCo model. + + Args: + mujoco_model (MjModel): MuJoCo model instance + use_mujoco_actuators (bool): If True, use the names of joints under the tags in the MuJoCo XML as joint names + (i.e., set 'joints_name_list' to actuator joint names). This is useful when you want the + KinDyn* instance to use actuator joint names instead of the default joint names in the xml. + Default is False. + + Returns: + KinDynFactoryMixin: An instance of the class initialized with the provided MuJoCo model + """ + if use_mujoco_actuators: + # use as joint names the names of joints under the tags in the mujoco xml + actuator_names = [ + mujoco_model.actuator(a).name for a in range(mujoco_model.nu) + ] + kwargs.setdefault("joints_name_list", actuator_names) + return cls(mujoco_model, *args, **kwargs) diff --git a/src/adam/model/mj_factory/__init__.py b/src/adam/model/mj_factory/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/adam/model/mj_factory/mujoco_model.py b/src/adam/model/mj_factory/mujoco_model.py new file mode 100644 index 00000000..11e077b3 --- /dev/null +++ b/src/adam/model/mj_factory/mujoco_model.py @@ -0,0 +1,287 @@ +from dataclasses import dataclass +from typing import Optional, TYPE_CHECKING + +import numpy as np +from scipy.spatial.transform import Rotation as R + +from adam.core.spatial_math import SpatialMath +from adam.model.abc_factories import Limits, ModelFactory +from adam.model.std_factories.std_joint import StdJoint +from adam.model.std_factories.std_link import StdLink + +# Type checking only - doesn't execute at runtime +if TYPE_CHECKING: + import mujoco + + +@dataclass +class MujocoOrigin: + xyz: np.ndarray + rpy: np.ndarray + + +@dataclass +class MujocoInertia: + ixx: float + ixy: float + ixz: float + iyy: float + iyz: float + izz: float + + +@dataclass +class MujocoInertial: + mass: float + inertia: MujocoInertia + origin: MujocoOrigin + + +@dataclass +class MujocoLink: + name: str + inertial: MujocoInertial + visuals: list + collisions: list + + +@dataclass +class MujocoJoint: + name: str + parent: str + child: str + joint_type: str + axis: Optional[np.ndarray] + origin: MujocoOrigin + limit: Optional[Limits] + + +def _normalize_quaternion(quat: np.ndarray) -> np.ndarray: + norm = np.linalg.norm(quat) + if norm == 0: + return np.array([1.0, 0.0, 0.0, 0.0], dtype=float) + return quat / norm + + +def _rotate_vector(quat: np.ndarray, vec: np.ndarray) -> np.ndarray: + """Rotate a vector using quaternion [w, x, y, z].""" + rot = R.from_quat(quat, scalar_first=True).as_matrix() + return rot @ vec + + +class MujocoModelFactory(ModelFactory): + """Factory that builds a model starting from a mujoco.MjModel.""" + + def __init__(self, mj_model: "mujoco.MjModel", math: SpatialMath): + self.math = math + self.mujoco = self._import_mujoco() + self.mj_model = self._model_exists(mj_model) + fallback_name = "mujoco_model" + self.name = getattr(self.mj_model, "name", None) or fallback_name + + self._links = self._build_links() + self._child_map = self._build_child_map() + self._joints = self._build_joints() + + def _import_mujoco(self): + try: + import mujoco + except ImportError as exc: # pragma: no cover - dependency optional + raise ImportError( + "The 'MuJoCo' package is required to load MuJoCo models." + ) from exc + return mujoco + + def _model_exists(self, mj_model): + if isinstance(mj_model, self.mujoco.MjModel): + return mj_model + + raise ValueError( + f"Expected a MuJoCo MjModel object, but got {type(mj_model).__name__}." + ) + + def _body_name(self, body_id: int) -> str: + name = self.mujoco.mj_id2name( + self.mj_model, self.mujoco.mjtObj.mjOBJ_BODY, body_id + ) + return name if name is not None else f"body_{body_id}" + + def _joint_name(self, joint_id: int) -> str: + name = self.mujoco.mj_id2name( + self.mj_model, self.mujoco.mjtObj.mjOBJ_JOINT, joint_id + ) + return name if name is not None else f"joint_{joint_id}" + + def _link_inertial(self, body_id: int) -> MujocoInertial: + mass = float(self.mj_model.body_mass[body_id]) + inertia_diagonal = self.mj_model.body_inertia[body_id] + inertia = MujocoInertia( + ixx=float(inertia_diagonal[0]), + ixy=0.0, + ixz=0.0, + iyy=float(inertia_diagonal[1]), + iyz=0.0, + izz=float(inertia_diagonal[2]), + ) + + ipos = np.array(self.mj_model.body_ipos[body_id], dtype=float) + iquat = _normalize_quaternion( + np.array(self.mj_model.body_iquat[body_id], dtype=float) + ) + origin = MujocoOrigin( + xyz=ipos, + rpy=R.from_quat(iquat, scalar_first=True).as_euler("xyz"), + ) + return MujocoInertial(mass=mass, inertia=inertia, origin=origin) + + def _build_links(self) -> list[StdLink]: + links: list[StdLink] = [] + for body_id in range(1, self.mj_model.nbody): + link = MujocoLink( + name=self._body_name(body_id), + inertial=self._link_inertial(body_id), + visuals=[], + collisions=[], + ) + links.append(StdLink(link, self.math)) + return links + + def _build_child_map(self) -> dict[str, list[str]]: + child_map: dict[str, list[str]] = {} + for body_id in range(1, self.mj_model.nbody): + parent_id = int(self.mj_model.body_parentid[body_id]) + parent_name = self._body_name(parent_id) if parent_id > 0 else None + if parent_name is None: + continue + child_map.setdefault(parent_name, []).append(self._body_name(body_id)) + return child_map + + def _joint_origin(self, body_id: int, joint_id: Optional[int]) -> MujocoOrigin: + body_pos = np.array(self.mj_model.body_pos[body_id], dtype=float) + body_quat = _normalize_quaternion( + np.array(self.mj_model.body_quat[body_id], dtype=float) + ) + xyz = body_pos + if joint_id is not None: + j_pos = np.array(self.mj_model.jnt_pos[joint_id], dtype=float) + if np.linalg.norm(j_pos) > 0.0: + xyz = xyz + _rotate_vector(body_quat, j_pos) + rpy = R.from_quat(body_quat, scalar_first=True).as_euler("xyz") + return MujocoOrigin(xyz=xyz, rpy=rpy) + + def _build_limits(self, joint_id: int, joint_type: str) -> Optional[Limits]: + if joint_type == "fixed": + return None + limited = bool(self.mj_model.jnt_limited[joint_id]) + if not limited: + return None + lower, upper = self.mj_model.jnt_range[joint_id] + return Limits(lower=lower, upper=upper, effort=None, velocity=None) + + def _joint_type(self, mj_type: int) -> str: + if mj_type == self.mujoco.mjtJoint.mjJNT_HINGE: + return "revolute" + if mj_type == self.mujoco.mjtJoint.mjJNT_SLIDE: + return "prismatic" + return "unsupported" + + def _build_joint( + self, + body_id: int, + joint_id: Optional[int], + parent_name: str, + joint_type: str, + ) -> StdJoint: + child_name = self._body_name(body_id) + name = ( + self._joint_name(joint_id) + if joint_id is not None + else f"{parent_name}_to_{child_name}_fixed" + ) + axis = ( + np.array(self.mj_model.jnt_axis[joint_id], dtype=float) + if joint_type != "fixed" and joint_id is not None + else None + ) + origin = self._joint_origin(body_id, joint_id) + limit = ( + self._build_limits(joint_id, joint_type) if joint_id is not None else None + ) + joint = MujocoJoint( + name=name, + parent=parent_name, + child=child_name, + joint_type=joint_type, + axis=axis, + origin=origin, + limit=limit, + ) + return StdJoint(joint, self.math) + + def _build_joints(self) -> list[StdJoint]: + joints: list[StdJoint] = [] + for body_id in range(1, self.mj_model.nbody): + parent_id = int(self.mj_model.body_parentid[body_id]) + if parent_id < 1: + continue + parent_name = self._body_name(parent_id) + joint_start = int(self.mj_model.body_jntadr[body_id]) + joint_num = int(self.mj_model.body_jntnum[body_id]) + + if joint_num == 0: + joints.append( + self._build_joint( + body_id=body_id, + joint_id=None, + parent_name=parent_name, + joint_type="fixed", + ) + ) + continue + + for joint_id in range(joint_start, joint_start + joint_num): + joint_type = self._joint_type(int(self.mj_model.jnt_type[joint_id])) + if joint_type == "unsupported": + # Skip free/ball joints; base pose is provided externally. + continue + joints.append( + self._build_joint( + body_id=body_id, + joint_id=joint_id, + parent_name=parent_name, + joint_type=joint_type, + ) + ) + return joints + + def build_joint(self, joint) -> StdJoint: # pragma: no cover - required by ABC + raise NotImplementedError("MujocoModelFactory does not build joints externally") + + def build_link(self, link) -> StdLink: # pragma: no cover - required by ABC + raise NotImplementedError("MujocoModelFactory does not build links externally") + + def get_joints(self) -> list[StdJoint]: + return self._joints + + def _has_non_fixed_joint(self, link_name: str) -> bool: + return any(j.child == link_name and j.type != "fixed" for j in self._joints) + + def get_links(self) -> list[StdLink]: + return [ + link + for link in self._links + if ( + float(link.inertial.mass.array) != 0.0 + or link.name in self._child_map.keys() + or self._has_non_fixed_joint(link.name) + ) + ] + + def get_frames(self) -> list[StdLink]: + return [ + link + for link in self._links + if float(link.inertial.mass.array) == 0.0 + and link.name not in self._child_map.keys() + and not self._has_non_fixed_joint(link.name) + ] diff --git a/src/adam/model/std_factories/std_model.py b/src/adam/model/std_factories/std_model.py index 7d70e9bc..a130224e 100644 --- a/src/adam/model/std_factories/std_model.py +++ b/src/adam/model/std_factories/std_model.py @@ -5,7 +5,9 @@ import urdf_parser_py.urdf from adam.core.spatial_math import SpatialMath -from adam.model import ModelFactory, StdJoint, StdLink +from adam.model.abc_factories import ModelFactory +from adam.model.std_factories.std_joint import StdJoint +from adam.model.std_factories.std_link import StdLink def urdf_remove_sensors_tags(xml_string): diff --git a/src/adam/numpy/computations.py b/src/adam/numpy/computations.py index 74c6942e..97c4a091 100644 --- a/src/adam/numpy/computations.py +++ b/src/adam/numpy/computations.py @@ -6,12 +6,13 @@ from adam.core.constants import Representations from adam.core.rbd_algorithms import RBDAlgorithms -from adam.model import Model, URDFModelFactory +from adam.model import Model, build_model_factory +from adam.model.kindyn_mixin import KinDynFactoryMixin from adam.numpy.numpy_like import SpatialMath from adam.core.array_api_math import spec_from_reference -class KinDynComputations: +class KinDynComputations(KinDynFactoryMixin): """This is a small class that retrieves robot quantities using NumPy for Floating Base systems.""" def __init__( @@ -24,13 +25,13 @@ def __init__( ) -> None: """ Args: - urdfstring (str): either path or string of the urdf + urdfstring (str): path/string of a URDF or a MuJoCo MjModel joints_name_list (list): list of the actuated joints root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None. """ ref = np.array(0.0, dtype=dtype) math = SpatialMath(spec_from_reference(ref)) - factory = URDFModelFactory(path=urdfstring, math=math) + factory = build_model_factory(description=urdfstring, math=math) model = Model.build(factory=factory, joints_name_list=joints_name_list) self.rbdalgos = RBDAlgorithms(model=model, math=math) self.NDoF = model.NDoF diff --git a/src/adam/pytorch/computation_batch.py b/src/adam/pytorch/computation_batch.py index 7b8c2911..768a0412 100644 --- a/src/adam/pytorch/computation_batch.py +++ b/src/adam/pytorch/computation_batch.py @@ -8,11 +8,12 @@ from adam.core.constants import Representations from adam.core.rbd_algorithms import RBDAlgorithms from adam.pytorch.torch_like import SpatialMath -from adam.model import Model, URDFModelFactory +from adam.model import Model, build_model_factory +from adam.model.kindyn_mixin import KinDynFactoryMixin from adam.core.array_api_math import spec_from_reference -class KinDynComputationsBatch: +class KinDynComputationsBatch(KinDynFactoryMixin): """ A PyTorch-based class for batch kinematic and dynamic computations on robotic systems. @@ -33,13 +34,14 @@ def __init__( ) -> None: """ Args: - urdfstring (str): path of the urdf + urdfstring (str): path/string of a URDF or a MuJoCo MjModel. + NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release. joints_name_list (list): list of the actuated joints root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None. """ ref = torch.tensor(0.0, dtype=dtype, device=device) math = SpatialMath(spec_from_reference(ref)) - factory = URDFModelFactory(path=urdfstring, math=math) + factory = build_model_factory(description=urdfstring, math=math) model = Model.build(factory=factory, joints_name_list=joints_name_list) self.rbdalgos = RBDAlgorithms(model=model, math=math) self.NDoF = self.rbdalgos.NDoF diff --git a/src/adam/pytorch/computations.py b/src/adam/pytorch/computations.py index ddc4b9f8..c476d1c0 100644 --- a/src/adam/pytorch/computations.py +++ b/src/adam/pytorch/computations.py @@ -8,12 +8,13 @@ from adam.core.constants import Representations from adam.core.rbd_algorithms import RBDAlgorithms -from adam.model import Model, URDFModelFactory +from adam.model import Model, build_model_factory +from adam.model.kindyn_mixin import KinDynFactoryMixin from adam.pytorch.torch_like import SpatialMath from adam.core.array_api_math import spec_from_reference -class KinDynComputations: +class KinDynComputations(KinDynFactoryMixin): """This is a small class that retrieves robot quantities using Pytorch for Floating Base systems.""" def __init__( @@ -29,14 +30,15 @@ def __init__( ) -> None: """ Args: - urdfstring (str): either path or string of the urdf + urdfstring (str): path/string of a URDF or a MuJoCo MjModel. + NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release. joints_name_list (list): list of the actuated joints root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None. """ ref = torch.tensor(0.0, dtype=dtype, device=device) spec = spec_from_reference(ref) math = SpatialMath(spec=spec) - factory = URDFModelFactory(path=urdfstring, math=math) + factory = build_model_factory(description=urdfstring, math=math) model = Model.build(factory=factory, joints_name_list=joints_name_list) self.rbdalgos = RBDAlgorithms(model=model, math=math) self.NDoF = self.rbdalgos.NDoF diff --git a/tests/test_mujoco.py b/tests/test_mujoco.py new file mode 100644 index 00000000..64afd512 --- /dev/null +++ b/tests/test_mujoco.py @@ -0,0 +1,339 @@ +import mujoco +import numpy as np +import pytest +from scipy.spatial.transform import Rotation as R +from robot_descriptions.loaders.mujoco import load_robot_description + +from adam import Representations +from adam.numpy.computations import KinDynComputations + + +DESCRIPTION_NAMES = ["g1_mj_description", "aliengo_mj_description"] + + +def _load_model(description: str) -> mujoco.MjModel: + print(f"Loading robot description '{description}'") + try: + model = load_robot_description(description) + except Exception as exc: # pragma: no cover - environment dependent + pytest.skip(f"Cannot load {description}: {exc}") + if model is None: # pragma: no cover - loader returns None on missing path + pytest.skip(f"{description} not available (loader returned None)") + return model + + +def _random_configuration( + model: mujoco.MjModel, rng: np.random.Generator +) -> np.ndarray: + q = rng.normal(scale=1.0, size=model.nq) + # Normalize the free joint quaternion (indices 3:7) + quat = q[3:7] + quat = quat / np.linalg.norm(quat) + q[3:7] = quat + return q + + +def _random_velocity(model: mujoco.MjModel, rng: np.random.Generator) -> np.ndarray: + return rng.normal(scale=1.0, size=model.nv) + + +def _set_state( + model: mujoco.MjModel, data: mujoco.MjData, qpos: np.ndarray, qvel: np.ndarray +): + data.qpos[:] = qpos + data.qvel[:] = qvel + mujoco.mj_forward(model, data) + + +def _base_velocity_transform( + base_rotation: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """ + MuJoCo encodes a free-joint qvel as [I v_B, B Ο‰_B] (I v_B = I \\dot{p}_B). + ADAM mixed vel representation uses [I v_B, I Ο‰_B]. Given R_IB = base_rotation, + I Ο‰_B = R_IB @ (B Ο‰_B). + Adj maps qvel_MJ -> nu_ADAM for the base; Adj_inv reverses it. + """ + Adj = np.block( + [ + [np.eye(3), np.zeros((3, 3))], + [np.zeros((3, 3)), base_rotation], + ] + ) + Adj_inv = np.block( + [ + [np.eye(3), np.zeros((3, 3))], + [np.zeros((3, 3)), base_rotation.T], + ] + ) + return Adj, Adj_inv + + +def _joint_permutation( + model: mujoco.MjModel, kd: KinDynComputations +) -> tuple[np.ndarray, list[str], list[str]]: + mj_joint_names = [ + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j_id) + for j_id in range(model.njnt) + if model.jnt_type[j_id] != mujoco.mjtJoint.mjJNT_FREE + ] + adam_joint_names = kd.rbdalgos.model.actuated_joints + perm = [mj_joint_names.index(name) for name in adam_joint_names] + P = np.eye(len(mj_joint_names))[:, perm] + # TODO: the P in this test is an identity matrix. The actuated joints and the MuJoCo joints are the same and in the same order. The cases in which this is not true should be handled by the user. + return P, mj_joint_names, adam_joint_names + + +def _compose_transforms( + base_rotation: np.ndarray, permutation: np.ndarray, n_act: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + # MuJoCo qvel layout (free joint first) is [I v_B, B Ο‰_B, joint_vel_XML]: + # - linear velocity of the base origin: I v_B = I \dot{p}_B + # - angular velocity of the base expressed in the body frame B + # - joint rates in the XML joint order + # ADAM mixed-velocity qdot is [I v_B, I Ο‰_B, joint_vel_ADAM]: + # - linear matches MuJoCo: I v_B = I \dot{p}_B + # - angular expressed in I (I Ο‰_B = R_IB B Ο‰_B) + # - joint rates reordered to ADAM's actuated_joints ordering + # Let R_IB = base_rotation. Then: + # I Ο‰_B = R_IB @ (B Ο‰_B) + # qdot_ADAM = [I v_B; I Ο‰_B; P^T qdot_XML] = S_inv @ qvel_MJ + # where + # C_inv = blkdiag(I3, R_IB^T), S_inv = [[C_inv, 0],[0, P]] + # For generalized forces (dual map): + # tau_MJ = S_inv.T @ tau_ADAM, i.e. base torques use R_IB and joints use P. + # Remember: the permutation matrix here is an identity. + Adj, Adj_inv = _base_velocity_transform(base_rotation) + S_inv = np.block( + [ + [Adj_inv, np.zeros((6, n_act))], + [np.zeros((n_act, 6)), permutation], + ] + ) + S_inv_force = np.block( + [ + [Adj_inv.T, np.zeros((6, n_act))], + [np.zeros((n_act, 6)), permutation.T], + ] + ) + return Adj, S_inv, S_inv_force + + +def _system_com(model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray: + masses = model.body_mass[:, None] + return np.sum(masses * data.xipos, axis=0) / np.sum(masses) + + +def _system_com_jacobian(model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray: + total_mass = np.sum(model.body_mass) + jac_total = np.zeros((3, model.nv)) + for body_id in range(1, model.nbody): + jacp = np.zeros((3, model.nv)) + jacr = np.zeros((3, model.nv)) + mujoco.mj_jacBodyCom(model, data, jacp, jacr, body_id) + jac_total += model.body_mass[body_id] * jacp + return jac_total / total_mass + + +def _bias_forces( + model: mujoco.MjModel, qpos: np.ndarray, qvel: np.ndarray, gravity: np.ndarray +) -> np.ndarray: + data = mujoco.MjData(model) + data.qpos[:] = qpos + data.qvel[:] = qvel + original_gravity = model.opt.gravity.copy() + model.opt.gravity[:] = gravity + mujoco.mj_forward(model, data) + model.opt.gravity[:] = original_gravity + return data.qfrc_bias.copy() + + +@pytest.fixture(scope="module", params=DESCRIPTION_NAMES) +def mujoco_setup(request): + model = _load_model(request.param) + + rng = np.random.default_rng(42) + qpos = _random_configuration(model, rng) + qvel = _random_velocity(model, rng) + + data = mujoco.MjData(model) + _set_state(model, data, qpos, qvel) + + kd = KinDynComputations.from_mujoco_model(model) + kd.set_frame_velocity_representation(Representations.MIXED_REPRESENTATION) + kd.g = np.concatenate([model.opt.gravity, np.zeros(3)]) + + base_rot = R.from_quat( + data.qpos[3:7], scalar_first=True + ).as_matrix() # free joint quaternion is at indices 3:7 + base_pos = data.qpos[:3] + + P, mj_joint_names, adam_joint_names = _joint_permutation(model, kd) + Adj, S_inv, S_inv_force = _compose_transforms(base_rot, P, kd.rbdalgos.NDoF) + base_transform = np.eye(4) + base_transform[:3, :3] = base_rot + base_transform[:3, 3] = base_pos + + q_joints = P.T @ qpos[7:] # reorder joints to ADAM ordering + base_vel = Adj @ qvel[:6] + qd = P.T @ qvel[6:] + return dict( + model=model, + data=data, + kd=kd, + base_transform=base_transform, + base_vel=base_vel, + q_joints=q_joints, + qd=qd, + qpos=qpos, + qvel=qvel, + joint_permutation=P, + base_transform_matrix=Adj, + velocity_transform=S_inv, + force_transform=S_inv_force, + mj_joint_names=mj_joint_names, + adam_joint_names=adam_joint_names, + ) + + +def test_forward_kinematics(mujoco_setup): + model = mujoco_setup["model"] + data = mujoco_setup["data"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + + for link_name in kd.rbdalgos.model.links: + body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, link_name) + if body_id < 0: + pytest.skip(f"Missing body {link_name} in mujoco model") + T_adam = kd.forward_kinematics(link_name, base_transform, q_joints) + R_body = data.xmat[body_id].reshape(3, 3) + p = data.xpos[body_id] + T_mj = np.eye(4) + T_mj[:3, :3] = R_body + T_mj[:3, 3] = p + np.testing.assert_allclose(T_adam, T_mj, atol=1e-4, rtol=1e-4) + + +def test_jacobian(mujoco_setup): + model = mujoco_setup["model"] + data = mujoco_setup["data"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + n_act = kd.rbdalgos.NDoF + C = mujoco_setup["base_transform_matrix"] + P = mujoco_setup["joint_permutation"] + + for link_name in kd.rbdalgos.model.links: + body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, link_name) + if body_id < 0: + pytest.skip(f"{DESCRIPTION_NAME} missing body {link_name} in mujoco model") + J_adam = kd.jacobian(link_name, base_transform, q_joints) + J_adam_j = J_adam[:, -n_act:] + jacp = np.zeros((3, model.nv)) + jacr = np.zeros((3, model.nv)) + mujoco.mj_jacBody(model, data, jacp, jacr, body_id) + J_mj = np.vstack([jacp, jacr]) + J_mj_joint = J_mj[:, 6:] + J_mj_joint_reordered = J_mj_joint @ P + J_adam_base = J_adam[:, :6] @ C + np.testing.assert_allclose(J_mj[:, :6], J_adam_base, atol=1e-4, rtol=1e-4) + np.testing.assert_allclose(J_mj_joint_reordered, J_adam_j, atol=1e-4, rtol=1e-4) + + +def test_mass_matrix(mujoco_setup): + model = mujoco_setup["model"] + data = mujoco_setup["data"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + S_inv = mujoco_setup["velocity_transform"] + M_mj = np.zeros((model.nv, model.nv)) + mujoco.mj_fullM(model, M_mj, data.qM) + M_adam = kd.mass_matrix(base_transform, q_joints) + # Remove MuJoCo joint armature (rotor inertia) from the full mass matrix. + M_mj_no_arm = M_mj.copy() + armature_joint = model.dof_armature[6:] + M_mj_no_arm[6:, 6:] -= np.diag(armature_joint) + # Change coordinates: qdot_adam = S qdot_mj, so M_adam = S^{-T} M_mj S^{-1} + M_expected = S_inv.T @ M_mj_no_arm @ S_inv + np.testing.assert_allclose(M_expected, M_adam, atol=2e-2, rtol=1e-3) + + +def test_total_mass(mujoco_setup): + model = mujoco_setup["model"] + kd = mujoco_setup["kd"] + np.testing.assert_allclose( + kd.get_total_mass(), np.sum(model.body_mass), rtol=1e-6, atol=1e-6 + ) + + +def test_com_position(mujoco_setup): + model = mujoco_setup["model"] + data = mujoco_setup["data"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + com_mj = _system_com(model, data) + com_adam = kd.CoM_position(base_transform, q_joints) + np.testing.assert_allclose(com_adam, com_mj, atol=1e-4, rtol=1e-4) + + +def test_com_jacobian(mujoco_setup): + model = mujoco_setup["model"] + data = mujoco_setup["data"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + Adj_inv = mujoco_setup["velocity_transform"] + + jac_mj = _system_com_jacobian(model, data) + jac_mj_expected = jac_mj @ Adj_inv + jac_adam = kd.CoM_jacobian(base_transform, q_joints) + np.testing.assert_allclose(jac_adam, jac_mj_expected, atol=1e-4, rtol=1e-4) + + +def test_bias_force(mujoco_setup): + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + base_vel = mujoco_setup["base_vel"] + qd = mujoco_setup["qd"] + bias_mj = mujoco_setup["data"].qfrc_bias.copy() + tau_expected = mujoco_setup["force_transform"] @ bias_mj + tau_adam = kd.bias_force(base_transform, q_joints, base_vel, qd) + np.testing.assert_allclose(tau_adam, tau_expected, atol=1e-3, rtol=1e-4) + + +def test_gravity_term(mujoco_setup): + model = mujoco_setup["model"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + qpos = mujoco_setup["qpos"] + zeros = np.zeros_like(mujoco_setup["qvel"]) + bias_mj = _bias_forces(model, qpos, zeros, model.opt.gravity.copy()) + tau_expected = mujoco_setup["force_transform"] @ bias_mj + tau_adam = kd.gravity_term(base_transform, q_joints) + np.testing.assert_allclose(tau_adam, tau_expected, atol=1e-3, rtol=1e-4) + + +def test_coriolis_term(mujoco_setup): + model = mujoco_setup["model"] + kd = mujoco_setup["kd"] + base_transform = mujoco_setup["base_transform"] + q_joints = mujoco_setup["q_joints"] + base_vel = mujoco_setup["base_vel"] + qd = mujoco_setup["qd"] + qpos = mujoco_setup["qpos"] + qvel = mujoco_setup["qvel"] + zero_gravity = np.zeros(3) + bias_mj = _bias_forces(model, qpos, qvel, zero_gravity) + tau_expected = mujoco_setup["force_transform"] @ bias_mj + original_g = kd.g + kd.g = np.concatenate([zero_gravity, np.zeros(3)]) + tau_adam = kd.coriolis_term(base_transform, q_joints, base_vel, qd) + kd.g = original_g + np.testing.assert_allclose(tau_adam, tau_expected, atol=1e-3, rtol=1e-4)