jaxsim 0.4.3.dev139__py3-none-any.whl → 0.4.3.dev143__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/model.py

@@ -1907,8 +1907,8 @@ def step(
     dt: jtp.FloatLike,
     integrator: jaxsim.integrators.Integrator,
     integrator_state: dict[str, Any] | None = None,
-    joint_forces: jtp.VectorLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
     **kwargs,
 ) -> tuple[js.data.JaxSimModelData, dict[str, Any]]:
     """
@@ -1920,10 +1920,10 @@ def step(
         dt: The time step to consider.
         integrator: The integrator to use.
         integrator_state: The state of the integrator.
-        joint_forces: The joint forces to consider.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
+        joint_force_references: The joint force references to consider.
         kwargs: Additional kwargs to pass to the integrator.
 
     Returns:
@@ -1953,7 +1953,7 @@ def step(
         params=integrator_state_x0,
         # 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_forces, link_forces).
+        # by the pair (joint_force_references, link_forces).
         # Note that the wrapper of the system dynamics will override (state_x0, t0)
         # inside the passed data even if it is not strictly needed. This logic is
         # necessary to re-use the jit-compiled step function of compatible pytrees
@@ -1962,7 +1962,7 @@ def step(
             dict(
                 model=model,
                 data=data,
-                joint_forces=joint_forces,
+                joint_force_references=joint_force_references,
                 link_forces=link_forces,
             )
             | integrator_kwargs

jaxsim/api/ode.py

@@ -86,8 +86,8 @@ def system_velocity_dynamics(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
+    joint_force_references: jtp.Vector | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector, dict[str, Any]]:
     """
     Compute the dynamics of the system velocity.
@@ -95,10 +95,10 @@ def system_velocity_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint force references to apply.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
+        joint_force_references: The joint force references to apply.
 
     Returns:
         A tuple containing the derivative of the base 6D velocity in inertial-fixed
