jaxsim 0.2.dev366__tar.gz → 0.2.dev376__tar.gz

{jaxsim-0.2.dev366 → jaxsim-0.2.dev376}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev366
+Version: 0.2.dev376
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo

{jaxsim-0.2.dev366 → jaxsim-0.2.dev376}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.dev366'
-__version_tuple__ = version_tuple = (0, 2, 'dev366')
+__version__ = version = '0.2.dev376'
+__version_tuple__ = version_tuple = (0, 2, 'dev376')

{jaxsim-0.2.dev366 → jaxsim-0.2.dev376}/src/jaxsim/api/model.py

@@ -52,6 +52,8 @@ class JaxSimModel(JaxsimDataclass):
     def build_from_model_description(
         model_description: str | pathlib.Path | rod.Model,
         model_name: str | None = None,
+        *,
+        terrain: jaxsim.terrain.Terrain | None = None,
         is_urdf: bool | None = None,
         considered_joints: list[str] | None = None,
     ) -> JaxSimModel:
@@ -65,6 +67,8 @@ class JaxSimModel(JaxsimDataclass):
             model_name:
                 The optional name of the model that overrides the one in
                 the description.
+            terrain:
+                The optional terrain to consider.
             is_urdf:
                 Whether the model description is a URDF or an SDF. This is
                 automatically inferred if the model description is a path to a file.
@@ -92,7 +96,9 @@ class JaxSimModel(JaxsimDataclass):
 
         # Build the model
         model = JaxSimModel.build(
-            model_description=intermediate_description, model_name=model_name
+            model_description=intermediate_description,
+            model_name=model_name,
+            terrain=terrain,
         )
 
         # Store the origin of the model, in case downstream logic needs it
@@ -105,6 +111,8 @@ class JaxSimModel(JaxsimDataclass):
     def build(
         model_description: jaxsim.parsers.descriptions.ModelDescription,
         model_name: str | None = None,
+        *,
+        terrain: jaxsim.terrain.Terrain | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from an intermediate model description.
@@ -115,6 +123,8 @@ class JaxSimModel(JaxsimDataclass):
                 of the model.
             model_name:
                 The optional name of the model overriding the physics model name.
+            terrain:
+                The optional terrain to consider.
 
         Returns:
             The built Model object.
@@ -130,6 +140,7 @@ class JaxSimModel(JaxsimDataclass):
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
+            terrain=terrain or JaxSimModel.__dataclass_fields__["terrain"].default,
         )
 
         return model

{jaxsim-0.2.dev366 → jaxsim-0.2.dev376}/src/jaxsim/mujoco/loaders.py

@@ -129,6 +129,8 @@ class RodModelToMjcf:
     def convert(
         rod_model: rod.Model,
         considered_joints: list[str] | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -198,8 +200,6 @@ class RodModelToMjcf:
         )
 
         urdf_string = ET.tostring(root, pretty_print=True).decode()
-        # print(urdf_string)
-        # raise
 
         # ------------------------------
         # Post-process all dummy visuals
@@ -358,6 +358,19 @@ class RodModelToMjcf:
             texuniform="true",
         )
 
+        _ = (
+            ET.SubElement(
+                asset_element,
+                "hfield",
+                name="terrain",
+                nrow="100",
+                ncol="100",
+                size="5 5 1 1",
+            )
+            if heightmap
+            else None
+        )
+
         # ----------------------------------
         # Populate the scene with the assets
         # ----------------------------------
@@ -368,12 +381,14 @@ class RodModelToMjcf:
             worldbody_scene_element,
             "geom",
             name="floor",
-            type="plane",
+            type="plane" if not heightmap else "hfield",
             size="0 0 0.05",
             material="plane_material",
             condim="3",
             contype="1",
             conaffinity="1",
