jaxsim 0.4.2.dev45__tar.gz → 0.4.2.dev50__tar.gz

{jaxsim-0.4.2.dev45 → jaxsim-0.4.2.dev50}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2.dev45
+Version: 0.4.2.dev50
 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.2.dev45 → jaxsim-0.4.2.dev50}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.2.dev45'
-__version_tuple__ = version_tuple = (0, 4, 2, 'dev45')
+__version__ = version = '0.4.2.dev50'
+__version_tuple__ = version_tuple = (0, 4, 2, 'dev50')

{jaxsim-0.4.2.dev45 → jaxsim-0.4.2.dev50}/src/jaxsim/api/references.py

@@ -9,7 +9,6 @@ import jax_dataclasses
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim import exceptions
-from jaxsim.math import Adjoint
 from jaxsim.utils.tracing import not_tracing
 
 from .common import VelRepr
@@ -493,9 +492,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Extract the frame indices.
         frame_idxs = js.frame.names_to_idxs(frame_names=frame_names, model=model)
-        parent_link_idxs = jax.vmap(
-            lambda frame_idx: js.frame.idx_of_parent_link, in_axes=(None,)
-        )(model, frame_idx=frame_idxs)
+        parent_link_idxs = jax.vmap(js.frame.idx_of_parent_link, in_axes=(None,))(
+            model, frame_index=frame_idxs
+        )
 
         exceptions.raise_value_error_if(
             condition=jnp.logical_not(data.valid(model=model)),
@@ -527,25 +526,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             case _:
                 raise ValueError("Invalid velocity representation.")
 
-        W_H_L = js.model.forward_kinematics(model=model, data=data)
-
-        def convert_to_link_force(
-            W_f_F: jtp.MatrixLike, W_H_F: jtp.MatrixLike, parent_link_idx: jtp.ArrayLike
-        ) -> jtp.Matrix:
-            L_Xf_W = Adjoint.from_transform(W_H_L[parent_link_idx]).T
-
-            return L_Xf_W @ W_f_F
-
-        W_f_L_i = jax.vmap(convert_to_link_force)(W_f_F, W_H_Fi, parent_link_idxs)
-
         # Sum the forces on the parent links.
         mask = parent_link_idxs[:, jnp.newaxis] == jnp.arange(model.number_of_links())
-        W_f_L = mask.T @ W_f_L_i
+        W_f_L = mask.T @ W_f_F
 
         with self.switch_velocity_representation(
             velocity_representation=VelRepr.Inertial
         ):
-            return self.apply_link_forces(
+            references = self.apply_link_forces(
                 model=model,
                 data=data,
                 link_names=js.link.idxs_to_names(
@@ -554,3 +542,8 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
                 forces=W_f_L,
                 additive=additive,
             )
+
+        with references.switch_velocity_representation(
+            velocity_representation=self.velocity_representation
+        ):
+            return references

{jaxsim-0.4.2.dev45 → jaxsim-0.4.2.dev50}/src/jaxsim/rbda/crba.py

@@ -59,10 +59,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
         return (i_X_0,), None
 
-    (i_X_0,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=forward_pass_carry,
-        xs=jnp.arange(start=1, stop=model.number_of_links()),
+    (i_X_0,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=forward_pass_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(i_X_0,), None]
     )
 
     # ===================
@@ -128,10 +132,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
                 operand=carry,
             )
 
-        (j, Fi, M), _ = jax.lax.scan(
-            f=inner_fn,
-            init=carry_inner_fn,
-            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        (j, Fi, M), _ = (
+            jax.lax.scan(
+                f=inner_fn,
+                init=carry_inner_fn,
+                xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+            )
+            if model.number_of_links() > 1
+            else [(j, Fi, M), None]
         )
 
         Fi = i_X_0[j].T @ Fi
@@ -143,10 +151,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
     # This scan performs the backward pass to compute Mbj, Mjb and Mjj, that
     # also includes a fake while loop implemented with a scan and two cond.
-    (Mc, M), _ = jax.lax.scan(
-        f=backward_pass,
-        init=backward_pass_carry,
-        xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+    (Mc, M), _ = (
+        jax.lax.scan(
+            f=backward_pass,
+            init=backward_pass_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(Mc, M), None]
     )
 
     # Store the locked 6D rigid-body inertia matrix Mbb ∈ ℝ⁶ˣ⁶.

{jaxsim-0.4.2.dev45 → jaxsim-0.4.2.dev50}/src/jaxsim/rbda/forward_kinematics.py

@@ -75,10 +75,14 @@ def forward_kinematics_model(
 
         return (W_X_i,), None
 
-    (W_X_i,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=jnp.arange(start=1, stop=model.number_of_links()),
+    (W_X_i,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_kinematics_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(W_X_i,), None]
     )
 
     return jax.vmap(Adjoint.to_transform)(W_X_i)

{jaxsim-0.4.2.dev45 → jaxsim-0.4.2.dev50}/src/jaxsim/rbda/jacobian.py

@@ -67,10 +67,14 @@ def jacobian(
 
         return (i_X_0,), None
 
-    (i_X_0,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=np.arange(start=1, stop=model.number_of_links()),
+    (i_X_0,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_kinematics_carry,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(i_X_0,), None]
     )
 
     # ============================
@@ -105,10 +109,14 @@ def jacobian(
 
         return J, None
 
-    L_J_WL_B, _ = jax.lax.scan(
-        f=compute_jacobian,
-        init=J,
-        xs=np.arange(start=1, stop=model.number_of_links()),
+    L_J_WL_B, _ = (
+        jax.lax.scan(
+            f=compute_jacobian,
+            init=J,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [J, None]
     )
 
     return L_J_WL_B
@@ -184,10 +192,14 @@ def jacobian_full_doubly_left(
 
         return (B_X_i, J), None
 
-    (B_X_i, J), _ = jax.lax.scan(
-        f=compute_full_jacobian,
-        init=compute_full_jacobian_carry,
-        xs=np.arange(start=1, stop=model.number_of_links()),
+    (B_X_i, J), _ = (
+        jax.lax.scan(
+            f=compute_full_jacobian,
+            init=compute_full_jacobian_carry,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(B_X_i, J), None]
     )
 
     # Convert adjoints to SE(3) transforms.

{jaxsim-0.4.2.dev45 → jaxsim-0.4.2.dev50}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2.dev45
+Version: 0.4.2.dev50
 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>