jaxsim 0.4.1.dev6__py3-none-any.whl → 0.4.1.dev11__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.1.dev6'
-__version_tuple__ = version_tuple = (0, 4, 1, 'dev6')
+__version__ = version = '0.4.1.dev11'
+__version_tuple__ = version_tuple = (0, 4, 1, 'dev11')

jaxsim/mujoco/loaders.py

@@ -160,6 +160,7 @@ class RodModelToMjcf:
         considered_joints: list[str] | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
+        heightmap_samples_xy: tuple[int, int] = (101, 101),
         cameras: list[dict[str, str]] | dict[str, str] | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """
@@ -170,10 +171,11 @@ class RodModelToMjcf:
             considered_joints: The list of joint names to consider in the conversion.
             plane_normal: The normal vector of the plane.
             heightmap: Whether to generate a heightmap.
+            heightmap_samples_xy: The number of points in the heightmap grid.
             cameras: The list of cameras to add to the scene.
 
         Returns:
-            tuple: A tuple containing the MJCF string and the assets dictionary.
+            tuple: A tuple containing the MJCF string and the dictionary of assets.
         """
 
         # -------------------------------------
@@ -404,9 +406,11 @@ class RodModelToMjcf:
                 asset_element,
                 "hfield",
                 name="terrain",
-                nrow="100",
-                ncol="100",
-                size="5 5 1 1",
+                nrow=f"{int(heightmap_samples_xy[0])}",
+                ncol=f"{int(heightmap_samples_xy[1])}",
+                # The following 'size' is a placeholder, it is updated dynamically
+                # when a hfield/heightmap is stored into MjData.
+                size="1 1 1 1",
             )
             if heightmap
             else None

jaxsim/mujoco/model.py

@@ -7,6 +7,7 @@ from typing import Any, Callable
 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
@@ -42,16 +43,27 @@ class 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 XML description and an optional assets dictionary.
+        Build a Mujoco model from an MJCF description.
 
         Args:
-            mjcf_description: A string containing the XML description of the Mujoco model
+            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
+            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.
         """
@@ -63,15 +75,61 @@ class MujocoModelHelper:
             else mjcf_description
         )
 
-        # Create the Mujoco model from the XML and, optionally, the assets dictionary.
+        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")
+
+            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
+            size[3] = max(0, -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)
 
-        if heightmap:
-            nrow = model.hfield_nrow.item()
-            ncol = model.hfield_ncol.item()
-            new_hfield = generate_hfield(heightmap, (nrow, ncol))
-            model.hfield_data = new_hfield
+        # 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)
 
@@ -385,10 +443,13 @@ class MujocoModelHelper:
 
 
 def generate_hfield(
-    heightmap: HeightmapCallable, size: tuple[int, int] = (10, 10)
+    heightmap: HeightmapCallable,
+    samples_xy: tuple[int, int] = (11, 11),
+    radius_xy: tuple[float, float] = (1.0, 1.0),
 ) -> npt.NDArray:
     """
-    Generates a numpy array representing the heightmap of
+    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]
@@ -398,17 +459,22 @@ def generate_hfield(
     ```
 
     Args:
-        heightmap: A function that takes two arguments (x, y) and returns the height
+        heightmap:
+            A function that takes two arguments (x, y) and returns the height
             at that point.
-        size: A tuple of two integers representing the size of the grid.
+        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:
-        np.ndarray: The terrain heightmap
+        A flat array of the sampled terrain heightmap.
     """
 
     # Generate the grid.
-    x = np.linspace(0, 1, size[0])
-    y = np.linspace(0, 1, size[1])
+    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-0.4.1.dev6.dist-info → jaxsim-0.4.1.dev11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.1.dev6
+Version: 0.4.1.dev11
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.1.dev6.dist-info → jaxsim-0.4.1.dev11.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
-jaxsim/_version.py,sha256=bHDDEq643CuzU_d4ZEpTqbsQTlTvvqbGfOA9vrbclqI,424
+jaxsim/_version.py,sha256=7ZFpPaz7M9u-cSnMJ8POHhkQ3_j7elMp-k9w58-1Rgw,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
@@ -31,8 +31,8 @@ jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=_5kSnfkS6_vxvjxdw50KeXMjvW8e1OGaumUlk1iGJgc,2969
 jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=FuGPPg9Iq2xLZSqIFzlY7YfPR5uy9rbGlb0cU4gio5A,21000
-jaxsim/mujoco/model.py,sha256=mObFS77EY97sTFdlFrl69_gf9S_0FHO0W3_hevxQNws,13511
+jaxsim/mujoco/loaders.py,sha256=hwT8cjMEZmVaPMs4RzS4UEgC0_c9Od9817UcbNANPBc,21345
+jaxsim/mujoco/model.py,sha256=EwUPg9BsNv1B7TdDfjZCpC022lDR16AyIAajPJGH7NU,16357
 jaxsim/mujoco/visualizer.py,sha256=9jfKXkuoaW7Ppo_0m8dogD3SxH13K2strFNDVLtG3hA,5154
 jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/parsers/kinematic_graph.py,sha256=88d0EmndVJWdcyFJsW25S78Z8F04cUt08RQMyoil1Xw,34734
@@ -61,8 +61,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=fLl1tY3DDb3lpIhG6BPqA5W34hM84oFzL-5cuz8k-68,11379
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=GOJQCJc5zwzoEGZB62wnWWGvUUQlXvDxz_A2Q-hFv7c,4027
-jaxsim-0.4.1.dev6.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.1.dev6.dist-info/METADATA,sha256=SHBMqYI6b449ER-BGP0ANpyrW84OjD41JgtMbfw41Ss,16778
-jaxsim-0.4.1.dev6.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-jaxsim-0.4.1.dev6.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.1.dev6.dist-info/RECORD,,
+jaxsim-0.4.1.dev11.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.1.dev11.dist-info/METADATA,sha256=RQcmfKhS454m3fgEmDhJP00XtQFwnxWeXaBkTWiUKaQ,16779
+jaxsim-0.4.1.dev11.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+jaxsim-0.4.1.dev11.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.1.dev11.dist-info/RECORD,,