jaxsim 0.4.1.dev6__tar.gz → 0.4.1.dev11__tar.gz

{jaxsim-0.4.1.dev6 → jaxsim-0.4.1.dev11}/PKG-INFO

@@ -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 → jaxsim-0.4.1.dev11}/src/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-0.4.1.dev6 → jaxsim-0.4.1.dev11}/src/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-0.4.1.dev6 → jaxsim-0.4.1.dev11}/src/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 → jaxsim-0.4.1.dev11}/src/jaxsim.egg-info/PKG-INFO

@@ -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 → jaxsim-0.4.1.dev11}/tests/test_simulations.py

@@ -6,7 +6,7 @@ import jaxsim.api as js
 import jaxsim.integrators
 import jaxsim.rbda
 from jaxsim import VelRepr
-from jaxsim.rbda.contacts.soft import SoftContactsParams
+from jaxsim.utils import Mutability
 
 
 def test_box_with_external_forces(
@@ -102,23 +102,19 @@ def test_box_with_zero_gravity(
 
     model = jaxsim_model_box
 
+    # Move the terrain (almost) infinitely far away from the box.
+    with model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
+        model.terrain = jaxsim.terrain.FlatTerrain.build(height=-1e9)
+
     # Split the PRNG key.
-    _, subkey, subkey2 = jax.random.split(prng_key, num=3)
+    _, subkey = jax.random.split(prng_key, num=2)
 
     # Build the data of the model.
     data0 = js.data.JaxSimModelData.build(
         model=model,
-        base_position=jax.random.uniform(subkey2, shape=(3,)),
+        base_position=jax.random.uniform(subkey, shape=(3,)),
         velocity_representation=velocity_representation,
         standard_gravity=0.0,
-        contacts_params=SoftContactsParams.build(K=0.0, D=0.0, mu=0.0),
-    )
-
-    # Generate a random linear force.
-    L_f = (
-        jax.random.uniform(subkey, shape=(model.number_of_links(), 6))
-        .at[:, 3:]
-        .set(jnp.zeros(3))
     )
 
     # Initialize a references object that simplifies handling external forces.
@@ -129,13 +125,29 @@ def test_box_with_zero_gravity(
     )
 
     # Apply a link forces to the base link.
-    references = references.apply_link_forces(
-        forces=jnp.atleast_2d(L_f),
-        link_names=model.link_names(),
-        model=model,
-        data=data0,
-        additive=False,
-    )
+    with references.switch_velocity_representation(jaxsim.VelRepr.Mixed):
+
+        # Generate a random linear force.
+        # We enforce them to be the same for all velocity representations so that
+        # we can compare their outcomes.
+        LW_f = 10.0 * (
+            jax.random.uniform(jax.random.key(0), shape=(model.number_of_links(), 6))
+            .at[:, 3:]
+            .set(jnp.zeros(3))
+        )
+
+        # Note that the context manager does not switch back the newly created
+        # `references` (that is not the yielded object) to the original representation.
+        # In the simulation loop below, we need to make sure that we switch both `data`
+        # and `references` to the same representation before extracting the information
+        # passed to the step function.
+        references = references.apply_link_forces(
+            forces=jnp.atleast_2d(LW_f),
+            link_names=model.link_names(),
+            model=model,
+            data=data0,
+            additive=False,
+        )
 
     # Create the integrator.
     integrator = jaxsim.integrators.fixed_step.RungeKutta4SO3.build(
@@ -145,8 +157,7 @@ def test_box_with_zero_gravity(
     )
 
     # Initialize the integrator.
-    tf = 1.0
-    dt = 0.010
+    tf, dt = 1.0, 0.010
     T = jnp.arange(start=0, stop=tf * 1e9, step=dt * 1e9, dtype=int)
     integrator_state = integrator.init(x0=data0.state, t0=0.0, dt=dt)
 
@@ -156,19 +167,28 @@ def test_box_with_zero_gravity(
     # ... and step the simulation.
     for t_ns in T:
 
-        data, integrator_state = js.model.step(
-            model=model,
-            data=data,
-            dt=dt,
-            integrator=integrator,
-            integrator_state=integrator_state,
-            link_forces=references.link_forces(model=model, data=data),
-        )
+        assert data.time() == t_ns / 1e9
+
+        with (
+            data.switch_velocity_representation(velocity_representation),
+            references.switch_velocity_representation(velocity_representation),
+        ):
+
+            data, integrator_state = js.model.step(
+                model=model,
+                data=data,
+                dt=dt,
+                integrator=integrator,
+                integrator_state=integrator_state,
+                link_forces=references.link_forces(model=model, data=data),
+            )
+
+    # Check the final simulation time.
+    assert data.time() == T[-1] / 1e9 + dt
 
     # Check that the box moved as expected.
-    assert data.time() == t_ns / 1e9 + dt
     assert data.base_position() == pytest.approx(
         data0.base_position()
-        + 0.5 * L_f[:, :3].squeeze() / js.model.total_mass(model=model) * tf**2,
+        + 0.5 * LW_f[:, :3].squeeze() / js.model.total_mass(model=model) * tf**2,
         abs=1e-3,
     )