jaxsim 0.4.3.dev169__py3-none-any.whl → 0.4.3.dev181__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.dev169'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev169')
+__version__ = version = '0.4.3.dev181'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev181')

jaxsim/api/data.py

@@ -38,12 +38,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     contacts_params: jaxsim.rbda.contacts.ContactsParams = dataclasses.field(repr=False)
 
-    time_ns: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(
-            0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
-        ),
-    )
-
     def __hash__(self) -> int:
 
         from jaxsim.utils.wrappers import HashedNumpyArray
@@ -52,7 +46,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             (
                 hash(self.state),
                 HashedNumpyArray.hash_of_array(self.gravity),
-                HashedNumpyArray.hash_of_array(self.time_ns),
                 hash(self.contacts_params),
             )
         )
@@ -115,7 +108,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
         contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
-        time: jtp.FloatLike | None = None,
         extended_ode_state: dict[str, jtp.PyTree] | None = None,
     ) -> JaxSimModelData:
         """
@@ -134,7 +126,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             standard_gravity: The standard gravity constant.
             contacts_params: The parameters of the soft contacts.
             velocity_representation: The velocity representation to use.
-            time: The time at which the state is created.
             extended_ode_state:
                 Additional user-defined state variables that are not part of the
                 standard `ODEState` object. Useful to extend the system dynamics
@@ -196,11 +187,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ).squeeze()
         )
 
-        time_ns = jnp.array(
-            time * 1e9 if time is not None else 0.0,
-            dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
-        )
-
         W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
             translation=base_position, quaternion=base_quaternion
         )
@@ -246,7 +232,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 contacts_params = model.contact_model.parameters
 
         return JaxSimModelData(
-            time_ns=time_ns,
             state=ode_state,
             gravity=gravity,
             contacts_params=contacts_params,
@@ -257,16 +242,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     # Extract quantities
     # ==================
 
-    def time(self) -> jtp.Float:
-        """
-        Get the simulated time.
-
-        Returns:
-            The simulated time in seconds.
-        """
-
-        return self.time_ns.astype(float) / 1e9
-
     def standard_gravity(self) -> jtp.Float:
         """
         Get the standard gravity constant.

jaxsim/api/model.py

@@ -32,6 +32,10 @@ class JaxSimModel(JaxsimDataclass):
 
     model_name: Static[str]
 
+    time_step: jaxsim.integrators.TimeStep = dataclasses.field(
+        default_factory=lambda: jnp.array(0.001, dtype=float),
+    )
+
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
         default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
     )
@@ -64,6 +68,9 @@ class JaxSimModel(JaxsimDataclass):
         if self.model_name != other.model_name:
             return False
 
+        if self.time_step != other.time_step:
+            return False
+
         if self.kin_dyn_parameters != other.kin_dyn_parameters:
             return False
 
@@ -74,6 +81,7 @@ class JaxSimModel(JaxsimDataclass):
         return hash(
             (
                 hash(self.model_name),
+                hash(float(self.time_step)),
                 hash(self.kin_dyn_parameters),
                 hash(self.contact_model),
             )
@@ -88,6 +96,7 @@ class JaxSimModel(JaxsimDataclass):
         model_description: str | pathlib.Path | rod.Model,
         model_name: str | None = None,
         *,
+        time_step: jtp.FloatLike | None = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         is_urdf: bool | None = None,
@@ -102,6 +111,9 @@ class JaxSimModel(JaxsimDataclass):
                 its content, or a pre-parsed/pre-built rod model.
             model_name:
                 The name of the model. If not specified, it is read from the description.
+            time_step:
+                The default time step to consider for the simulation. It can be
+                manually overridden in the function that steps the simulation.
             terrain: The terrain to consider (the default is a flat infinite plane).
             contact_model:
                 The contact model to consider.
@@ -135,6 +147,7 @@ class JaxSimModel(JaxsimDataclass):
         model = JaxSimModel.build(
             model_description=intermediate_description,
             model_name=model_name,
+            time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
         )
@@ -150,6 +163,7 @@ class JaxSimModel(JaxsimDataclass):
         model_description: ModelDescription,
         model_name: str | None = None,
         *,
+        time_step: jtp.FloatLike | None = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
     ) -> JaxSimModel:
@@ -162,6 +176,9 @@ class JaxSimModel(JaxsimDataclass):
                 of the model.
             model_name:
                 The name of the model. If not specified, it is read from the description.
+            time_step:
+                The default time step to consider for the simulation. It can be
+                manually overridden in the function that steps the simulation.
             terrain: The terrain to consider (the default is a flat infinite plane).
             contact_model:
                 The contact model to consider.
@@ -179,6 +196,11 @@ class JaxSimModel(JaxsimDataclass):
             terrain or JaxSimModel.__dataclass_fields__["terrain"].default_factory()
         )
 
