Import Mujco mjcf files

src/adam/casadi/computations.py

@@ -2,14 +2,16 @@
 # This software may be modified and distributed under the terms of the
 # GNU Lesser General Public License v2.1 or any later version.
 
+from pathlib import Path
+from typing import Union
+
 import casadi as cs
 import numpy as np
-from typing import Union
 
 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 MJModelFactory, Model, utils
 
 
 class KinDynComputations:
@@ -19,26 +21,70 @@ class KinDynComputations:
 
     def __init__(
         self,
-        urdfstring: str,
+        model_string: Union[str, Path],
         joints_name_list: list = None,
         root_link: str = "root_link",
         gravity: np.array = np.array([0.0, 0.0, -9.80665, 0.0, 0.0, 0.0]),
         f_opts: dict = dict(jit=False, jit_options=dict(flags="-Ofast"), cse=True),
     ) -> None:
         """
         Args:
-            urdfstring (str): either path or string of the urdf
+            model_string (Union[str, Path]): path or string of the URDF model, or the Mujoco XML file
             joints_name_list (list): list of the actuated joints
             root_link (str, optional): the first link. Defaults to 'root_link'.
         """
         math = SpatialMath()
-        factory = URDFModelFactory(path=urdfstring, math=math)
+        factory = utils.get_factory_from_string(model_string=model_string, 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
         self.g = gravity
         self.f_opts = f_opts
 
+    @staticmethod
+    def from_urdf(
+        urdf_string: Union[str, Path],
+        joints_name_list: list = None,
+        gravity: np.array = np.array([0.0, 0.0, -9.80665, 0.0, 0.0, 0.0]),
+    ) -> "KinDynComputations":
+        """Creates a KinDynComputations object from a URDF string
+
+        Args:
+            urdf_string (Union[str, Path]): The URDF string or path
+            joints_name_list (list, optional): List of the actuated joints. Defaults to None.
+            gravity (np.array, optional): The gravity vector. Defaults to np.array([0.0, 0.0, -9.80665, 0.0, 0.0, 0.0]).
+
+        Returns:
+            KinDynComputations: The KinDynComputations object
+        """
+        return KinDynComputations(
+            model_string=urdf_string,
+            joints_name_list=joints_name_list,
+            gravity=gravity,
+        )
+
+    @staticmethod
+    def from_mujoco_xml(
+        xml_string: Union[str, Path],
+        joints_name_list: list = None,
+        gravity: np.array = np.array([0.0, 0.0, -9.80665, 0.0, 0.0, 0.0]),
+    ) -> "KinDynComputations":
+        """Creates a KinDynComputations object from a Mujoco XML string
+
+        Args:
+            xml_string (Union[str, Path]): The Mujoco XML path
+            joints_name_list (list, optional): List of the actuated joints. Defaults to None.
+            gravity (np.array, optional): The gravity vector. Defaults to np.array([0.0, 0.0, -9.80665, 0.0, 0.0, 0.0]).
+
+        Returns:
+            KinDynComputations: The KinDynComputations object
+        """
+        return KinDynComputations(
+            model_string=xml_string,
+            joints_name_list=joints_name_list,
+            gravity=gravity,
+        )
+
     def set_frame_velocity_representation(
         self, representation: Representations
     ) -> None:

src/adam/core/rbd_algorithms.py

@@ -16,9 +16,8 @@ class RBDAlgorithms:
     def __init__(self, model: Model, math: SpatialMath) -> None:
         """
         Args:
-            urdfstring (str): path of the urdf
-            joints_name_list (list): list of the actuated joints
-            root_link (str, optional): the first link. Defaults to 'root_link'.
+            model (Model): The robot model
+            math (SpatialMath): The spatial math object
         """
 
         self.model = model

src/adam/jax/computations.py

@@ -2,38 +2,83 @@
 # This software may be modified and distributed under the terms of the
 # GNU Lesser General Public License v2.1 or any later version.
 
+from pathlib import Path
+from typing import Union
+
 import jax.numpy as jnp
 import numpy as np
 
 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, utils
 
 
 class KinDynComputations:
     """This is a small class that retrieves robot quantities using Jax for Floating Base systems."""
 
     def __init__(
         self,
-        urdfstring: str,
+        model_string: str,
         joints_name_list: list = None,
         root_link: str = "root_link",
         gravity: np.array = jnp.array([0, 0, -9.80665, 0, 0, 0]),
     ) -> None:
         """
         Args:
-            urdfstring (str): either path or string of the urdf
+            model_string (Union[str, Path]): the path or string of the URDF, or MuJoCo XML file
             joints_name_list (list): list of the actuated joints
             root_link (str, optional): the first link. Defaults to 'root_link'.
         """
         math = SpatialMath()
-        factory = URDFModelFactory(path=urdfstring, math=math)
+        factory = utils.get_factory_from_string(model_string, math)
         model = Model.build(factory=factory, joints_name_list=joints_name_list)
         self.rbdalgos = RBDAlgorithms(model=model, math=math)
         self.NDoF = self.rbdalgos.NDoF
         self.g = gravity
 
+    @staticmethod
+    def from_urdf(
+        urdf_string: Union[str, Path],
+        joints_name_list: list = None,
+        gravity: np.array = jnp.array([0, 0, -9.80665, 0, 0, 0]),
+    ):
+        """Creates a KinDynComputations object from a URDF string
+
+        Args:
+            urdf_string (Union[str, Path]): The URDF path or string
+            joints_name_list (list, optional): The actuated joints. Defaults to None.
+            root_link (str, optional): The root link. Defaults to "root_link".
+            gravity (np.array, optional): The gravity vector. Defaults to jnp.array([0, 0, -9.80665, 0, 0, 0]).
+
+        Returns:
+            KinDynComputations: The KinDynComputations object
+        """
+        return KinDynComputations(
+            model_string=urdf_string, joints_name_list=joints_name_list, gravity=gravity
+        )
+
+    @staticmethod
+    def from_mujoco_xml(
+        xml_string: Union[str, Path],
+        joints_name_list: list = None,
+        gravity: np.array = jnp.array([0, 0, -9.80665, 0, 0, 0]),
+    ):
+        """Creates a KinDynComputations object from a MuJoCo XML string
+
+        Args:
+            xml_string (Union[str, Path]): The MuJoCo XML path
+            joints_name_list (list, optional): The actuated joints. Defaults to None.
+            root_link (str, optional): The root link. Defaults to "root_link".
+            gravity (np.array, optional): The gravity vector. Defaults to jnp.array([0, 0, -9.80665, 0, 0, 0]).
+
+        Returns:
+            KinDynComputations: The KinDynComputations object
+        """
+        return KinDynComputations(
+            model_string=xml_string, joints_name_list=joints_name_list, gravity=gravity
+        )
+
     def set_frame_velocity_representation(
         self, representation: Representations
     ) -> None:

src/adam/model/__init__.py

@@ -1,4 +1,7 @@
-from .abc_factories import Joint, Link, ModelFactory, Inertial, Pose
+from .abc_factories import Inertia, Inertial, Joint, Link, ModelFactory, Pose
+from .mj_factories.mj_joint import MJJoint
+from .mj_factories.mj_link import MJLink
+from .mj_factories.mj_model import MJModelFactory
 from .model import Model
 from .std_factories.std_joint import StdJoint
 from .std_factories.std_link import StdLink

src/adam/model/abc_factories.py

@@ -54,33 +54,84 @@ class Joint(abc.ABC):
     Abstract base class for all joints.
     """
 
-    @abc.abstractmethod
-    def spatial_transform(self, q: npt.ArrayLike) -> npt.ArrayLike:
+    def homogeneous(self, q: npt.ArrayLike) -> npt.ArrayLike:
         """
         Args:
-            q (npt.ArrayLike): joint motion
+            q (npt.ArrayLike): joint value
 
         Returns:
-            npt.ArrayLike: spatial transform of the joint given q
+            npt.ArrayLike: the homogenous transform of a joint, given q
         """
-        pass
 
-    @abc.abstractmethod
-    def motion_subspace(self) -> npt.ArrayLike:
+        if self.type == "fixed":
+            xyz = self.origin.xyz
+            rpy = self.origin.rpy
+            return self.math.H_from_Pos_RPY(xyz, rpy)
+        elif self.type in ["revolute", "continuous"]:
+            return self.math.H_revolute_joint(
+                self.origin.xyz,
+                self.origin.rpy,
+                self.axis,
+                q,
+            )
+        elif self.type in ["prismatic"]:
+            return self.math.H_prismatic_joint(
+                self.origin.xyz,
+                self.origin.rpy,
+                self.axis,
+                q,
+            )
+
+    def spatial_transform(self, q: npt.ArrayLike) -> npt.ArrayLike:
         """
+        Args:
+            q (npt.ArrayLike): joint motion
+
         Returns:
-            npt.ArrayLike: motion subspace of the joint
+            npt.ArrayLike: spatial transform of the joint given q
         """
+        if self.type == "fixed":
+            return self.math.X_fixed_joint(self.origin.xyz, self.origin.rpy)
+        elif self.type in ["revolute", "continuous"]:
+            return self.math.X_revolute_joint(
+                self.origin.xyz, self.origin.rpy, self.axis, q
+            )
+        elif self.type in ["prismatic"]:
+            return self.math.X_prismatic_joint(
+                self.origin.xyz,
+                self.origin.rpy,
+                self.axis,
+                q,
+            )
 
-    @abc.abstractmethod
-    def homogeneous(self, q: npt.ArrayLike) -> npt.ArrayLike:
+    def motion_subspace(self) -> npt.ArrayLike:
         """
         Args:
-            q (npt.ArrayLike): joint value
+            joint (Joint): Joint
+
         Returns:
-            npt.ArrayLike: homogeneous transform given the joint value
+            npt.ArrayLike: motion subspace of the joint
         """
-        pass
+        if self.type == "fixed":
+            return self.math.vertcat(0, 0, 0, 0, 0, 0)
+        elif self.type in ["revolute", "continuous"]:
+            return self.math.vertcat(
+                0,
+                0,
+                0,
+                self.axis[0],
+                self.axis[1],
+                self.axis[2],
+            )
+        elif self.type in ["prismatic"]:
+            return self.math.vertcat(
+                self.axis[0],
+                self.axis[1],
+                self.axis[2],
+                0,
+                0,
+                0,
+            )
 
 
 @dataclasses.dataclass
@@ -130,22 +181,27 @@ class Link(abc.ABC):
     inertial: Inertial
     collisions: List
 
-    @abc.abstractmethod
     def spatial_inertia(self) -> npt.ArrayLike:
         """
         Returns:
             npt.ArrayLike: the 6x6 inertia matrix expressed at
                            the origin of the link (with rotation)
         """
-        pass
+        I = self.inertial.inertia
+        mass = self.inertial.mass
+        o = self.inertial.origin.xyz
+        rpy = self.inertial.origin.rpy
+        return self.math.spatial_inertia(I, mass, o, rpy)
 
-    @abc.abstractmethod
     def homogeneous(self) -> npt.ArrayLike:
         """
         Returns:
             npt.ArrayLike: the homogeneous transform of the link
         """
-        pass
+        return self.math.H_from_Pos_RPY(
+            self.inertial.origin.xyz,
+            self.inertial.origin.rpy,
+        )
 
 
 @dataclasses.dataclass

src/adam/model/mj_factories/mj_joint.py

@@ -0,0 +1,36 @@
+from typing import Dict, List, Optional
+
+from scipy.spatial.transform import Rotation
+
+from adam.core.spatial_math import SpatialMath
+from adam.model import Joint, Pose
+
+
+def convert_data(data: str) -> List[float]:
+    """Convert strings to floats."""
+    return list(map(float, data.split()))
+
+
+def wxyz_to_xyzw(quat: List[float]) -> List[float]:
+    """Convert quaternion from Mujoco (wxyz) to scipy (xyzw) convention."""
+    return [quat[1], quat[2], quat[3], quat[0]]
+
+
+class MJJoint(Joint):
+    """MuJoCo Joint class."""
+
+    def __init__(
+        self, joint: Dict, math: SpatialMath, idx: Optional[int] = None
+    ) -> None:
+        self.math = math
+        self.name = joint["name"]
+        self.parent = joint["parent"]
+        self.child = joint["child"]
+        self.type = joint["joint_type"]
+        self.axis = convert_data(joint["axis"])
+        joint_ori_xyz = convert_data(joint["origin"])
+        joint_quat = wxyz_to_xyzw(convert_data(joint["orientation"]))
+        joint_rpy = Rotation.from_quat(joint_quat).as_euler("xyz")
+        self.origin = Pose(xyz=joint_ori_xyz, rpy=joint_rpy)
+        self.limit = joint["limit"]
+        self.idx = idx