jaxsim 0.4.3.dev80__tar.gz → 0.4.3.dev90__tar.gz

{jaxsim-0.4.3.dev80 → jaxsim-0.4.3.dev90}/PKG-INFO

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

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

{jaxsim-0.4.3.dev80 → jaxsim-0.4.3.dev90}/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.dev80 → jaxsim-0.4.3.dev90}/src/jaxsim/api/kin_dyn_parameters.py

@@ -422,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.dev80 → jaxsim-0.4.3.dev90}/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.dev80 → jaxsim-0.4.3.dev90}/src/jaxsim/rbda/crba.py

@@ -94,48 +94,48 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
         j = i
 
-        CarryInnerFn = tuple[jtp.Int, jtp.Matrix, jtp.Matrix]
-        carry_inner_fn = (j, Fi, M)
+        FakeWhileCarry = tuple[jtp.Int, jtp.Vector, jtp.Matrix]
+        fake_while_carry = (j, Fi, M)
 
-        def while_loop_body(carry: CarryInnerFn) -> CarryInnerFn:
-            j, Fi, M = carry
+        # This internal for loop implements the while loop of the CRBA algorithm
+        # to compute off-diagonal blocks of the mass matrix M.
+        # In pseudocode it is implemented as a while loop. However, in order to enable
+        # applying reverse-mode AD, we implement it as a nested for loop with a fixed
+        # number of iterations and a branching model to skip for loop iterations.
+        def fake_while_loop(
+            carry: FakeWhileCarry, i: jtp.Int
+        ) -> tuple[FakeWhileCarry, None]:
 
-            Fi = i_X_λi[j].T @ Fi
-            j = λ[j]
-            jj = j - 1
+            def compute(carry: FakeWhileCarry) -> FakeWhileCarry:
 
-            M_ij = Fi.T @ S[j]
+                j, Fi, M = carry
 
-            M = M.at[ii + 6, jj + 6].set(M_ij.squeeze())
-            M = M.at[jj + 6, ii + 6].set(M_ij.squeeze())
+                Fi = i_X_λi[j].T @ Fi
+                j = λ[j]
 
-            return j, Fi, M
+                M_ij = Fi.T @ S[j]
 
-        # The following functions are part of a (rather messy) workaround for computing
-        # a while loop using a for loop with fixed number of iterations.
-        def inner_fn(carry: CarryInnerFn, k: jtp.Int) -> tuple[CarryInnerFn, None]:
-            def compute_inner(carry: CarryInnerFn) -> tuple[CarryInnerFn, None]:
-                j, _, _ = carry
-                out = jax.lax.cond(
-                    pred=(λ[j] > 0),
-                    true_fun=while_loop_body,
-                    false_fun=lambda carry: carry,
-                    operand=carry,
-                )
-                return out, None
+                jj = j - 1
+                M = M.at[ii + 6, jj + 6].set(M_ij.squeeze())
+                M = M.at[jj + 6, ii + 6].set(M_ij.squeeze())
+
+                return j, Fi, M
 
             j, _, _ = carry
-            return jax.lax.cond(
-                pred=(k == j),
-                true_fun=compute_inner,
-                false_fun=lambda carry: (carry, None),
+
+            j, Fi, M = jax.lax.cond(
+                pred=jnp.logical_and(i == λ[j], λ[j] > 0),
+                true_fun=compute,
+                false_fun=lambda carry: carry,
                 operand=carry,
             )
 
+            return (j, Fi, M), None
+
         (j, Fi, M), _ = (
             jax.lax.scan(
-                f=inner_fn,
-                init=carry_inner_fn,
+                f=fake_while_loop,
+                init=fake_while_carry,
                 xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
             )
             if model.number_of_links() > 1

{jaxsim-0.4.3.dev80 → jaxsim-0.4.3.dev90}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev80
+Version: 0.4.3.dev90
 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.dev80 → jaxsim-0.4.3.dev90}/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.dev80 → jaxsim-0.4.3.dev90}/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