Add MujocoModelFactory

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

ci_env.yml

@@ -21,3 +21,5 @@ dependencies:
   - requests
   - array-api-compat
   - liecasadi
+  - mujoco
+  - robot_descriptions

ci_env_win.yml

@@ -19,3 +19,5 @@ dependencies:
   - requests
   - liecasadi
   - array-api-compat
+  - mujoco
+  - robot_descriptions

pyproject.toml

@@ -43,6 +43,8 @@ test = [
   "jaxlib",
   "casadi",
   "torch",
+  "mujoco",
+  "robot_descriptions",
   "pytest",
   "idyntree",
   "icub-models",

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

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

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

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 "<robot" in s[:2048].lower():
+            return True
+        try:
+            return pathlib.Path(obj).suffix.lower() == ".urdf"
+        except Exception:
+            return False
+
+    return False
+
+
+def build_model_factory(description: Any, math) -> 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}"
+    )

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 <actuator> 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 <actuator> 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)