+            zaxis=" ".join(map(str, plane_normal)),
+            **({"hfield": "terrain"} if heightmap else {}),
         )
 
         _ = ET.SubElement(
@@ -407,13 +422,14 @@ class RodModelToMjcf:
             raise RuntimeError("Failed to find the <worldbody> element of the model")
 
         # Camera attached to the model
+        # It can be manually copied from `python -m mujoco.viewer --mjcf=<URDF_PATH>`
         _ = ET.SubElement(
             worldbody_element,
             "camera",
             name="track",
             mode="trackcom",
-            pos="1 0 5",
-            zaxis="0 0 1",
+            pos="1.930 -2.279 0.556",
+            xyaxes="0.771 0.637 0.000 -0.116 0.140 0.983",
             fovy="60",
         )
 
@@ -449,6 +465,8 @@ class UrdfToMjcf:
         urdf: str | pathlib.Path,
         considered_joints: list[str] | None = None,
         model_name: str | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -461,7 +479,10 @@ class UrdfToMjcf:
 
         # Convert the ROD model to MJCF.
         return RodModelToMjcf.convert(
-            rod_model=rod_model, considered_joints=considered_joints
+            rod_model=rod_model,
+            considered_joints=considered_joints,
+            plane_normal=plane_normal,
+            heightmap=heightmap,
         )
 
 
@@ -471,6 +492,8 @@ class SdfToMjcf:
         sdf: str | pathlib.Path,
         considered_joints: list[str] | None = None,
         model_name: str | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -483,5 +506,8 @@ class SdfToMjcf:
 
         # Convert the ROD model to MJCF.
         return RodModelToMjcf.convert(
-            rod_model=rod_model, considered_joints=considered_joints
+            rod_model=rod_model,
+            considered_joints=considered_joints,
+            plane_normal=plane_normal,
+            heightmap=heightmap,
         )

{jaxsim-0.2.dev366 → jaxsim-0.2.dev376}/src/jaxsim/mujoco/model.py

@@ -1,12 +1,14 @@
 import functools
 import pathlib
-from typing import Any
+from typing import Any, Callable
 
 import mujoco as mj
 import numpy as np
 import numpy.typing as npt
 from scipy.spatial.transform import Rotation
 
+HeightmapCallable = Callable[[jtp.FloatLike, jtp.FloatLike], jtp.FloatLike]
+
 
 class MujocoModelHelper:
     """
@@ -27,7 +29,9 @@ class MujocoModelHelper:
 
     @staticmethod
     def build_from_xml(
-        mjcf_description: str | pathlib.Path, assets: dict[str, Any] = None
+        mjcf_description: str | pathlib.Path,
+        assets: dict[str, Any] = None,
+        heightmap: HeightmapCallable | None = None,
     ) -> "MujocoModelHelper":
         """"""
 
@@ -40,6 +44,13 @@ class MujocoModelHelper:
 
         # Create the Mujoco model from the XML and, optionally, the assets dictionary
         model = mj.MjModel.from_xml_string(xml=mjcf_description, assets=assets)  # noqa
+        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
 
         return MujocoModelHelper(model=model, data=mj.MjData(model))
 
@@ -350,3 +361,33 @@ class MujocoModelHelper:
                 ]
             ).squeeze()
         )
+
+
+def generate_hfield(
+    heightmap: HeightmapCallable, size: tuple[int, int] = (10, 10)
+) -> npt.NDArray:
+    """
+    Generates a numpy array representing the heightmap of
+    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.
+        size: A tuple of two integers representing the size of the grid.
+
+    Returns:
+        np.ndarray: The terrain heightmap
+    """
+
+    # Generate the grid.
+    x = np.linspace(0, 1, size[0])
+    y = np.linspace(0, 1, size[1])
+
+    # Generate the heightmap.
+    return np.array([[heightmap(xi, yi) for xi in x] for yi in y]).flatten()

jaxsim-0.2.dev376/src/jaxsim/terrain/__init__.py

@@ -0,0 +1,2 @@
+from . import terrain
+from .terrain import FlatTerrain, PlaneTerrain, Terrain

{jaxsim-0.2.dev366 → jaxsim-0.2.dev376}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev366
+Version: 0.2.dev376
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo

jaxsim-0.2.dev366/src/jaxsim/terrain/__init__.py

@@ -1,2 +0,0 @@
-from . import terrain
-from .terrain import FlatTerrain, Terrain