jaxsim 0.4.3.dev64__py3-none-any.whl → 0.4.3.dev70__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.dev64'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev64')
+__version__ = version = '0.4.3.dev70'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev70')

jaxsim/api/data.py

@@ -6,10 +6,11 @@ from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp
+import jax.scipy.spatial.transform
 import jax_dataclasses
-import jaxlie
 
 import jaxsim.api as js
+import jaxsim.math
 import jaxsim.rbda
 import jaxsim.typing as jtp
 from jaxsim.rbda.contacts.soft import SoftContacts
@@ -195,10 +196,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         else:
             contacts_params = model.contact_model.parameters
 
-        W_H_B = jaxlie.SE3.from_rotation_and_translation(
-            translation=base_position,
-            rotation=jaxlie.SO3(wxyz=base_quaternion),
-        ).as_matrix()
+        W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
+            translation=base_position, quaternion=base_quaternion
+        )
 
         v_WB = JaxSimModelData.other_representation_to_inertial(
             array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
@@ -384,7 +384,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
         )
 
-        return (W_Q_B if not dcm else jaxlie.SO3(wxyz=W_Q_B).as_matrix()).astype(float)
+        return (W_Q_B if not dcm else jaxsim.math.Quaternion.to_dcm(W_Q_B)).astype(
+            float
+        )
 
     @jax.jit
     def base_transform(self) -> jtp.Matrix:
