jaxsim 0.4.3.dev177__py3-none-any.whl → 0.4.3.dev186__py3-none-any.whl

jaxsim/__init__.py

@@ -8,21 +8,35 @@ def _jnp_options() -> None:
 
     import jax
 
-    # Enable by default 64bit precision in JAX.
-    if os.environ.get("JAX_ENABLE_X64", "1") != "0":
-
-        logging.info("Enabling JAX to use 64bit precision")
+    # Check if running on TPU
+    is_tpu = jax.devices()[0].platform == "tpu"
+
+    # Enable by default 64-bit precision to get accurate physics.
+    # Users can enforce 32-bit precision by setting the following variable to 0.
+    use_x64 = os.environ.get("JAX_ENABLE_X64", "1") != "0"
+
+    # Notify the user if unsupported 64-bit precision was enforced on TPU.
+    if is_tpu and use_x64:
+        msg = "64-bit precision is not allowed on TPU. Enforcing 32bit precision."
+        logging.warning(msg)
+        use_x64 = False
+
+    # Enable 64-bit precision in JAX.
+    if use_x64:
+        logging.info("Enabling JAX to use 64-bit precision")
         jax.config.update("jax_enable_x64", True)
 
         import jax.numpy as jnp
         import numpy as np
 
+        # Verify that 64-bit precision is correctly set.
         if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
-            logging.warning("Failed to enable 64bit precision in JAX")
+            logging.warning("Failed to enable 64-bit precision in JAX")
 
+    # Warn about experimental usage of 32-bit precision.
     else:
         logging.warning(
-            "Using 32bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
+            "Using 32-bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
         )
 
 

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev177'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev177')
+__version__ = version = '0.4.3.dev186'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev186')

jaxsim/exceptions.py

@@ -17,6 +17,10 @@ def raise_if(
             format string (fmt), whose fields are filled with the args and kwargs.
     """
 
+    # Disable host callback if running on TPU.
+    if jax.devices()[0].platform == "tpu":
+        return
+
     # Check early that the format string is well-formed.
     try:
         _ = msg.format(*args, **kwargs)

jaxsim/mujoco/__init__.py

@@ -1,3 +1,4 @@
 from .loaders import RodModelToMjcf, SdfToMjcf, UrdfToMjcf
 from .model import MujocoModelHelper
+from .utils import mujoco_data_from_jaxsim
 from .visualizer import MujocoVideoRecorder, MujocoVisualizer

jaxsim/mujoco/loaders.py

@@ -646,7 +646,7 @@ class MujocoCamera:
     def build(cls, **kwargs) -> MujocoCamera:
 
         if not all(isinstance(value, str) for value in kwargs.values()):
-            raise ValueError("Values must be strings")
+            raise ValueError(f"Values must be strings: {kwargs}")
 
         return cls(**kwargs)
 

jaxsim/mujoco/model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import functools
 import pathlib
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import mujoco as mj
@@ -107,7 +107,8 @@ class MujocoModelHelper:
             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))
+            # 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)
@@ -315,7 +316,7 @@ 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."""
 

jaxsim/mujoco/utils.py

@@ -0,0 +1,101 @@
+import mujoco as mj
+import numpy as np
+
+from . 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

jaxsim/mujoco/visualizer.py

@@ -89,7 +89,7 @@ class MujocoVideoRecorder:
         if not exist_ok and path.is_file():
             raise FileExistsError(f"The file '{path}' already exists.")
 
-        media.write_video(path=path, images=self.frames, fps=self.fps)
+        media.write_video(path=path, images=np.array(self.frames), fps=self.fps)
 
     @staticmethod
     def compute_down_sampling(original_fps: int, target_min_fps: int) -> int:

{jaxsim-0.4.3.dev177.dist-info → jaxsim-0.4.3.dev186.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev177
+Version: 0.4.3.dev186
 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.3.dev177.dist-info → jaxsim-0.4.3.dev186.dist-info}/RECORD

@@ -1,6 +1,6 @@
-jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=SzGoIDpeznpZHWyfdxXEtnO3y8zLaXJORHRJRSUcxsU,428
-jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
+jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
+jaxsim/_version.py,sha256=K0Qt3IiihQ28Vnsxow16UNJNyDGyw4M94790Je5aXw8,428
+jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
@@ -29,11 +29,12 @@ jaxsim/math/quaternion.py,sha256=_WA7W3iv7px83sWO1V1n0-J78hqAlO4SL1-jofE-UZ4,475
 jaxsim/math/rotation.py,sha256=k-nwT79zmWrys3NNAB-lGWxat7Kqm_6JnFRoimJ8rBg,2156
 jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=KXzQgOnCfAtbXCwxhplpJ3F0JT3oEyeLVby1_uRAryQ,2892
-jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
+jaxsim/mujoco/__init__.py,sha256=fZyRWre49pIhOrYdf6yJk_hOax8qWGe8OCmoq-dMVq8,201
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=qT7Le_L7z2prXKA7O9x5rkbbh-_lIrrmLXTjgoAjhZ4,25339
-jaxsim/mujoco/model.py,sha256=AQksXemXWACJ3yvefV2G5HLwwBU9ISoJrOD1wlxdY5w,16386
-jaxsim/mujoco/visualizer.py,sha256=T1vU-w4NKSmgEkZ0FqVcGmIvYrYO0len2UBSsU4MOZ0,6978
+jaxsim/mujoco/loaders.py,sha256=CkFGydgOku5P_Pz7wdWlM2SCJRs71ePF-vsY9i90-I0,25350
+jaxsim/mujoco/model.py,sha256=5_7rWk_WBkNKDHqeewIFj0t2ZGqJpE6RDXHSbRvw4e4,16493
+jaxsim/mujoco/utils.py,sha256=bGbLMSzcdqbinIwHHJHt8ZN1uup_6DLdB2dWqKiXwO4,3955
+jaxsim/mujoco/visualizer.py,sha256=nD6SNWmn-nxjjjIY9oPAHvL2j8q93DJDjZeepzke_DQ,6988
 jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/parsers/kinematic_graph.py,sha256=wT2bgaCS8VQJTHy2H9sENkVPDOiMkRikxEF1t_WaahQ,34748
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
@@ -64,8 +65,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev177.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev177.dist-info/METADATA,sha256=BMlT_szB4WbLIZmCucz750x_aRXoe2n7ycwlzI3l-sk,17276
-jaxsim-0.4.3.dev177.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev177.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev177.dist-info/RECORD,,
+jaxsim-0.4.3.dev186.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev186.dist-info/METADATA,sha256=YYb7FonjyOeyop0Ni-0dB0ijfk215c0-PZo6k9v6JAo,17276
+jaxsim-0.4.3.dev186.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev186.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev186.dist-info/RECORD,,