@@ -120,7 +120,7 @@ def system_velocity_dynamics(
     references = js.references.JaxSimModelReferences.build(
         model=model,
         link_forces=O_f_L,
-        joint_force_references=joint_forces,
+        joint_force_references=joint_force_references,
         data=data,
         velocity_representation=data.velocity_representation,
     )
@@ -192,7 +192,10 @@ def system_velocity_dynamics(
         f_L_total = references.link_forces(model=model, data=data)
 
         v̇_WB, s̈ = system_acceleration(
-            model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
+            model=model,
+            data=data,
+            joint_force_references=joint_force_references,
+            link_forces=f_L_total,
         )
 
     return v̇_WB, s̈, aux_data
@@ -202,8 +205,8 @@ def system_acceleration(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.VectorLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
     Compute the system acceleration in the active representation.
@@ -211,12 +214,13 @@ def system_acceleration(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
         link_forces:
-            The 6D forces to apply to the links expressed in the same representation of data.
+            The 6D forces to apply to the links expressed in the same
+            velocity representation of data.
+        joint_force_references: The joint force references to apply.
 
     Returns:
-        A tuple containing the base 6D acceleration in in the active representation
+        A tuple containing the base 6D acceleration in the active representation
         and the joint accelerations.
     """
 
@@ -232,9 +236,9 @@ def system_acceleration(
     ).astype(float)
 
     # Build joint torques if not provided.
-    τ = (
-        jnp.atleast_1d(joint_forces.squeeze())
-        if joint_forces is not None
+    τ_references = (
+        jnp.atleast_1d(joint_force_references.squeeze())
+        if joint_force_references is not None
         else jnp.zeros_like(data.joint_positions())
     ).astype(float)
 
@@ -243,15 +247,16 @@ def system_acceleration(
     # ====================
 
     # TODO: enforce joint limits
-    τ_position_limit = jnp.zeros_like(τ).astype(float)
+    τ_position_limit = jnp.zeros_like(τ_references).astype(float)
 
     # ====================
     # Joint friction model
     # ====================
 
-    τ_friction = jnp.zeros_like(τ).astype(float)
+    τ_friction = jnp.zeros_like(τ_references).astype(float)
 
     if model.dofs() > 0:
+
         # Static and viscous joint friction parameters
         kc = jnp.array(
             model.kin_dyn_parameters.joint_parameters.friction_static
@@ -271,22 +276,27 @@ def system_acceleration(
     # ========================
 
     # Compute the total joint forces.
-    τ_total = τ + τ_friction + τ_position_limit
+    τ_total = τ_references + τ_friction + τ_position_limit
 
+    # Store the link forces in a references object.
     references = js.references.JaxSimModelReferences.build(
         model=model,
         data=data,
         velocity_representation=data.velocity_representation,
-        joint_force_references=τ_total,
         link_forces=f_L,
     )
 
+    # Compute forward dynamics.
+    #
     # - Joint accelerations: s̈ ∈ ℝⁿ
     # - Base acceleration: v̇_WB ∈ ℝ⁶
+    #
+    # Note that ABA returns the base acceleration in the velocity representation
+    # stored in the `data` object.
     v̇_WB, s̈ = js.model.forward_dynamics_aba(
         model=model,
         data=data,
-        joint_forces=references.joint_force_references(model=model),
+        joint_forces=τ_total,
         link_forces=references.link_forces(model=model, data=data),
     )
 
@@ -337,8 +347,8 @@ def system_dynamics(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
+    joint_force_references: jtp.Vector | None = None,
     baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[ODEState, dict[str, Any]]:
     """
@@ -347,10 +357,10 @@ def system_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
+        joint_force_references: The joint force references to apply.
         baumgarte_quaternion_regularization:
             The Baumgarte regularization coefficient used to adjust the norm of the
             quaternion (only used in integrators not operating on the SO(3) manifold).
@@ -368,7 +378,7 @@ def system_dynamics(
     W_v̇_WB, s̈, aux_dict = system_velocity_dynamics(
         model=model,
         data=data,
-        joint_forces=joint_forces,
+        joint_force_references=joint_force_references,
         link_forces=link_forces,
     )
 

jaxsim/api/ode_data.py

@@ -38,16 +38,16 @@ class ODEInput(JaxsimDataclass):
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
-        joint_forces: jtp.VectorLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
     ) -> ODEInput:
         """
         Build an `ODEInput` from a `JaxSimModel`.
 
         Args:
             model: The `JaxSimModel` associated with the ODE input.
-            joint_forces: The vector of joint forces.
             link_forces: The matrix of external forces applied to the links.
+            joint_force_references: The vector of joint force references.
 
         Returns:
             The `ODEInput` built from the `JaxSimModel`.
@@ -60,8 +60,8 @@ class ODEInput(JaxsimDataclass):
         return ODEInput.build(
             physics_model_input=PhysicsModelInput.build_from_jaxsim_model(
                 model=model,
-                joint_forces=joint_forces,
                 link_forces=link_forces,
+                joint_force_references=joint_force_references,
             ),
             model=model,
         )
@@ -526,16 +526,16 @@ class PhysicsModelInput(JaxsimDataclass):
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
-        joint_forces: jtp.VectorLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
     ) -> PhysicsModelInput:
         """
         Build a `PhysicsModelInput` from a `JaxSimModel`.
 
         Args:
             model: The `JaxSimModel` associated with the input.
-            joint_forces: The vector of joint forces.
             link_forces: The matrix of external forces applied to the links.
+            joint_force_references: The vector of joint force references.
 
         Returns:
             A `PhysicsModelInput` instance.
@@ -546,7 +546,7 @@ class PhysicsModelInput(JaxsimDataclass):
         """
 
         return PhysicsModelInput.build(
-            joint_forces=joint_forces,
+            joint_force_references=joint_force_references,
             link_forces=link_forces,
             number_of_dofs=model.dofs(),
             number_of_links=model.number_of_links(),
@@ -554,8 +554,8 @@ class PhysicsModelInput(JaxsimDataclass):
 
     @staticmethod
     def build(
-        joint_forces: jtp.VectorLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
         number_of_dofs: jtp.Int | None = None,
         number_of_links: jtp.Int | None = None,
     ) -> PhysicsModelInput:
@@ -563,8 +563,8 @@ class PhysicsModelInput(JaxsimDataclass):
         Build a `PhysicsModelInput`.
 
         Args:
-            joint_forces: The vector of joint forces.
             link_forces: The matrix of external forces applied to the links.
+            joint_force_references: The vector of joint force references.
             number_of_dofs: The number of degrees of freedom of the model.
             number_of_links: The number of links of the model.
 
@@ -572,19 +572,21 @@ class PhysicsModelInput(JaxsimDataclass):
             A `PhysicsModelInput` instance.
         """
 
-        joint_forces = (
-            joint_forces if joint_forces is not None else jnp.zeros(number_of_dofs)
-        )
+        joint_force_references = jnp.atleast_1d(
+            jnp.array(joint_force_references, dtype=float).squeeze()
+            if joint_force_references is not None
+            else jnp.zeros(number_of_dofs)
+        ).astype(float)
 
-        link_forces = (
-            link_forces
+        link_forces = jnp.atleast_2d(
+            jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
             else jnp.zeros(shape=(number_of_links, 6))
-        )
+        ).astype(float)
 
         return PhysicsModelInput(
-            tau=jnp.array(joint_forces, dtype=float),
-            f_ext=jnp.array(link_forces, dtype=float),
+            tau=joint_force_references,
+            f_ext=link_forces,
         )
 
     @staticmethod

jaxsim/api/references.py

@@ -55,8 +55,8 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
-        joint_force_references: jtp.Vector | None = None,
-        link_forces: jtp.Matrix | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+        link_forces: jtp.MatrixLike | None = None,
         data: js.data.JaxSimModelData | None = None,
         velocity_representation: VelRepr | None = None,
     ) -> JaxSimModelReferences:
@@ -78,14 +78,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Create or adjust joint force references.
         joint_force_references = jnp.atleast_1d(
-            joint_force_references.squeeze()
+            jnp.array(joint_force_references, dtype=float).squeeze()
             if joint_force_references is not None
             else jnp.zeros(model.dofs())
         ).astype(float)
 
         # Create or adjust link forces.
         f_L = jnp.atleast_2d(
-            link_forces.squeeze()
+            jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
             else jnp.zeros((model.number_of_links(), 6))
         ).astype(float)
@@ -299,9 +299,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             A new `JaxSimModelReferences` object with the given joint force references.
         """
 
-        forces = jnp.array(forces)
+        forces = jnp.atleast_1d(jnp.array(forces, dtype=float).squeeze())
 
-        def replace(forces: jtp.VectorLike) -> JaxSimModelReferences:
+        def replace(forces: jtp.Vector) -> JaxSimModelReferences:
             return self.replace(
                 validate=True,
                 input=self.input.replace(

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -274,7 +274,9 @@ class RelaxedRigidContacts(ContactModel):
                     model=model,
                     data=data,
                     link_forces=references.link_forces(model=model, data=data),
-                    joint_forces=references.joint_force_references(model=model),
+                    joint_force_references=references.joint_force_references(
+                        model=model
+                    ),
                 )
             )
             BW_ν = data.generalized_velocity()

jaxsim/rbda/contacts/rigid.py

@@ -313,8 +313,10 @@ class RigidContacts(ContactModel):
                 js.ode.system_acceleration(
                     model=model,
                     data=data,
-                    joint_forces=references.joint_force_references(model=model),
                     link_forces=references.link_forces(model=model, data=data),
+                    joint_force_references=references.joint_force_references(
+                        model=model
+                    ),
                 )
             )
 

{jaxsim-0.4.3.dev139.dist-info → jaxsim-0.4.3.dev143.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=3AziyYsOidZySopkrxUFcVGUEMK5LFxmckSLgq6qtds,428
+jaxsim/_version.py,sha256=AUMtPwp-9k4PBssfGUg1S2TuNzVQkbYfnNIjR31wCZY,428
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -12,10 +12,10 @@ jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=ElahFk_RCcLvjTidH2qDOsY-m1gN1hXitCv4SvfgGYY,29260
 jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
-jaxsim/api/model.py,sha256=hDb8lVoRMhAWOMn1Q-MFHAa4wJY431-_nuLTKY0E3-Q,68200
-jaxsim/api/ode.py,sha256=gYSbtHWGCDP-IkUzQlH3t0fBKnK8qmxwhIvsbLG9lwU,13616
-jaxsim/api/ode_data.py,sha256=k1hVU1x8vuTVYdkf0cLhZ-oqeGJocXN2lLey-pS1_vo,20166
-jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
+jaxsim/api/model.py,sha256=Y0A9p7uKpfBdZLC9o_pTw7WQrmAaNJJBK1G97YkDRjQ,68260
+jaxsim/api/ode.py,sha256=B9CANQGLbrKAkmTxvWvkb0UGKC__r1nlHf-bGw963q0,13999
+jaxsim/api/ode_data.py,sha256=6q3IECaA351hLSGRovrOSKLo_P8IhJrW-p2-ZDlzGG8,20436
+jaxsim/api/references.py,sha256=fW77LitZ8DYgT6ZmUInJfm5luBV1mTcqcNRiC_i79og,20862
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=78MBs89GxsL0wU2yAexjvBZt3HEtfZoGVIN9f0a8yTc,20305
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
@@ -54,8 +54,8 @@ jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
 jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
 jaxsim/rbda/contacts/__init__.py,sha256=0UnO9ZR3BwdjQa276jOFbPi90pporr32LSc0qa9UUm4,369
 jaxsim/rbda/contacts/common.py,sha256=-eM8d1kvJ2E_2_kAgZJk4s3x8vDZHNSyOAinwPmRmEk,3469
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=8kytUPYUmYXVrPEoHbCduFp5KOmOFPK4Vmqv3KhDqy8,15738
-jaxsim/rbda/contacts/rigid.py,sha256=6cU8kM8LMjEFbt8dtSg5nnz_uh4aD50sKw_svCzYUms,15633
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=G3NoFrIGqzr_hrScLtsbmliTSqcVxDxheA1w2HRcU7A,15794
+jaxsim/rbda/contacts/rigid.py,sha256=WzPPo_UACOo_7olv_R1PelRjdjGqdWHkQ6hkSSfdBZk,15689
 jaxsim/rbda/contacts/soft.py,sha256=NzzCYw5rvK8Fx_qH3fiMzPgey-KoxmRe9xkF3fluidE,18866
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=Y0TGnUAGPuaeeSN8vbaSFjMXP5GYy3nxMCETjpUIMSA,5009
@@ -63,8 +63,8 @@ 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.dev139.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev139.dist-info/METADATA,sha256=uWmLplGU4SQyTl8iOGoW-isZe14lCh-mTkNXk1EThJM,17276
-jaxsim-0.4.3.dev139.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev139.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev139.dist-info/RECORD,,
+jaxsim-0.4.3.dev143.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev143.dist-info/METADATA,sha256=jZMIz7kCovjDtiHxhrp18oWpaYUhdX8t056wtbTv1us,17276
+jaxsim-0.4.3.dev143.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev143.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev143.dist-info/RECORD,,