jaxsim 0.4.3.dev77__tar.gz → 0.4.3.dev88__tar.gz

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev77
+Version: 0.4.3.dev88
 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>
@@ -60,9 +60,9 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
-Requires-Dist: jax>=0.4.13
+Requires-Dist: jax>=0.4.26
 Requires-Dist: jaxopt>=0.8.0
-Requires-Dist: jaxlib>=0.4.13
+Requires-Dist: jaxlib>=0.4.26
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax_dataclasses>=1.4.0
 Requires-Dist: pptree

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/environment.yml

@@ -7,9 +7,9 @@ dependencies:
   # ===========================
   - python >= 3.12.0
   - coloredlogs
-  - jax >= 0.4.13
+  - jax >= 0.4.26
   - jaxopt >= 0.8.0
-  - jaxlib >= 0.4.13
+  - jaxlib >= 0.4.26
   - jaxlie >= 1.3.0
   - jax-dataclasses >= 1.4.0
   - pptree

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/pyproject.toml

@@ -44,9 +44,9 @@ classifiers = [
 ]
 dependencies = [
     "coloredlogs",
-    "jax >= 0.4.13",
+    "jax >= 0.4.26",
     "jaxopt >= 0.8.0",
-    "jaxlib >= 0.4.13",
+    "jaxlib >= 0.4.26",
     "jaxlie >= 1.3.0",
     "jax_dataclasses >= 1.4.0",
     "pptree",

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/_version.py

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

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/api/contact.py

@@ -372,11 +372,9 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     """
 
     # Get the transforms of the parent link of all collidable points.
-    W_H_L = jax.vmap(
-        lambda parent_link_idx: js.link.transform(
-            model=model, data=data, link_index=parent_link_idx
-        )
-    )(jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int))
+    W_H_L = js.model.forward_kinematics(model=model, data=data)[
+        jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
+    ]
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
@@ -427,9 +425,9 @@ def jacobian(
     # Compute the contact Jacobian.
     # In inertial-fixed output representation, the Jacobian of the parent link is also
     # the Jacobian of the frame C implicitly associated with the collidable point.
-    W_J_WC = jax.vmap(lambda parent_link_idx: W_J_WL[parent_link_idx])(
+    W_J_WC = W_J_WL[
         jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
-    )
+    ]
 
     # Adjust the output representation.
     match output_vel_repr:

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/api/kin_dyn_parameters.py

@@ -98,9 +98,7 @@ class KynDynParameters(JaxsimDataclass):
         ]
 
         # Create a vectorized object of link parameters.
-        link_parameters = jax.tree_util.tree_map(
-            lambda *l: jnp.stack(l), *link_parameters_list
-        )
+        link_parameters = jax.tree.map(lambda *l: jnp.stack(l), *link_parameters_list)
 
         # =================
         # Joints properties
@@ -114,7 +112,7 @@ class KynDynParameters(JaxsimDataclass):
 
         # Create a vectorized object of joint parameters.
         joint_parameters = (
-            jax.tree_util.tree_map(lambda *l: jnp.stack(l), *joint_parameters_list)
+            jax.tree.map(lambda *l: jnp.stack(l), *joint_parameters_list)
             if len(ordered_joints) > 0
             else JointParameters(
                 index=jnp.array([], dtype=int),
@@ -424,9 +422,7 @@ class KynDynParameters(JaxsimDataclass):
         # Note that here we include also the index 0 since suc_H_child[0] stores the
         # optional pose of the base link w.r.t. the root frame of the model.
         # This is supported by SDF when the base link <pose> element is defined.
-        suc_H_i = jax.vmap(lambda i: self.joint_model.successor_H_child(joint_index=i))(
-            jnp.arange(0, 1 + self.number_of_joints())
-        )
+        suc_H_i = self.joint_model.suc_H_i[jnp.arange(0, 1 + self.number_of_joints())]
 
         # Compute the overall transforms from the parent to the child of each joint by
         # composing all the components of our joint model.

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/api/link.py

@@ -154,7 +154,7 @@ def spatial_inertia(
         idx=link_index,
     )
 
-    link_parameters = jax.tree_util.tree_map(
+    link_parameters = jax.tree.map(
         lambda l: l[link_index], model.kin_dyn_parameters.link_parameters
     )
 

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/api/model.py

@@ -395,13 +395,7 @@ def total_mass(model: JaxSimModel) -> jtp.Float:
         The total mass of the model.
     """
 
-    return (
-        jax.vmap(lambda idx: js.link.mass(model=model, link_index=idx))(
-            jnp.arange(model.number_of_links())
-        )
-        .sum()
-        .astype(float)
-    )
+    return model.kin_dyn_parameters.link_parameters.mass.sum().astype(float)
 
 
 @jax.jit
@@ -974,7 +968,7 @@ def free_floating_coriolis_matrix(
             lambda link_index: js.link.jacobian_derivative(
                 model=model, data=data, link_index=link_index
             )
-        )(js.link.names_to_idxs(model=model, link_names=model.link_names()))
+        )(jnp.arange(model.number_of_links()))
 
     L_M_L = link_spatial_inertia_matrices(model=model)
 

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/integrators/common.py

