jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/mujoco/model.py

@@ -1,12 +1,20 @@
+from __future__ import annotations
+
 import functools
 import pathlib
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import mujoco as mj
 import numpy as np
 import numpy.typing as npt
+import xmltodict
 from scipy.spatial.transform import Rotation
 
+import jaxsim.typing as jtp
+
+HeightmapCallable = Callable[[jtp.FloatLike, jtp.FloatLike], jtp.FloatLike]
+
 
 class MujocoModelHelper:
     """
@@ -14,34 +22,118 @@ class MujocoModelHelper:
     """
 
     def __init__(self, model: mj.MjModel, data: mj.MjData | None = None) -> None:
-        """"""
+        """
+        Initialize the MujocoModelHelper object.
+
+        Args:
+            model: A Mujoco model object.
+            data: A Mujoco data object. If None, a new one will be created.
+        """
 
         self.model = model
         self.data = data if data is not None else mj.MjData(self.model)
 
-        # Populate the data with kinematics
+        # Populate the data with kinematics.
         mj.mj_forward(self.model, self.data)
 
-        # Keep the cache of this method local to improve GC
+        # Keep the cache of this method local to improve GC.
         self.mask_qpos = functools.cache(self._mask_qpos)
 
     @staticmethod
     def build_from_xml(
-        mjcf_description: str | pathlib.Path, assets: dict[str, Any] = None
-    ) -> "MujocoModelHelper":
-        """"""
+        mjcf_description: str | pathlib.Path,
+        assets: dict[str, Any] | None = None,
+        heightmap: HeightmapCallable | None = None,
+        heightmap_name: str = "terrain",
+        heightmap_radius_xy: tuple[float, float] = (1.0, 1.0),
+    ) -> MujocoModelHelper:
+        """
+        Build a Mujoco model from an MJCF description.
 
-        # Read the XML description if it's a path to file
+        Args:
+            mjcf_description:
+                A string containing the XML description of the Mujoco model
+                or a path to a file containing the XML description.
+            assets: An optional dictionary containing the assets of the model.
+            heightmap:
+                A function in two variables that returns the height of a terrain
+                in the specified coordinate point.
+            heightmap_name:
+                The default name of the heightmap in the MJCF description
+                to load the corresponding configuration.
+            heightmap_radius_xy:
+                The extension of the heightmap in the x-y surface corresponding to the
+                plane over which the grid of the sampled heightmap is generated.
+
+        Returns:
+            A MujocoModelHelper object.
+        """
+
+        # Read the XML description if it is a path to file.
         mjcf_description = (
             mjcf_description.read_text()
             if isinstance(mjcf_description, pathlib.Path)
             else mjcf_description
         )
 
-        # Create the Mujoco model from the XML and, optionally, the assets dictionary
-        model = mj.MjModel.from_xml_string(xml=mjcf_description, assets=assets)  # noqa
+        if heightmap is None:
+            hfield = None
+
+        else:
+
+            mjcf_description_dict = xmltodict.parse(xml_input=mjcf_description)
+
+            # Create a dictionary of all hfield configurations from the MJCF.
+            hfields = mjcf_description_dict["mujoco"]["asset"].get("hfield", [])
+            hfields = hfields if isinstance(hfields, list) else [hfields]
+            hfields_dict = {hfield["@name"]: hfield for hfield in hfields}
+
+            if heightmap_name not in hfields_dict:
+                raise ValueError(f"Heightmap '{heightmap_name}' not found in MJCF")
 
