jaxsim 0.4.3.dev18__py3-none-any.whl → 0.4.3.dev28__py3-none-any.whl

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/_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/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/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.dist-info → jaxsim-0.4.3.dev28.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
-jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
-jaxsim/_version.py,sha256=SFJGfO84uy3oOc6jDmWWev6VuJIqHL3tI8_OvaYfdsA,426
+jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
+jaxsim/_version.py,sha256=QnmczpVtiPFU-Ox0iyWHC4BmRC1wb9rfmHqwmKCR7W4,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
@@ -7,12 +7,12 @@ jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
 jaxsim/api/contact.py,sha256=HyEAjF7BySDDOlRahN0l7V15IPB0HPXuoM0twamuEW0,20913
-jaxsim/api/data.py,sha256=CUh9lvhVk3_clNQ26BUBGpjvFSsK_PrVWVMEWpMdHRM,27206
+jaxsim/api/data.py,sha256=p44Q7pNQEWQz2KgwNrk47_m0D-vbVtT8xuy4fLD8T58,27465
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=L81bQe-noPT6_54KOSF7KBjRmEPAS433ULn2EcXI8vI,5115
 jaxsim/api/kin_dyn_parameters.py,sha256=CcfSg5Mc8qb1mZeMQ4AK_ffZIsK5yOl7tu397pFhcDA,29369
 jaxsim/api/link.py,sha256=qPRtc8qqMRjZxUCZYXJMygbB6huDXBfIT1b1b8Durkw,18631
-jaxsim/api/model.py,sha256=HXoqCtQ3KStGoxhgvFm8P_Sc-lbEM4l5No2MoHzNlOk,65558
+jaxsim/api/model.py,sha256=Tq0CBjmjipGT02q4LgvVv2PH2uEFFHrDr_OiaiQNIXQ,65872
 jaxsim/api/ode.py,sha256=Vb2sN4zwpXnaJDD9-ziz2qvfmfa4jvIQ0fONbBIRGmU,13368
 jaxsim/api/ode_data.py,sha256=U7F6TL6bENAxpQQl4PupPoDG7d7VfTTFqDAs3xwu6Hs,20003
 jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
@@ -62,8 +62,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=5xJbY0G8d7C0OTNIW9T4vQxiDak6TGZT9gpNOvRykFI,11373
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=JhLUh1g8iU-lhjbuZRfkscPZhYlLCOorVM2Xl3ulRBI,4054
-jaxsim-0.4.3.dev18.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev18.dist-info/METADATA,sha256=aLpRkfa9CC7GVzXMKX3LY5DkCHEmOr4CE-u3Vbt5fx8,17247
-jaxsim-0.4.3.dev18.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-jaxsim-0.4.3.dev18.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev18.dist-info/RECORD,,
+jaxsim-0.4.3.dev28.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev28.dist-info/METADATA,sha256=EktoHZgQXhOsv1CbRnKOcrXmncWEiEgzvmzQVlkYB0M,17247
+jaxsim-0.4.3.dev28.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+jaxsim-0.4.3.dev28.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev28.dist-info/RECORD,,