jaxsim 0.2.dev366__py3-none-any.whl → 0.2.dev376__py3-none-any.whl

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/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/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/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/terrain/__init__.py

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

{jaxsim-0.2.dev366.dist-info → jaxsim-0.2.dev376.dist-info}/METADATA

@@ -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.dist-info → jaxsim-0.2.dev376.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=ncT8dupmzd5NQtuHjYN5ljex6-14cJKuIIfJTJVSFHw,423
+jaxsim/_version.py,sha256=nAy-z_b2P4v0CvZeZkkzcsMlBO4LKnDdLnFE2Kq0a4U,423
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
 jaxsim/api/__init__.py,sha256=fNTCPUeDfOAizRd4RsW3Epv0sLTu0KJGoFRSEsi75VM,162
@@ -10,7 +10,7 @@ jaxsim/api/data.py,sha256=1AJyKjmXdsWEf_CkvOXRmvRsDZamG706mAAg1Gttll0,26773
 jaxsim/api/joint.py,sha256=q31Kp3Cqv-yTcxijjzbj_QADFnGQyjb2al9fYZtzedo,4763
 jaxsim/api/kin_dyn_parameters.py,sha256=G4mtSi8fElYe0yttLgsxSOPf7vcK-yqTu06Aa5SSrYg,26012
 jaxsim/api/link.py,sha256=LZVcQhQsTKsfR13KewFtEMYu4siVJl7mqoDwYsoFFes,9240
-jaxsim/api/model.py,sha256=mFdEwVuzIR7Lvj4oiIsA1n1oxpRZXWyB0IDzIxcG33Q,43689
+jaxsim/api/model.py,sha256=_waDChFZsHRbKOX6dkajRQxGkhqXFRMiMp3JXMqZ-MU,44089
 jaxsim/api/ode.py,sha256=rbSruK0Dkp09oBgHbB_-NrZS4o2tY9geK0yLLJfXzpM,9821
 jaxsim/api/ode_data.py,sha256=dwRFVPJ30XMmdUbPXEu7YxsQ97jZP4L4fd5ZzhrO5Ys,22184
 jaxsim/api/references.py,sha256=Lvskf17r619KKxwCJP7hAAty2kaXgDXJX1uKqoDIDgo,15483
@@ -29,8 +29,8 @@ jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=nqH6ofde6VWjfHihmYdXvDxKFChpOPH6AsoqkUI1Og0,2928
 jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=8sXc_tsDFWBYl8nesgFarYd3hA-PESLMrXsnR3Siz1Y,16400
-jaxsim/mujoco/model.py,sha256=0kG2GERxjVFqWZ1K3352rgUNfchB4kRtIrsvv4pS4oc,10766
+jaxsim/mujoco/loaders.py,sha256=IGhchLIzINXDm8nDIsy1qIrffoiOpr8Hs36JTOOYVhg,17392
+jaxsim/mujoco/model.py,sha256=GgDlr7fSSEnkyqziPLLQaippMr4wRhfTo39DstD9Qik,12088
 jaxsim/mujoco/visualizer.py,sha256=-qg26t5tleTva6zzQmc5SdnlC8XZ1ZAwZ_lDjdwHJ0A,4400
 jaxsim/parsers/__init__.py,sha256=sonYi-bBWAoB04kp1mxT4uIORxjb7SdZ0ukGPmVx98Y,44
 jaxsim/parsers/kinematic_graph.py,sha256=2B5gtUboiSVJIm6PegbwHb5g_iXltG0R9_7h8RJ-92M,23785
@@ -51,14 +51,14 @@ jaxsim/rbda/jacobian.py,sha256=9LGGy9ya5m5U0mBmV1NFH5XYZpEMYbx74qnYBvZs7Ok,6360
 jaxsim/rbda/rnea.py,sha256=DjwkvXQVUSUclM3Uy3UPZ2tao91R5dGd4o7TsS2qObI,7650
 jaxsim/rbda/soft_contacts.py,sha256=2EZ9Lw4nFWqXTMEeYsirl17H61s82SmTZllKVsP1Yek,10759
 jaxsim/rbda/utils.py,sha256=zpbFM2Iq8cntku0BFVu9nfEqZhInCWi9D2INT6MFEI8,5003
-jaxsim/terrain/__init__.py,sha256=dzekq9yyj3DKTsCARteqc81lAw3OSnl6EhXn8_Q_ozI,64
+jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=q0xkWqEShVq-p1j2abTLZq8sEhjyJwquxQKm80PaHhM,2161
 jaxsim/utils/__init__.py,sha256=tnQq1_CavdfeKaLYt3pmO7Jk4MU2RwwQU_qICkjyoTY,197
 jaxsim/utils/hashless.py,sha256=bFIwKeo9KiWwsY8QM55duEGGQOyyJ4jQyPcuqTLEp5k,297
 jaxsim/utils/jaxsim_dataclass.py,sha256=9FtzgXQPSa6AN26FleDgXGsJMse4dtVACDESrn0g3bw,11359
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim-0.2.dev366.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
-jaxsim-0.2.dev366.dist-info/METADATA,sha256=D-Ca1OQtpKyTen4FF58yJK_FdC9iAA2yYZj85blGCHg,7630
-jaxsim-0.2.dev366.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.dev366.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.dev366.dist-info/RECORD,,
+jaxsim-0.2.dev376.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
+jaxsim-0.2.dev376.dist-info/METADATA,sha256=eWcinEhSdothvx37-yafMhBddfT5s0fIo4veqCLNnI8,7630
+jaxsim-0.2.dev376.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.dev376.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.dev376.dist-info/RECORD,,