-        return MujocoModelHelper(model=model, data=mj.MjData(model))
+            hfield_element = hfields_dict[heightmap_name]
+
+            # Generate the hfield by sampling the heightmap function.
+            hfield = generate_hfield(
+                heightmap=heightmap,
+                samples_xy=(int(hfield_element["@nrow"]), int(hfield_element["@ncol"])),
+                radius_xy=heightmap_radius_xy,
+            )
+
+            # Update dynamically the '/asset/hfield[@name=heightmap_name]@size' attribute
+            # with the information of the sampled points.
+            # This is necessary for correctly rendering the heightmap over the
+            # specified xy area with the correct z elevation.
+            size = [float(el) for el in hfield_element["@size"].split(" ")]
+            size[0], size[1] = heightmap_radius_xy
+            size[2] = 1.0
+            # The following could be zero but Mujoco complains if it's exactly zero.
+            size[3] = max(0.000_001, -min(hfield))
+
+            # Replace the 'size' attribute.
+            hfields_dict[heightmap_name]["@size"] = " ".join(str(el) for el in size)
+
+            # Update the hfield elements of the original MJCF.
+            # Only the hfield corresponding to 'heightmap_name' was actually edited.
+            mjcf_description_dict["mujoco"]["asset"]["hfield"] = list(
+                hfields_dict.values()
+            )
+
+            # Serialize the updated MJCF to XML.
+            mjcf_description = xmltodict.unparse(
+                input_dict=mjcf_description_dict, pretty=True
+            )
+
+        # Create the Mujoco model from the XML and, optionally, the dictionary of assets.
+        model = mj.MjModel.from_xml_string(xml=mjcf_description, assets=assets)
+        data = mj.MjData(model)
+
+        # Store the sampled heightmap into the Mujoco model.
+        if heightmap is not None:
+            assert hfield is not None
+            model.hfield_data = hfield
+
+        return MujocoModelHelper(model=model, data=data)
 
     def time(self) -> float:
         """Return the simulation time."""
@@ -148,9 +240,9 @@ class MujocoModelHelper:
             raise ValueError("The orientation is not a valid element of SO(3)")
 
         W_Q_B = (
-            Rotation.from_matrix(orientation).as_quat(canonical=True)[
-                np.array([3, 0, 1, 2])
-            ]
+            Rotation.from_matrix(orientation).as_quat(
+                canonical=True, scalar_first=False
+            )
             if dcm
             else orientation
         )
@@ -162,17 +254,17 @@ class MujocoModelHelper:
     # ==================
 
     def number_of_joints(self) -> int:
-        """"""
+        """Return the number of joints in the model."""
 
         return self.model.njnt
 
     def number_of_dofs(self) -> int:
-        """"""
+        """Return the number of DoFs in the model."""
 
         return self.model.nq
 
     def joint_names(self) -> list[str]:
-        """"""
+        """Return the names of the joints in the model."""
 
         return [
             mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_JOINT, idx)
@@ -180,7 +272,7 @@ class MujocoModelHelper:
         ]
 
     def joint_dofs(self, joint_name: str) -> int:
-        """"""
+        """Return the number of DoFs of a joint."""
 
         if joint_name not in self.joint_names():
             raise ValueError(f"Joint '{joint_name}' not found")
@@ -188,7 +280,7 @@ class MujocoModelHelper:
         return self.data.joint(joint_name).qpos.size
 
     def joint_position(self, joint_name: str) -> npt.NDArray:
-        """"""
+        """Return the position of a joint."""
 
         if joint_name not in self.joint_names():
             raise ValueError(f"Joint '{joint_name}' not found")
@@ -196,7 +288,7 @@ class MujocoModelHelper:
         return self.data.joint(joint_name).qpos
 
     def joint_positions(self, joint_names: list[str] | None = None) -> npt.NDArray:
-        """"""
+        """Return the positions of the joints."""
 
         joint_names = joint_names if joint_names is not None else self.joint_names()
 
@@ -207,7 +299,7 @@ class MujocoModelHelper:
     def set_joint_position(
         self, joint_name: str, position: npt.NDArray | float
     ) -> None:
-        """"""
+        """Set the position of a joint."""
 
         position = np.atleast_1d(np.array(position).squeeze())
 
@@ -224,9 +316,9 @@ class MujocoModelHelper:
         self.data.qpos[sl] = position
 
     def set_joint_positions(
-        self, joint_names: list[str], positions: npt.NDArray | list[npt.NDArray]
+        self, joint_names: Sequence[str], positions: npt.NDArray | list[npt.NDArray]
     ) -> None:
-        """"""
+        """Set the positions of multiple joints."""
 
         mask = self.mask_qpos(joint_names=tuple(joint_names))
         self.data.qpos[mask] = positions