@@ -173,9 +173,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
         # Make sure that all leafs of the dictionary are JAX arrays.
         # Also, since these are dummy parameters, set them all to zero.
-        params_after_init = jax.tree_util.tree_map(
-            lambda l: jnp.zeros_like(l), integrator.params
-        )
+        params_after_init = jax.tree.map(lambda l: jnp.zeros_like(l), integrator.params)
 
         # Mark the next step as first step after initialization.
         params_after_init = params_after_init | {
@@ -290,7 +288,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         z, aux_dict = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
 
         # The next state is the batch element located at the configured index of solution.
-        next_state = jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
+        next_state = jax.tree.map(lambda l: l[self.row_index_of_solution], z)
 
         return next_state, aux_dict
 
@@ -327,7 +325,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         """
 
         op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf
-        return jax.tree_util.tree_map(op, x0, k)
+        return jax.tree.map(op, x0, k)
 
     @classmethod
     def post_process_state(
@@ -374,7 +372,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
         # Initialize the carry of the for loop with the stacked kᵢ vectors.
-        carry0 = jax.tree_util.tree_map(
+        carry0 = jax.tree.map(
             lambda l: jnp.repeat(jnp.zeros_like(l)[jnp.newaxis, ...], c.size, axis=0),
             x0,
         )
@@ -398,7 +396,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
                 # Compute ∑ⱼ aᵢⱼ kⱼ.
                 op_sum_ak = lambda k: jnp.einsum("s,s...->...", A[i], k)
-                sum_ak = jax.tree_util.tree_map(op_sum_ak, K)
+                sum_ak = jax.tree.map(op_sum_ak, K)
 
                 # Compute the next state for the kᵢ evaluation.
                 # Note that this is not a Δt integration since aᵢⱼ could be fractional.
@@ -419,7 +417,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
             # Store the kᵢ derivative in K.
             op = lambda l_k, l_ki: l_k.at[i].set(l_ki)
-            K = jax.tree_util.tree_map(op, K, ki)
+            K = jax.tree.map(op, K, ki)
 
             carry = K
             return carry, aux_dict
@@ -433,14 +431,12 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree_util.tree_map(
-                lambda l: l[self.index_of_fsal], K
-            )
+            self.params["dxdt0"] = jax.tree.map(lambda l: l[self.index_of_fsal], K)
 
         # Compute the output state.
         # Note that z contains as many new states as the rows of `b.T`.
         op = lambda x0, k: x0 + Δt * jnp.einsum("zs,s...->z...", b.T, k)
-        z = jax.tree_util.tree_map(op, x0, K)
+        z = jax.tree.map(op, x0, K)
 
         # Transform the final state of the integration.
         # This allows to inject custom logic, if needed.

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/integrators/variable_step.py

@@ -87,13 +87,13 @@ def estimate_step_size(
 
     # Compute the scaling factors of the initial state and its derivative.
     compute_scale = lambda x: atol + jnp.abs(x) * rtol
-    scale0 = jax.tree_util.tree_map(compute_scale, x0)
-    scale1 = jax.tree_util.tree_map(compute_scale, ẋ0)
+    scale0 = jax.tree.map(compute_scale, x0)
+    scale1 = jax.tree.map(compute_scale, ẋ0)
 
     # Scale the initial state and its derivative.
     scale_pytree = lambda x, scale: jnp.abs(x) / scale
-    x0_scaled = jax.tree_util.tree_map(scale_pytree, x0, scale0)
-    ẋ0_scaled = jax.tree_util.tree_map(scale_pytree, ẋ0, scale1)
+    x0_scaled = jax.tree.map(scale_pytree, x0, scale0)
+    ẋ0_scaled = jax.tree.map(scale_pytree, ẋ0, scale1)
 
     # Get the maximum of the scaled pytrees.
     d0 = jnp.linalg.norm(flatten(x0_scaled), ord=jnp.inf)
@@ -103,16 +103,16 @@ def estimate_step_size(
     h0 = jnp.where(jnp.minimum(d0, d1) <= 1e-5, 1e-6, 0.01 * d0 / d1)
 
     # Compute the next state (explicit Euler step) and its derivative.
-    x1 = jax.tree_util.tree_map(lambda x0, ẋ0: x0 + h0 * ẋ0, x0, ẋ0)
+    x1 = jax.tree.map(lambda x0, ẋ0: x0 + h0 * ẋ0, x0, ẋ0)
     ẋ1 = f(x1, t0 + h0)[0]
 
     # Compute the scaling factor of the state derivatives.
     compute_scale_2 = lambda ẋ0, ẋ1: atol + jnp.maximum(jnp.abs(ẋ0), jnp.abs(ẋ1)) * rtol
-    scale2 = jax.tree_util.tree_map(compute_scale_2, ẋ0, ẋ1)
+    scale2 = jax.tree.map(compute_scale_2, ẋ0, ẋ1)
 
     # Scale the difference of the state derivatives.
     scale_ẋ_difference = lambda ẋ0, ẋ1, scale: jnp.abs((ẋ0 - ẋ1) / scale)
-    ẋ_difference_scaled = jax.tree_util.tree_map(scale_ẋ_difference, ẋ0, ẋ1, scale2)
+    ẋ_difference_scaled = jax.tree.map(scale_ẋ_difference, ẋ0, ẋ1, scale2)
 
     # Get the maximum of the scaled derivatives difference.
     d2 = jnp.linalg.norm(flatten(ẋ_difference_scaled), ord=jnp.inf) / h0
@@ -151,11 +151,11 @@ def compute_pytree_scale(
     """
 
     # Consider a zero second pytree, if not given.
-    x2 = jax.tree_util.tree_map(lambda l: jnp.zeros_like(l), x1) if x2 is None else x2
+    x2 = jax.tree.map(lambda l: jnp.zeros_like(l), x1) if x2 is None else x2
 
     # Compute the scaling factors of the initial state and its derivative.
     compute_scale = lambda l1, l2: atol + jnp.maximum(jnp.abs(l1), jnp.abs(l2)) * rtol
-    scale = jax.tree_util.tree_map(compute_scale, x1, x2)
+    scale = jax.tree.map(compute_scale, x1, x2)
 
     return scale
 
@@ -198,14 +198,14 @@ def local_error_estimation(
 
     # Consider a zero estimated final state, if not given.
     xf_estimate = (
-        jax.tree_util.tree_map(lambda l: jnp.zeros_like(l), xf)
+        jax.tree.map(lambda l: jnp.zeros_like(l), xf)
         if xf_estimate is None
         else xf_estimate
     )
 
     # Estimate the error.
     estimate_error = lambda l, l̂, sc: jnp.abs(l - l̂) / sc
-    error_estimate = jax.tree_util.tree_map(estimate_error, xf, xf_estimate, scale)
+    error_estimate = jax.tree.map(estimate_error, xf, xf_estimate, scale)
 
     # Return the highest element of the error estimate.
     return jnp.linalg.norm(flatten(error_estimate), ord=norm_ord)
@@ -359,10 +359,8 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                 params_next = integrator.params
 
             # Extract the high-order solution xf and the low-order estimate x̂f.
-            xf = jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
-            x̂f = jax.tree_util.tree_map(
-                lambda l: l[self.row_index_of_solution_estimate], z
-            )
+            xf = jax.tree.map(lambda l: l[self.row_index_of_solution], z)
+            x̂f = jax.tree.map(lambda l: l[self.row_index_of_solution_estimate], z)
 
             # Calculate the local integration error.
             local_error = local_error_estimation(

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -230,7 +230,7 @@ class RelaxedRigidContacts(ContactModel):
         )
 
         def _detect_contact(x: jtp.Array, y: jtp.Array, z: jtp.Array) -> jtp.Array:
-            x, y, z = jax.tree_map(jnp.squeeze, (x, y, z))
+            x, y, z = jax.tree.map(jnp.squeeze, (x, y, z))
 
             n̂ = self.terrain.normal(x=x, y=y).squeeze()
             h = jnp.array([0, 0, z - model.terrain.height(x=x, y=y)])

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim/utils/jaxsim_dataclass.py

@@ -298,7 +298,7 @@ class JaxsimDataclass(abc.ABC):
         """
 
         # Make a copy calling tree_map.
-        obj = jax.tree_util.tree_map(lambda leaf: leaf, self)
+        obj = jax.tree.map(lambda leaf: leaf, self)
 
         # Make sure that the copied object and all the copied leaves have the same
         # mutability of the original object.

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev77
+Version: 0.4.3.dev88
 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>
@@ -60,9 +60,9 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
-Requires-Dist: jax>=0.4.13
+Requires-Dist: jax>=0.4.26
 Requires-Dist: jaxopt>=0.8.0
-Requires-Dist: jaxlib>=0.4.13
+Requires-Dist: jaxlib>=0.4.26
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax_dataclasses>=1.4.0
 Requires-Dist: pptree

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/src/jaxsim.egg-info/requires.txt

@@ -1,7 +1,7 @@
 coloredlogs
-jax>=0.4.13
+jax>=0.4.26
 jaxopt>=0.8.0
-jaxlib>=0.4.13
+jaxlib>=0.4.26
 jaxlie>=1.3.0
 jax_dataclasses>=1.4.0
 pptree

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/tests/test_api_link.py

@@ -74,7 +74,7 @@ def test_link_inertial_properties(
 
     for link_name, link_idx in zip(
         model.link_names(),
-        js.link.names_to_idxs(model=model, link_names=model.link_names()),
+        jnp.arange(model.number_of_links()),
         strict=True,
     ):
         if link_name == model.base_link():
@@ -164,7 +164,7 @@ def test_link_jacobians(
 
     for link_name, link_idx in zip(
         model.link_names(),
-        js.link.names_to_idxs(model=model, link_names=model.link_names()),
+        jnp.arange(model.number_of_links()),
         strict=True,
     ):
         v_WL_idt = kin_dyn.frame_velocity(frame_name=link_name)
@@ -185,7 +185,7 @@ def test_link_jacobians(
 
         for link_name, link_idx in zip(
             model.link_names(),
-            js.link.names_to_idxs(model=model, link_names=model.link_names()),
+            jnp.arange(model.number_of_links()),
             strict=True,
         ):
             v_WL_idt = kin_dyn_other_repr.frame_velocity(frame_name=link_name)
@@ -220,7 +220,7 @@ def test_link_bias_acceleration(
 
     for name, index in zip(
         model.link_names(),
-        js.link.names_to_idxs(model=model, link_names=model.link_names()),
+        jnp.arange(model.number_of_links()),
         strict=True,
     ):
         Jν_idt = kin_dyn.frame_bias_acc(frame_name=name)
@@ -240,11 +240,7 @@ def test_link_bias_acceleration(
 
             W_H_L = js.model.forward_kinematics(model=model, data=data)
 
-            W_a_bias_WL = jax.vmap(
-                lambda index: js.link.bias_acceleration(
-                    model=model, data=data, link_index=index
-                )
-            )(jnp.arange(model.number_of_links()))
+            W_a_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
 
             with data.switch_velocity_representation(VelRepr.Body):
 
@@ -252,11 +248,7 @@ def test_link_bias_acceleration(
                     lambda W_H_L: jaxsim.math.Adjoint.from_transform(transform=W_H_L)
                 )(W_H_L)
 
-                L_a_bias_WL = jax.vmap(
-                    lambda index: js.link.bias_acceleration(
-                        model=model, data=data, link_index=index
-                    )
-                )(jnp.arange(model.number_of_links()))
+                L_a_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
 
                 W_a_bias_WL_converted = jax.vmap(
                     lambda W_X_L, L_a_bias_WL: W_X_L @ L_a_bias_WL
@@ -269,11 +261,7 @@ def test_link_bias_acceleration(
 
             W_H_L = js.model.forward_kinematics(model=model, data=data)
 
-            L_a_bias_WL = jax.vmap(
-                lambda index: js.link.bias_acceleration(
-                    model=model, data=data, link_index=index
-                )
-            )(jnp.arange(model.number_of_links()))
+            L_a_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
 
             with data.switch_velocity_representation(VelRepr.Inertial):
 
@@ -283,11 +271,7 @@ def test_link_bias_acceleration(
                     )
                 )(W_H_L)
 
-                W_a_bias_WL = jax.vmap(
-                    lambda index: js.link.bias_acceleration(
-                        model=model, data=data, link_index=index
-                    )
-                )(jnp.arange(model.number_of_links()))
+                W_a_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
 
                 L_a_bias_WL_converted = jax.vmap(
                     lambda L_X_W, W_a_bias_WL: L_X_W @ W_a_bias_WL
@@ -323,14 +307,10 @@ def test_link_jacobian_derivative(
         lambda link_index: js.link.jacobian_derivative(
             model=model, data=data, link_index=link_index
         )
-    )(js.link.names_to_idxs(model=model, link_names=model.link_names()))
+    )(jnp.arange(model.number_of_links()))
 
     # Compute the product J̇ν.
-    O_a_bias_WL = jax.vmap(
-        lambda link_index: js.link.bias_acceleration(
-            model=model, data=data, link_index=link_index
-        )
-    )(js.link.names_to_idxs(model=model, link_names=model.link_names()))
+    O_a_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
 
     # Compare the two computations.
     assert jnp.einsum("l6g,g->l6", O_J̇_WL_I, I_ν) == pytest.approx(

{jaxsim-0.4.3.dev77 → jaxsim-0.4.3.dev88}/tests/test_automatic_differentiation.py

@@ -263,7 +263,7 @@ def test_ad_jacobian(
     # ====
 
     # Get the link indices.
-    link_indices = js.link.names_to_idxs(model=model, link_names=model.link_names())
+    link_indices = jnp.arange(model.number_of_links())
 
     # Get a closure exposing only the parameters to be differentiated.
     # We differentiate the jacobian of the last link, likely among those