+        # Consider the default time step if not specified.
+        time_step = (
+            time_step or JaxSimModel.__dataclass_fields__["time_step"].default_factory()
+        )
+
         # Create the default contact model.
         # It will be populated with an initial estimation of good parameters.
         # While these might not be the best, they are a good starting point.
@@ -192,6 +214,7 @@ class JaxSimModel(JaxsimDataclass):
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
+            time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
             # The following is wrapped as hashless since it's a static argument, and we
@@ -1915,8 +1938,9 @@ def step(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    dt: jtp.FloatLike,
     integrator: jaxsim.integrators.Integrator,
+    t0: jtp.FloatLike = 0.0,
+    dt: jtp.FloatLike | None = None,
     integrator_state: dict[str, Any] | None = None,
     link_forces: jtp.MatrixLike | None = None,
     joint_force_references: jtp.VectorLike | None = None,
@@ -1928,9 +1952,10 @@ def step(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        dt: The time step to consider.
         integrator: The integrator to use.
         integrator_state: The state of the integrator.
+        t0: The initial time to consider. Only relevant for time-dependent dynamics.
+        dt: The time step to consider. If not specified, it is read from the model.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
@@ -1951,17 +1976,20 @@ def step(
 
     integrator_state = integrator_state if integrator_state is not None else dict()
 
-    # Extract the initial resources.
-    t0_ns = data.time_ns
+    # Initialize the time-related variables.
     state_t0 = data.state
-    integrator_state_x0 = integrator_state
+    t0 = jnp.array(t0, dtype=float)
+    dt = jnp.array(dt if dt is not None else model.time_step).astype(float)
+
+    # Rename the integrator state.
+    integrator_state_t0 = integrator_state
 
     # Step the dynamics forward.
     state_tf, integrator_state_tf = integrator.step(
         x0=state_t0,
-        t0=jnp.array(t0_ns / 1e9).astype(float),
+        t0=t0,
         dt=dt,
-        params=integrator_state_x0,
+        params=integrator_state_t0,
         # Always inject the current (model, data) pair into the system dynamics
         # considered by the integrator, and include the input variables represented
         # by the pair (joint_force_references, link_forces).
@@ -1980,24 +2008,8 @@ 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=tf_ns,
-        )
-    )
+    # Store the new state of the model.
+    data_tf = data.replace(state=state_tf)
 
     # Post process the simulation state, if needed.
     match model.contact_model:
@@ -2064,7 +2076,4 @@ def step(
                 velocity_representation=data.velocity_representation, validate=False
             )
 
-    return (
-        data_tf,
-        integrator_state_tf,
-    )
+    return data_tf, integrator_state_tf

jaxsim/api/ode.py

@@ -63,7 +63,6 @@ def wrap_system_dynamics_for_integration(
         # Update the state and time stored inside data.
         with data_f.editable(validate=True) as data_rw:
             data_rw.state = x
-            data_rw.time_ns = jnp.array(t * 1e9).astype(data_rw.time_ns.dtype)
 
         # Evaluate the system dynamics, allowing to override the kwargs originally
         # passed when the closure was created.

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/rbda/contacts/visco_elastic.py

@@ -263,7 +263,7 @@ class ViscoElasticContacts(common.ContactModel):
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         *,
-        dt: jtp.FloatLike,
+        dt: jtp.FloatLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
         joint_force_references: jtp.VectorLike | None = None,
     ) -> tuple[jtp.Vector, tuple[Any, ...]]:
@@ -273,7 +273,7 @@ class ViscoElasticContacts(common.ContactModel):
         Args:
             model: The robot model considered by the contact model.
             data: The data of the considered model.
-            dt: The integration time step.
+            dt: The time step to consider. If not specified, it is read from the model.
             link_forces:
                 The 6D forces to apply to the links expressed in the frame corresponding
                 to the velocity representation of `data`.
@@ -305,13 +305,16 @@ class ViscoElasticContacts(common.ContactModel):
             model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
         )
 
+        # Initialize the time step.
+        dt = dt if dt is not None else model.time_step
+
         # Compute the average contact linear forces in mixed representation by
         # integrating the contact dynamics in the continuous time domain.
         CW_f̅l, CW_fl̿, m_tf = (
             ViscoElasticContacts._compute_contact_forces_with_exponential_integration(
                 model=model,
                 data=data,
-                dt=dt,
+                dt=jnp.array(dt).astype(float),
                 joint_force_references=joint_force_references,
                 link_forces=link_forces,
                 indices_of_enabled_collidable_points=indices_of_enabled_collidable_points,
@@ -923,14 +926,10 @@ class ViscoElasticContacts(common.ContactModel):
         )
 
         # Create the data at the final time.
-        with data.editable(validate=True) as data_tf:
-            data_tf: js.data.JaxSimModelData
-            data_tf.time_ns = data.time_ns + (dt * 1e9).astype(data.time_ns.dtype)
-
+        data_tf = data.copy()
         data_tf = data_tf.reset_joint_positions(q_plus[7:])
         data_tf = data_tf.reset_base_position(q_plus[0:3])
         data_tf = data_tf.reset_base_quaternion(q_plus[3:7])
-
         data_tf = data_tf.reset_joint_velocities(W_ν_plus[6:])
         data_tf = data_tf.reset_base_velocity(
             W_ν_plus[0:6], velocity_representation=jaxsim.VelRepr.Inertial
@@ -946,7 +945,7 @@ def step(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    dt: jtp.FloatLike,
+    dt: jtp.FloatLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
     joint_force_references: jtp.VectorLike | None = None,
 ) -> tuple[js.data.JaxSimModelData, dict[str, Any]]:
@@ -956,7 +955,7 @@ def step(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        dt: The time step to consider.
+        dt: The time step to consider. If not specified, it is read from the model.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
@@ -970,11 +969,14 @@ def step(
     assert isinstance(model.contact_model, ViscoElasticContacts)
     assert isinstance(data.contacts_params, ViscoElasticContactsParams)
 
+    # Initialize the time step.
+    dt = dt if dt is not None else model.time_step
+
     # Compute the contact forces with the exponential integrator.
     W_f̅_C, (W_f̿_C, m_tf) = model.contact_model.compute_contact_forces(
         model=model,
         data=data,
-        dt=dt,
+        dt=jnp.array(dt).astype(float),
         link_forces=link_forces,
         joint_force_references=joint_force_references,
     )

{jaxsim-0.4.3.dev169.dist-info → jaxsim-0.4.3.dev181.dist-info}/METADATA

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

@@ -1,19 +1,19 @@
-jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=gkpl6zmn3SM0X-pLcufsS6MQ7L76BZH7LaWidJd7jfY,428
-jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
+jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
+jaxsim/_version.py,sha256=2YMpT461ObsI3rceAHrUVR-OPYPsQ43SHkzRH20mlDY,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
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
 jaxsim/api/contact.py,sha256=2qBIStXWxIJTrW3Eyx6UPQcTXDXalOyasgD7Pk1_v1E,24486
-jaxsim/api/data.py,sha256=QumLL91DDdGtqgjTex4d4GAk23sMiV2gHztETwkR9Xw,29695
+jaxsim/api/data.py,sha256=sfu_uJtYRQIf_sg9IWzR95McRdZgtHwArAuzF6KD-1A,28939
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=thJbz9XhpXgom23S6MXX2ugxGoAD-k947ZMAHDisy2w,29620
 jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
-jaxsim/api/model.py,sha256=IcJhFdVjswSncurVpQXoa7uXVqSsvAuZVs-DXDsXi08,69021
-jaxsim/api/ode.py,sha256=aXcw1hT8cQlpUK9tmCf6JBngtvL7IcEbLGpw2qRMilI,14290
+jaxsim/api/model.py,sha256=-Au3Xdm3TJLyKN_r06pr9G99zmzjNhDI4KZz4xox7iE,69783
+jaxsim/api/ode.py,sha256=J_WuaoPl3ZY-yvTrCQun-rQoIAv_duynSXAGxqx93sg,14211
 jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
 jaxsim/api/references.py,sha256=fW77LitZ8DYgT6ZmUInJfm5luBV1mTcqcNRiC_i79og,20862
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
@@ -57,15 +57,15 @@ jaxsim/rbda/contacts/common.py,sha256=_yrxTM16Je9ck5aM95ndk8Kwu_oijxG9Jaf1jEjHEY
 jaxsim/rbda/contacts/relaxed_rigid.py,sha256=Ob5LdKe3D7tGlIdT4LamJ6_F0j5pzUmWNYoWqy8Di98,17169
 jaxsim/rbda/contacts/rigid.py,sha256=1TTiGXSOipO8l5FDTtxqRNo1ArCNtDg-Yr3olPgBLGs,17588
 jaxsim/rbda/contacts/soft.py,sha256=TMCUDtFmNIae04LCla57iXMjdt9F5qTFjYEnP5NdLFg,16809
-jaxsim/rbda/contacts/visco_elastic.py,sha256=Uj0bsrO5VQs9XCExREW18u853Ndj5FUP7YSOwwNEtzg,39590
+jaxsim/rbda/contacts/visco_elastic.py,sha256=tfcH4_JFox6_6PyR29kLlqc8pYN8WCslYnClxV7TnSU,39780
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=K91HEzPqTSyNrc_j1KfAAEF_5oDeuk_-jnnZGrcMEcY,5015
 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.dev169.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev169.dist-info/METADATA,sha256=QfLskOn1M2_p0j-wYEbrfZqJGDdiWEVzPmgMiFgD7CM,17276
-jaxsim-0.4.3.dev169.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev169.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev169.dist-info/RECORD,,
+jaxsim-0.4.3.dev181.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev181.dist-info/METADATA,sha256=w-b474j6ugFST4V5QJ2h9DTtrb2xRH0cdCRAMyBG8wg,17276
+jaxsim-0.4.3.dev181.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev181.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev181.dist-info/RECORD,,