@@ -236,12 +328,12 @@ class MujocoModelHelper:
     # ==================
 
     def number_of_bodies(self) -> int:
-        """"""
+        """Return the number of bodies in the model."""
 
         return self.model.nbody
 
     def body_names(self) -> list[str]:
-        """"""
+        """Return the names of the bodies in the model."""
 
         return [
             mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_BODY, idx)
@@ -249,7 +341,7 @@ class MujocoModelHelper:
         ]
 
     def body_position(self, body_name: str) -> npt.NDArray:
-        """"""
+        """Return the position of a body."""
 
         if body_name not in self.body_names():
             raise ValueError(f"Body '{body_name}' not found")
@@ -257,7 +349,7 @@ class MujocoModelHelper:
         return self.data.body(body_name).xpos
 
     def body_orientation(self, body_name: str, dcm: bool = False) -> npt.NDArray:
-        """"""
+        """Return the orientation of a body."""
 
         if body_name not in self.body_names():
             raise ValueError(f"Body '{body_name}' not found")
@@ -271,12 +363,12 @@ class MujocoModelHelper:
     # ======================
 
     def number_of_geometries(self) -> int:
-        """"""
+        """Return the number of geometries in the model."""
 
         return self.model.ngeom
 
     def geometry_names(self) -> list[str]:
-        """"""
+        """Return the names of the geometries in the model."""
 
         return [
             mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_GEOM, idx)
@@ -284,7 +376,7 @@ class MujocoModelHelper:
         ]
 
     def geometry_position(self, geometry_name: str) -> npt.NDArray:
-        """"""
+        """Return the position of a geometry."""
 
         if geometry_name not in self.geometry_names():
             raise ValueError(f"Geometry '{geometry_name}' not found")
@@ -294,7 +386,7 @@ class MujocoModelHelper:
     def geometry_orientation(
         self, geometry_name: str, dcm: bool = False
     ) -> npt.NDArray:
-        """"""
+        """Return the orientation of a geometry."""
 
         if geometry_name not in self.geometry_names():
             raise ValueError(f"Geometry '{geometry_name}' not found")
@@ -304,8 +396,8 @@ class MujocoModelHelper:
         if dcm:
             return R
 
-        q_xyzw = Rotation.from_matrix(R).as_quat(canonical=True)
-        return q_xyzw[[3, 0, 1, 2]]
+        q_xyzw = Rotation.from_matrix(R).as_quat(canonical=True, scalar_first=False)
+        return q_xyzw
 
     # ===============
     # Private methods
@@ -346,7 +438,45 @@ class MujocoModelHelper:
                             for i in range(self.joint_dofs(joint_name=joint_name))
                         ]
                     )
-                    for idx, joint_name in zip(idxs, joint_names)
+                    for idx, joint_name in zip(idxs, joint_names, strict=True)
                 ]
             ).squeeze()
         )
+
+
+def generate_hfield(
+    heightmap: HeightmapCallable,
+    samples_xy: tuple[int, int] = (11, 11),
+    radius_xy: tuple[float, float] = (1.0, 1.0),
+) -> npt.NDArray:
+    """
+    Generate an array with elevation points sampled from a heightmap function.
+
+    The map will have the following format:
+    ```
+    heightmap[0, 0] heightmap[0, 1] ... heightmap[0, size[1]-1]
+    heightmap[1, 0] heightmap[1, 1] ... heightmap[1, size[1]-1]
+    ...
+    heightmap[size[0]-1, 0] heightmap[size[0]-1, 1] ... heightmap[size[0]-1, size[1]-1]
+    ```
+
+    Args:
+        heightmap:
+            A function that takes two arguments (x, y) and returns the height
+            at that point.
+        samples_xy: A tuple of two integers representing the size of the grid.
+        radius_xy:
+            A tuple of two floats representing extension of the heightmap in the
+            x-y surface corresponding to the area over which the grid of the sampled
+            heightmap is generated.
+
+    Returns:
+        A flat array of the sampled terrain heightmap.
+    """
+
+    # Generate the grid.
+    x = np.linspace(-radius_xy[0], radius_xy[0], samples_xy[0])
+    y = np.linspace(-radius_xy[1], radius_xy[1], samples_xy[1])
+
+    # Generate the heightmap.
+    return np.array([[heightmap(xi, yi) for xi in x] for yi in y]).flatten()

