jaxsim 0.4.3.dev18__tar.gz → 0.4.3.dev28__tar.gz

{jaxsim-0.4.3.dev18 → jaxsim-0.4.3.dev28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev18
+Version: 0.4.3.dev28
 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.dev18 → jaxsim-0.4.3.dev28}/src/jaxsim/__init__.py

@@ -20,6 +20,11 @@ def _jnp_options() -> None:
         if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
             logging.warning("Failed to enable 64bit precision in JAX")
 
+    else:
+        logging.warning(
+            "Using 32bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
+        )
+
 
 def _np_options() -> None:
     import numpy as np

{jaxsim-0.4.3.dev18 → jaxsim-0.4.3.dev28}/src/jaxsim/_version.py

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

{jaxsim-0.4.3.dev18 → jaxsim-0.4.3.dev28}/src/jaxsim/api/data.py

@@ -39,7 +39,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     contacts_params: jaxsim.rbda.ContactsParams = dataclasses.field(repr=False)
 
     time_ns: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(0, dtype=jnp.uint64)
+        default_factory=lambda: jnp.array(
+            0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
+        ),
     )
 
     def __hash__(self) -> int:
@@ -172,9 +174,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         )
 
         time_ns = (
-            jnp.array(time * 1e9, dtype=jnp.uint64)
+            jnp.array(
+                time * 1e9,
+                dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
+            )
             if time is not None
-            else jnp.array(0, dtype=jnp.uint64)
+            else jnp.array(
+                0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
+            )
         )
 
         if isinstance(model.contact_model, SoftContacts):

{jaxsim-0.4.3.dev18 → jaxsim-0.4.3.dev28}/src/jaxsim/api/model.py

@@ -1931,11 +1931,22 @@ def step(
         ),
     )
 
+    tf_ns = t0_ns + jnp.array(dt * 1e9, dtype=t0_ns.dtype)
+    tf_ns = jnp.where(tf_ns >= t0_ns, tf_ns, jnp.array(0, dtype=t0_ns.dtype))
+
+    jax.lax.cond(
+        pred=tf_ns >= t0_ns,
+        true_fun=lambda: jax.debug.print(
+            "The simulation time overflowed, resetting simulation time to 0."
+        ),
+        false_fun=lambda: None,
+    )
+
     data_tf = (
         # Store the new state of the model and the new time.
         data.replace(
             state=state_tf,
-            time_ns=t0_ns + jnp.array(dt * 1e9).astype(jnp.uint64),
+            time_ns=tf_ns,
         )
     )
 

{jaxsim-0.4.3.dev18 → jaxsim-0.4.3.dev28}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev18
+Version: 0.4.3.dev28
 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.dev18 → jaxsim-0.4.3.dev28}/tests/conftest.py

@@ -386,3 +386,28 @@ def jaxsim_models_fixed_base(request) -> pathlib.Path | str:
 
     model_name: str = request.param
     return get_jaxsim_model_fixture(model_name=model_name, request=request)
+
+
+@pytest.fixture(scope="function")
+def set_jax_32bit(monkeypatch):
+    """
+    Fixture that temporarily sets JAX precision to 32-bit for the duration of the test.
+    """
+
+    del globals()["jaxsim"]
+    del globals()["js"]
+
+    # Temporarily disable x64
+    monkeypatch.setenv("JAX_ENABLE_X64", "0")
+
+
+@pytest.fixture(scope="function")
+def jaxsim_model_box_32bit(set_jax_32bit, request) -> js.model.JaxSimModel:
+    """
+    Fixture providing the JaxSim model of a box with 32-bit precision.
+
+    Returns:
+        The JaxSim model of a box with 32-bit precision.
+    """
+
+    return get_jaxsim_model_fixture(model_name="box", request=request)