jaxsim 0.2.1.dev101__py3-none-any.whl → 0.2.1.dev113__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev101'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev101')
+__version__ = version = '0.2.1.dev113'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev113')

jaxsim/api/common.py

@@ -87,7 +87,8 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
         array: jtp.Array,
         other_representation: VelRepr,
         transform: jtp.Matrix,
-        is_force: bool = False,
+        *,
+        is_force: bool,
     ) -> jtp.Array:
         r"""
         Convert a 6D quantity from inertial-fixed to another representation.
@@ -153,7 +154,8 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
         array: jtp.Array,
         other_representation: VelRepr,
         transform: jtp.Matrix,
-        is_force: bool = False,
+        *,
+        is_force: bool,
     ) -> jtp.Array:
         r"""
         Convert a 6D quantity from another representation to inertial-fixed.

jaxsim/api/model.py

@@ -451,27 +451,36 @@ def generalized_free_floating_jacobian(
     )
 
     # Compute the doubly-left free-floating full jacobian.
-    B_J_full_WX_B, _ = jaxsim.rbda.jacobian_full_doubly_left(
+    B_J_full_WX_B, B_H_L = jaxsim.rbda.jacobian_full_doubly_left(
         model=model,
         joint_positions=data.joint_positions(),
     )
 
-    # Update the input velocity representation such that `J_WL_I @ I_ν`.
+    # ======================================================================
+    # Update the input velocity representation such that v_WL = J_WL_I @ I_ν
+    # ======================================================================
+
     match data.velocity_representation:
+
         case VelRepr.Inertial:
+
             W_H_B = data.base_transform()
             B_X_W = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
+
             B_J_full_WX_I = B_J_full_WX_W = B_J_full_WX_B @ jax.scipy.linalg.block_diag(
                 B_X_W, jnp.eye(model.dofs())
             )
 
         case VelRepr.Body:
+
             B_J_full_WX_I = B_J_full_WX_B
 
         case VelRepr.Mixed:
+
             W_R_B = data.base_orientation(dcm=True)
             BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
             B_X_BW = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
+
             B_J_full_WX_I = B_J_full_WX_BW = (
                 B_J_full_WX_B
                 @ jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
@@ -480,32 +489,61 @@ def generalized_free_floating_jacobian(
         case _:
             raise ValueError(data.velocity_representation)
 
-    # Update the output velocity representation such that `O_v_WL_I = O_J_WL_I @ I_ν`.
+    # ====================================================================
+    # Create stacked Jacobian for each link by filtering the full Jacobian
+    # ====================================================================
+
+    κ_bool = model.kin_dyn_parameters.support_body_array_bool
+
+    # Keep only the columns of the full Jacobian corresponding to the support
+    # body array of each link.
+    B_J_WL_I = jax.vmap(
+        lambda κ: jnp.where(
+            jnp.hstack([jnp.ones(5), κ]), B_J_full_WX_I, jnp.zeros_like(B_J_full_WX_I)
+        )
+    )(κ_bool)
+
+    # =======================================================================
+    # Update the output velocity representation such that O_v_WL = O_J_WL @ ν
+    # =======================================================================
+
     match output_vel_repr:
+
         case VelRepr.Inertial:
+
             W_H_B = data.base_transform()
-            W_X_B = jaxlie.SE3.from_matrix(W_H_B).adjoint()
-            O_J_full_WX_I = W_J_full_WX_I = W_X_B @ B_J_full_WX_I
+            W_X_B = jaxsim.math.Adjoint.from_transform(W_H_B)
+
+            O_J_WL_I = W_J_WL_I = jax.vmap(lambda B_J_WL_I: W_X_B @ B_J_WL_I)(B_J_WL_I)
 
         case VelRepr.Body:
-            O_J_full_WX_I = B_J_full_WX_I
+
+            O_J_WL_I = L_J_WL_I = jax.vmap(
+                lambda B_H_L, B_J_WL_I: jaxsim.math.Adjoint.from_transform(
+                    B_H_L, inverse=True
+                )
+                @ B_J_WL_I
+            )(B_H_L, B_J_WL_I)
 
         case VelRepr.Mixed:
-            W_R_B = data.base_orientation(dcm=True)
-            BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
-            BW_X_B = jaxlie.SE3.from_matrix(BW_H_B).adjoint()
-            O_J_full_WX_I = BW_J_full_WX_I = BW_X_B @ B_J_full_WX_I
 
-        case _:
-            raise ValueError(output_vel_repr)
+            W_H_B = data.base_transform()
 
-    κ_bool = model.kin_dyn_parameters.support_body_array_bool
+            LW_H_L = jax.vmap(
+                lambda B_H_L: (W_H_B @ B_H_L).at[0:3, 3].set(jnp.zeros(3))
+            )(B_H_L)
 
-    O_J_WL_I = jax.vmap(
-        lambda κ: jnp.where(
-            jnp.hstack([jnp.ones(5), κ]), O_J_full_WX_I, jnp.zeros_like(O_J_full_WX_I)
-        )
-    )(κ_bool)
+            LW_H_B = jax.vmap(
+                lambda LW_H_L, B_H_L: LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
+            )(LW_H_L, B_H_L)
+
+            O_J_WL_I = LW_J_WL_I = jax.vmap(
+                lambda LW_H_B, B_J_WL_I: jaxsim.math.Adjoint.from_transform(LW_H_B)
+                @ B_J_WL_I
+            )(LW_H_B, B_J_WL_I)
+
+        case _:
+            raise ValueError(output_vel_repr)
 
     return O_J_WL_I
 

jaxsim/api/ode.py

@@ -113,7 +113,7 @@ def system_velocity_dynamics(
     ).astype(float)
 
     # Build link forces if not provided
-    W_f_L = (
+    O_f_L = (
         jnp.atleast_2d(link_forces.squeeze())
         if link_forces is not None
         else jnp.zeros((model.number_of_links(), 6))
@@ -125,7 +125,7 @@ def system_velocity_dynamics(
 
     # Initialize the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
     # with the terrain.
-    W_f_Li_terrain = jnp.zeros_like(W_f_L).astype(float)
+    W_f_Li_terrain = jnp.zeros_like(O_f_L).astype(float)
 
     # Initialize the 6D contact forces W_f ∈ ℝ^{n_c × 6} applied to collidable points,
     # expressed in the world frame.
@@ -183,7 +183,7 @@ def system_velocity_dynamics(
 
         # Compute the joint friction torque
         τ_friction = -(
-            jnp.diag(kc) @ jnp.sign(data.state.physics_model.joint_positions)
+            jnp.diag(kc) @ jnp.sign(data.state.physics_model.joint_velocities)
             + jnp.diag(kv) @ data.state.physics_model.joint_velocities
         )
 
@@ -194,6 +194,17 @@ def system_velocity_dynamics(
     # Compute the total joint forces
     τ_total = τ + τ_friction + τ_position_limit
 
+    references = js.references.JaxSimModelReferences.build(
+        model=model,
+        joint_force_references=τ_total,
+        link_forces=O_f_L,
+        data=data,
+        velocity_representation=data.velocity_representation,
+    )
+
+    with references.switch_velocity_representation(VelRepr.Inertial):
+        W_f_L = references.link_forces(model=model, data=data)
+
     # Compute the total external 6D forces applied to the links
     W_f_L_total = W_f_L + W_f_Li_terrain
 

jaxsim/api/references.py

@@ -202,17 +202,22 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         if not_tracing(self.input.physics_model.f_ext) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
-        # Helper function to convert a single 6D force to the active representation.
-        def convert(f_L: jtp.Vector) -> jtp.Vector:
-            return JaxSimModelReferences.inertial_to_other_representation(
-                array=f_L,
-                other_representation=self.velocity_representation,
-                transform=data.base_transform(),
-                is_force=True,
-            )
+        # Helper function to convert a single 6D force to the active representation
+        # considering as body the link (i.e. L_f_L and LW_f_L).
+        def convert(W_f_L: jtp.MatrixLike, W_H_L: jtp.ArrayLike) -> jtp.Matrix:
+
+            return jax.vmap(
+                lambda W_f_L, W_H_L: JaxSimModelReferences.inertial_to_other_representation(
+                    array=W_f_L,
+                    other_representation=self.velocity_representation,
+                    transform=W_H_L,
+                    is_force=True,
+                )
+            )(W_f_L, W_H_L)
 
-        # Convert to the desired representation.
-        f_L = jax.vmap(convert)(W_f_L[link_idxs, :])
+        # The f_L output is either L_f_L or LW_f_L, depending on the representation.
+        W_H_L = js.model.forward_kinematics(model=model, data=data)
+        f_L = convert(W_f_L=W_f_L[link_idxs, :], W_H_L=W_H_L[link_idxs, :, :])
 
         return f_L
 
@@ -319,7 +324,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         forces: jtp.MatrixLike,
         model: js.model.JaxSimModel | None = None,
         data: js.data.JaxSimModelData | None = None,
-        link_names: tuple[str, ...] | None = None,
+        link_names: tuple[str, ...] | str | None = None,
         additive: bool = False,
     ) -> Self:
         """
@@ -345,7 +350,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             Then, we always convert and store forces in inertial-fixed representation.
         """
 
-        f_L = jnp.array(forces)
+        f_L = jnp.atleast_2d(forces).astype(float)
 
         # Helper function to replace the link forces.
         def replace(forces: jtp.MatrixLike) -> JaxSimModelReferences:
@@ -380,6 +385,15 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # If we have the model, we can extract the link names if not provided.
         link_names = link_names if link_names is not None else model.link_names()
+
+        # Make sure that the link names are a tuple if they are provided by the user.
+        link_names = (link_names,) if isinstance(link_names, str) else link_names
+
+        if len(link_names) != f_L.shape[0]:
+            msg = "The number of link names ({}) must match the number of forces ({})"
+            raise ValueError(msg.format(len(link_names), f_L.shape[0]))
+
+        # Extract the link indices.
         link_idxs = js.link.names_to_idxs(link_names=link_names, model=model)
 
         # Compute the bias depending on whether we either set or add the link forces.
@@ -405,16 +419,24 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         if not_tracing(forces) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
-        # Helper function to convert a single 6D force to the inertial representation.
-        def convert(f_L: jtp.Vector) -> jtp.Vector:
-            return JaxSimModelReferences.other_representation_to_inertial(
-                array=f_L,
-                other_representation=self.velocity_representation,
-                transform=data.base_transform(),
-                is_force=True,
-            )
+        # Helper function to convert a single 6D force to the inertial representation
+        # considering as body the link (i.e. L_f_L and LW_f_L).
+        def convert_using_link_frame(
+            f_L: jtp.MatrixLike, W_H_L: jtp.ArrayLike
+        ) -> jtp.Matrix:
+
+            return jax.vmap(
+                lambda f_L, W_H_L: JaxSimModelReferences.other_representation_to_inertial(
+                    array=f_L,
+                    other_representation=self.velocity_representation,
+                    transform=W_H_L,
+                    is_force=True,
+                )
+            )(f_L, W_H_L)
 
-        W_f_L = jax.vmap(convert)(f_L)
+        # The f_L input is either L_f_L or LW_f_L, depending on the representation.
+        W_H_L = js.model.forward_kinematics(model=model, data=data)
+        W_f_L = convert_using_link_frame(f_L=f_L, W_H_L=W_H_L[link_idxs, :, :])
 
         return replace(
             forces=self.input.physics_model.f_ext.at[link_idxs, :].set(W_f0_L + W_f_L)

{jaxsim-0.2.1.dev101.dist-info → jaxsim-0.2.1.dev113.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev101
+Version: 0.2.1.dev113
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.2.1.dev101.dist-info → jaxsim-0.2.1.dev113.dist-info}/RECORD

@@ -1,20 +1,20 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=G1D0JLzyL0ACAlwNnaPkzuJCqBQM59tregHl11EAaSo,428
+jaxsim/_version.py,sha256=tWgdOeciml9_rSgehjRMlzuVy3GjT9rU5sEmQnyWbSM,428
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=Yof6otFi-mLWAs1rqjmeNJTOWIH9gn7BdU5EIjiL6Ts,13481
-jaxsim/api/common.py,sha256=DV-WZG28sikXopNv458aYvpLjmiAtFr5LRscOwXusuk,6640
+jaxsim/api/common.py,sha256=bqQ__pIQZbh-j8rkoHUkYHAgGiJnDzjHG-q4Ny0OOYQ,6646
 jaxsim/api/contact.py,sha256=Cvr-EfQtHP3nymtWdo-9WWU24Bkta-2Pp3nKsdjo6uc,12778
 jaxsim/api/data.py,sha256=xfKJz6Rw0YTk-EHCGiT8BFQrs_ggOz01lRi1Qh1mb28,27256
 jaxsim/api/frame.py,sha256=0YXOrGmx3cSQqa4_Ky-n6zyup3I3xvXNEgub-Bc5xUw,6222
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
 jaxsim/api/kin_dyn_parameters.py,sha256=zMca7OmCsCWK_cavLTSZSeYh9Qu1-409cdsyWvWPAUQ,26090
 jaxsim/api/link.py,sha256=rypTwkMf9HJ5UuAtHRJh0LqqdJWcLKTtTjWcjduEsF0,9842
-jaxsim/api/model.py,sha256=Ii17tBzCkMHCa_G7plKPzNcW-bj3QGuGI9jCklPQStM,54294
-jaxsim/api/ode.py,sha256=6l-6i2YHagsQvR8Ac-_fmO6P0hBVT6NkHhwXnrdITEg,9785
+jaxsim/api/model.py,sha256=1HlQ5FMzeJAk-cE1pmELgVjzMYUX9-iipw3N4WssAL4,55435
+jaxsim/api/ode.py,sha256=BfvV_14uu0szWecoDiV8rTu-dvSFLK7eyrO38ZqHB_w,10157
 jaxsim/api/ode_data.py,sha256=D6FzMkvY_qNuoFEImyp7sxAk-0pJOd3oZeSr9bBTcLk,23089
-jaxsim/api/references.py,sha256=Lvskf17r619KKxwCJP7hAAty2kaXgDXJX1uKqoDIDgo,15483
+jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=9HXRVFo95Mpt6RcVhBrOfvOO7mDxqbkXeg_lKUibEFY,20693
 jaxsim/integrators/fixed_step.py,sha256=JXaEyEzfSiYea0GnPA7l27J3X0YPB0e25D4qfrxAvzQ,2766
@@ -58,8 +58,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=EJMcblYKUjxw9HJShVf81Ig3pHUJno6Dx6h-RnY--wM,2040
-jaxsim-0.2.1.dev101.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.2.1.dev101.dist-info/METADATA,sha256=e2Rvemht7C6PBWz899beW2a-92wHtkPzxKpaRelY-eg,9745
-jaxsim-0.2.1.dev101.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.1.dev101.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.1.dev101.dist-info/RECORD,,
+jaxsim-0.2.1.dev113.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.2.1.dev113.dist-info/METADATA,sha256=NH5olEy-GNgTb5Ibe-q3FH2bu_WOuY0hbQBJZCcOiS4,9745
+jaxsim-0.2.1.dev113.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.1.dev113.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.1.dev113.dist-info/RECORD,,