jaxsim/mujoco/utils.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+
+import mujoco as mj
+import numpy as np
+import numpy.typing as npt
+from scipy.spatial.transform import Rotation
+
+from .model import MujocoModelHelper
+
+
+def mujoco_data_from_jaxsim(
+    mujoco_model: mj.MjModel,
+    jaxsim_model,
+    jaxsim_data,
+    mujoco_data: mj.MjData | None = None,
+    update_removed_joints: bool = True,
+) -> mj.MjData:
+    """
+    Create a Mujoco data object from a JaxSim model and data objects.
+
+    Args:
+        mujoco_model: The Mujoco model object corresponding to the JaxSim model.
+        jaxsim_model: The JaxSim model object from which the Mujoco model was created.
+        jaxsim_data: The JaxSim data object containing the state of the model.
+        mujoco_data: An optional Mujoco data object. If None, a new one will be created.
+        update_removed_joints:
+            If True, the positions of the joints that have been removed during the
+            model reduction process will be set to their initial values.
+
+    Returns:
+        The Mujoco data object containing the state of the JaxSim model.
+
+    Note:
+        This method is useful to initialize a Mujoco data object used for visualization
+        with the state of a JaxSim model. In particular, this function takes care of
+        initializing the positions of the joints that have been removed during the
+        model reduction process. After the initial creation of the Mujoco data object,
+        it's faster to update the state using an external MujocoModelHelper object.
+    """
+
+    # The package `jaxsim.mujoco` is supposed to be jax-independent.
+    # We import all the JaxSim resources privately.
+    import jaxsim.api as js
+
+    if not isinstance(jaxsim_model, js.model.JaxSimModel):
+        raise ValueError("The `jaxsim_model` argument must be a JaxSimModel object.")
+
+    if not isinstance(jaxsim_data, js.data.JaxSimModelData):
+        raise ValueError("The `jaxsim_data` argument must be a JaxSimModelData object.")
+
+    # Create the helper to operate on the Mujoco model and data.
+    model_helper = MujocoModelHelper(model=mujoco_model, data=mujoco_data)
+
+    # If the model is fixed-base, the Mujoco model won't have the joint corresponding
+    # to the floating base, and the helper would raise an exception.
+    if jaxsim_model.floating_base():
+
+        # Set the model position.
+        model_helper.set_base_position(position=np.array(jaxsim_data.base_position()))
+
+        # Set the model orientation.
+        model_helper.set_base_orientation(
+            orientation=np.array(jaxsim_data.base_orientation())
+        )
+
+    # Set the joint positions.
+    if jaxsim_model.dofs() > 0:
+
+        model_helper.set_joint_positions(
+            joint_names=list(jaxsim_model.joint_names()),
+            positions=np.array(
+                jaxsim_data.joint_positions(
+                    model=jaxsim_model, joint_names=jaxsim_model.joint_names()
+                )
+            ),
+        )
+
+    # Updating these joints is not necessary after the first time.
+    # Users can disable this update after initialization.
+    if update_removed_joints:
+
+        # Create a dictionary with the joints that have been removed for various reasons
+        # (like link lumping due to model reduction).
+        joints_removed_dict = {
+            j.name: j
+            for j in jaxsim_model.description._joints_removed
+            if j.name not in set(jaxsim_model.joint_names())
+        }
+
+        # Set the positions of the removed joints.
+        _ = [
+            model_helper.set_joint_position(
+                position=joints_removed_dict[joint_name].initial_position,
+                joint_name=joint_name,
+            )
+            # Select all original joint that have been removed from the JaxSim model
+            # that are still present in the Mujoco model.
+            for joint_name in joints_removed_dict
+            if joint_name in model_helper.joint_names()
+        ]
+
+    # Return the mujoco data with updated kinematics.
+    mj.mj_forward(mujoco_model, model_helper.data)
+
+    return model_helper.data
+
+
+@dataclasses.dataclass
+class MujocoCamera:
+    """
+    Helper class storing parameters of a Mujoco camera.
+
+    Refer to the official documentation for more details:
+    https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-camera
+    """
+
+    mode: str = "fixed"
+
+    target: str | None = None
+    fovy: str = "45"
+    pos: str = "0 0 0"
+
+    quat: str | None = None
+    axisangle: str | None = None
+    xyaxes: str | None = None
+    zaxis: str | None = None
+    euler: str | None = None
+
+    name: str | None = None
+
+    @classmethod
+    def build(cls, **kwargs) -> MujocoCamera:
+        """
+        Build a Mujoco camera from a dictionary.
+        """
+
+        if not all(isinstance(value, str) for value in kwargs.values()):
+            raise ValueError(f"Values must be strings: {kwargs}")
+
+        return cls(**kwargs)
+
+    @staticmethod
+    def build_from_target_view(
+        camera_name: str,
+        lookat: Sequence[float | int] | npt.NDArray = (0, 0, 0),
+        distance: float | int | npt.NDArray = 3,
+        azimuth: float | int | npt.NDArray = 90,
+        elevation: float | int | npt.NDArray = -45,
+        fovy: float | int | npt.NDArray = 45,
+        degrees: bool = True,
+        **kwargs,
+    ) -> MujocoCamera:
+        """
+        Create a custom camera that looks at a target point.
+
+        Note:
+            The choice of the parameters is easier if we imagine to consider a target
+            frame `T` whose origin is located over the lookat point and having the same
+            orientation of the world frame `W`. We also introduce a camera frame `C`
+            whose origin is located over the lower-left corner of the image, and having
+            the x-axis pointing right and the y-axis pointing up in image coordinates.
+            The camera renders what it sees in the -z direction of frame `C`.
+
+        Args:
+            camera_name: The name of the camera.
+            lookat: The target point to look at (origin of `T`).
+            distance:
+                The distance from the target point (displacement between the origins
+                of `T` and `C`).
+            azimuth:
+                The rotation around z of the camera. With an angle of 0, the camera
+                would loot at the target point towards the positive x-axis of `T`.
+            elevation:
+                The rotation around the x-axis of the camera frame `C`. Note that if
+                you want to lift the view angle, the elevation is negative.
+            fovy: The field of view of the camera.
+            degrees: Whether the angles are in degrees or radians.
+            **kwargs: Additional camera parameters.
+
+        Returns:
+            The custom camera.
+        """
+
+        # Start from a frame whose origin is located over the lookat point.
+        # We initialize a -90 degrees rotation around the z-axis because due to
+        # the default camera coordinate system (x pointing right, y pointing up).
+        W_H_C = np.eye(4)
+        W_H_C[0:3, 3] = np.array(lookat)
+        W_H_C[0:3, 0:3] = Rotation.from_euler(
+            seq="ZX", angles=[-90, 90], degrees=True
+        ).as_matrix()
+
+        # Process the azimuth.
+        R_az = Rotation.from_euler(seq="Y", angles=azimuth, degrees=degrees).as_matrix()
+        W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_az
+
+        # Process elevation.
+        R_el = Rotation.from_euler(
+            seq="X", angles=elevation, degrees=degrees
+        ).as_matrix()
+        W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_el
+
+        # Process distance.
+        tf_distance = np.eye(4)
+        tf_distance[2, 3] = distance
+        W_H_C = W_H_C @ tf_distance
+
+        # Extract the position and the quaternion.
+        p = W_H_C[0:3, 3]
+        Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat(scalar_first=True)
+
+        return MujocoCamera.build(
+            name=camera_name,
+            mode="fixed",
+            fovy=f"{fovy if degrees else np.rad2deg(fovy)}",
+            pos=" ".join(p.astype(str).tolist()),
+            quat=" ".join(Q.astype(str).tolist()),
+            **kwargs,
+        )
+
+    def asdict(self) -> dict[str, str]:
+        """
+        Convert the camera to a dictionary.
+        """
+        return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}