@@ -737,6 +739,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         )
 
 
+@functools.partial(jax.jit, static_argnames=["velocity_representation", "base_rpy_seq"])
 def random_model_data(
     model: js.model.JaxSimModel,
     *,
@@ -746,6 +749,11 @@ def random_model_data(
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
     ] = ((-1, -1, 0.5), 1.0),
+    base_rpy_bounds: tuple[
+        jtp.FloatLike | Sequence[jtp.FloatLike],
+        jtp.FloatLike | Sequence[jtp.FloatLike],
+    ] = (-jnp.pi, jnp.pi),
+    base_rpy_seq: str = "XYZ",
     joint_pos_bounds: (
         tuple[
             jtp.FloatLike | Sequence[jtp.FloatLike],
@@ -778,6 +786,10 @@ def random_model_data(
         key: The random key.
         velocity_representation: The velocity representation to use.
         base_pos_bounds: The bounds for the base position.
+        base_rpy_bounds:
+            The bounds for the euler angles used to build the base orientation.
+        base_rpy_seq:
+            The sequence of axes for rotation (using `Rotation` from scipy).
         joint_pos_bounds:
             The bounds for the joint positions (reading the joint limits if None).
         base_vel_lin_bounds: The bounds for the base linear velocity.
@@ -794,6 +806,8 @@ def random_model_data(
 
     p_min = jnp.array(base_pos_bounds[0], dtype=float)
     p_max = jnp.array(base_pos_bounds[1], dtype=float)
+    rpy_min = jnp.array(base_rpy_bounds[0], dtype=float)
+    rpy_max = jnp.array(base_rpy_bounds[1], dtype=float)
     v_min = jnp.array(base_vel_lin_bounds[0], dtype=float)
     v_max = jnp.array(base_vel_lin_bounds[1], dtype=float)
     ω_min = jnp.array(base_vel_ang_bounds[0], dtype=float)
@@ -819,9 +833,14 @@ def random_model_data(
             key=k1, shape=(3,), minval=p_min, maxval=p_max
         )
 
-        physics_model_state.base_quaternion = jaxlie.SO3.from_rpy_radians(
-            *jax.random.uniform(key=k2, shape=(3,), minval=0, maxval=2 * jnp.pi)
-        ).wxyz
+        physics_model_state.base_quaternion = jaxsim.math.Quaternion.to_wxyz(
+            xyzw=jax.scipy.spatial.transform.Rotation.from_euler(
+                seq=base_rpy_seq,
+                angles=jax.random.uniform(
+                    key=k2, shape=(3,), minval=rpy_min, maxval=rpy_max
+                ),
+            ).as_quat()
+        )
 
         if model.number_of_joints() > 0:
 

jaxsim/api/model.py

@@ -1747,14 +1747,18 @@ def link_contact_forces(
         data: The data of the considered model.
 
     Returns:
-        A (nL, 6) array containing the stacked 6D contact forces of the links,
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
         expressed in the frame corresponding to the active representation.
     """
 
+    # Note: the following code should be kept in sync with the function
+    # `jaxsim.api.ode.system_velocity_dynamics`. We cannot merge them since
+    # there we need to get also aux_data.
+
     # Compute the 6D forces applied to each collidable point expressed in the
     # inertial frame.
     with data.switch_velocity_representation(VelRepr.Inertial):
-        W_f_Ci = js.contact.collidable_point_forces(model=model, data=data)
+        W_f_C = js.contact.collidable_point_forces(model=model, data=data)
 
     # Construct the vector defining the parent link index of each collidable point.
     # We use this vector to sum the 6D forces of all collidable points rigidly
@@ -1763,29 +1767,28 @@ def link_contact_forces(
         model.kin_dyn_parameters.contact_parameters.body, dtype=int
     )
 
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
     # Sum the forces of all collidable points rigidly attached to a body.
-    # Since the contact forces W_f_Ci are expressed in the world frame,
+    # Since the contact forces W_f_C are expressed in the world frame,
     # we don't need any coordinate transformation.
-    W_f_Li = jax.vmap(
-        lambda nc: (
-            jnp.vstack(
-                jnp.equal(parent_link_index_of_collidable_points, nc).astype(int)
-            )
-            * W_f_Ci
-        ).sum(axis=0)
-    )(jnp.arange(model.number_of_links()))
-
-    # Convert the 6D forces to the active representation.
-    f_Li = jax.vmap(
-        lambda W_f_L: data.inertial_to_other_representation(
-            array=W_f_L,
-            other_representation=data.velocity_representation,
-            transform=data.base_transform(),
-            is_force=True,
-        )
-    )(W_f_Li)
+    W_f_L = mask.T @ W_f_C
+
+    # Create a references object to store the link forces.
+    references = js.references.JaxSimModelReferences.build(
+        model=model, link_forces=W_f_L, velocity_representation=VelRepr.Inertial
+    )
+
+    # Use the references object to convert the link forces to the velocity
+    # representation of data.
+    with references.switch_velocity_representation(data.velocity_representation):
+        f_L = references.link_forces(model=model, data=data)
 
-    return f_Li
+    return f_L
 
 
 # ======

jaxsim/api/ode.py

@@ -132,9 +132,16 @@ def system_velocity_dynamics(
     # with the terrain.
     W_f_Li_terrain = jnp.zeros_like(O_f_L).astype(float)
 
+    # Initialize a dictionary of auxiliary data.
+    # This dictionary is used to store additional data computed by the contact model.
     aux_data = {}
+
     if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
 
+        # Note: the following code should be kept in sync with the function
+        # `jaxsim.api.model.link_contact_forces`. We cannot merge them since
+        # here we need to get also aux_data.
+
         # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
         # along with contact-specific auxiliary states.
         with data.switch_velocity_representation(VelRepr.Inertial):

{jaxsim-0.4.3.dev64.dist-info → jaxsim-0.4.3.dev70.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=lLNskxtfHW1HqvnLRuhux3LlK89fMiZFUWknSYopw7k,426
+jaxsim/_version.py,sha256=TimgvJoa-WOzfWYDSSYt46dywfeE4QnnNq6VHB1jyaQ,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -7,13 +7,13 @@ jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
 jaxsim/api/contact.py,sha256=C_PgMjWYYiqpA7Oz3IxHeFgrp855-xG6AQr6Ze98CtI,21863
-jaxsim/api/data.py,sha256=mFUw2mj8AIXduW6HnkGN7eooZHfJhwnWbtYZfLF6gk4,28206
+jaxsim/api/data.py,sha256=QldUHniJqKrdNtAcXuRaS9UyeslJ0Rjvb17UA0Ca5Tw,29008
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=CcfSg5Mc8qb1mZeMQ4AK_ffZIsK5yOl7tu397pFhcDA,29369
 jaxsim/api/link.py,sha256=qPRtc8qqMRjZxUCZYXJMygbB6huDXBfIT1b1b8Durkw,18631
-jaxsim/api/model.py,sha256=K0q8-j-04f6B3MEXsctDGtWiuWlN3HbDrsS7zoPYStk,65871
-jaxsim/api/ode.py,sha256=VuOLvCFoyGLmhNf2vFP5BI9BAPz78V_RW5tJ4hrizsw,13041
+jaxsim/api/model.py,sha256=TLjgacgTXm-2YRGDA0Id9pe9nxIem28KoAls6Tdk9WM,66241
+jaxsim/api/ode.py,sha256=AE_obhb1u_xMGs8b4OlOXTvvQrbXOODqIDUWV5VGrJI,13376
 jaxsim/api/ode_data.py,sha256=7RSoBhfCJdP6P9InQbDwdBVpClPMMuetewI-6AWm-_0,20276
 jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
@@ -63,8 +63,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=FSiUvdnq4Y1T9Jaa_mw4ZBQJe8H7deLr3Kupxtlh4iI,11322
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev64.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev64.dist-info/METADATA,sha256=0-JS1eJjFMSaMzwqbCSpWYU2GcrZkxT1LBDo7lhWICo,17276
-jaxsim-0.4.3.dev64.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-jaxsim-0.4.3.dev64.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev64.dist-info/RECORD,,
+jaxsim-0.4.3.dev70.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev70.dist-info/METADATA,sha256=G6d-f4k63c6A1nozzlueaBz-fSgvz4zfpQIBecmGAiA,17276
+jaxsim-0.4.3.dev70.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev70.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev70.dist-info/RECORD,,

{jaxsim-0.4.3.dev64.dist-info → jaxsim-0.4.3.dev70.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.1.2)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any