jaxsim 0.4.3.dev352__py3-none-any.whl → 0.4.3.dev362__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.dev352'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev352')
+__version__ = version = '0.4.3.dev362'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev362')

jaxsim/api/contact.py

@@ -151,7 +151,7 @@ def collidable_point_dynamics(
         kwargs: Additional keyword arguments to pass to the active contact model.
 
     Returns:
-        The 6D force applied to each eneabled collidable point and additional data based
+        The 6D force applied to each enabled collidable point and additional data based
         on the contact model configured:
         - Soft: the material deformation rate.
         - Rigid: no additional data.

jaxsim/api/model.py

@@ -2235,7 +2235,7 @@ def step(
         # by the pair (f_L, τ_references).
         # 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
+        # necessary to reuse the jit-compiled step function of compatible pytrees
         # of model and data produced e.g. by parameterized applications.
         **(
             dict(

jaxsim/rbda/contacts/common.py

@@ -324,10 +324,10 @@ class ContactModel(JaxsimDataclass):
             The initialized model and data objects.
         """
 
-        with model.editable(validate=validate) as model_out:
-            model_out.contact_model = self
+        with self.editable(validate=validate) as contact_model:
+            contact_model.parameters = data.contacts_params
 
-        with data.editable(validate=validate) as data_out:
-            data_out.contacts_params = data.contacts_params
+        with model.editable(validate=validate) as model_out:
+            model_out.contact_model = contact_model
 
-        return model_out, data_out
+        return model_out, data

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -10,6 +10,7 @@ import jax_dataclasses
 import optax
 
 import jaxsim.api as js
+import jaxsim.rbda.contacts
 import jaxsim.typing as jtp
 from jaxsim import logging
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
@@ -314,22 +315,20 @@ class RelaxedRigidContacts(common.ContactModel):
             joint_force_references=joint_force_references,
         )
 
-        def detect_contact(x: jtp.Array, y: jtp.Array, z: jtp.Array) -> jtp.Array:
-            x, y, z = jax.tree.map(jnp.squeeze, (x, y, z))
-
-            n̂ = model.terrain.normal(x=x, y=y).squeeze()
-            h = jnp.array([0, 0, z - model.terrain.height(x=x, y=y)])
-
-            return jnp.dot(h, n̂)
-
         # Compute the position and linear velocities (mixed representation) of
         # all collidable points belonging to the robot.
         position, velocity = js.contact.collidable_point_kinematics(
             model=model, data=data
         )
 
-        # Compute the activation state of the collidable points
-        δ = jax.vmap(detect_contact)(*position.T)
+        # Compute the penetration depth and velocity of the collidable points.
+        # Note that this function considers the penetration in the normal direction.
+        δ, _, n̂ = jax.vmap(common.compute_penetration_data, in_axes=(0, 0, None))(
+            position, velocity, model.terrain
+        )
+
+        # Compute the position in the constraint frame.
+        position_constraint = jax.vmap(lambda δ, n̂: -δ * n̂)(δ, n̂)
 
         # Compute the transforms of the implicit frames corresponding to the
         # collidable points.
@@ -356,24 +355,22 @@ class RelaxedRigidContacts(common.ContactModel):
             M = js.model.free_floating_mass_matrix(model=model, data=data)
 
             Jl_WC = jnp.vstack(
-                jax.vmap(lambda J, height: J * (height < 0))(
-                    js.contact.jacobian(model=model, data=data)[:, :3, :],
-                    δ,
+                jax.vmap(lambda J, δ: J * (δ > 0))(
+                    js.contact.jacobian(model=model, data=data)[:, :3, :], δ
                 )
             )
 
             J̇_WC = jnp.vstack(
-                jax.vmap(lambda J̇, height: J̇ * (height < 0))(
-                    js.contact.jacobian_derivative(model=model, data=data)[:, :3],
-                    δ,
+                jax.vmap(lambda J̇, δ: J̇ * (δ > 0))(
+                    js.contact.jacobian_derivative(model=model, data=data)[:, :3], δ
                 ),
             )
 
         # Compute the regularization terms.
-        a_ref, R, K, D = self._regularizers(
+        a_ref, R, *_ = self._regularizers(
             model=model,
-            penetration=δ,
-            velocity=velocity,
+            position_constraint=position_constraint,
+            velocity_constraint=velocity,
             parameters=data.contacts_params,
         )
 
@@ -435,6 +432,7 @@ class RelaxedRigidContacts(common.ContactModel):
 
                 return params, state
 
+            # TODO: maybe fix the number of iterations and switch to scan?
             def continuing_criterion(carry: OptimizationCarry) -> jtp.Bool:
 
                 _, state = carry
@@ -456,10 +454,20 @@ class RelaxedRigidContacts(common.ContactModel):
         # ======================================
 
         # Initialize the optimized forces with a linear Hunt/Crossley model.
-        init_params = (
-            K[:, jnp.newaxis] * jnp.zeros_like(position).at[:, 2].set(δ)
-            + D[:, jnp.newaxis] * velocity
-        ).flatten()
+        init_params = jax.vmap(
+            lambda p, v: jaxsim.rbda.contacts.SoftContacts.hunt_crossley_contact_model(
+                position=p,
+                velocity=v,
+                terrain=model.terrain,
+                K=1e6,
+                D=2e3,
+                p=0.0,
+                q=0.0,
+                # No tangential initial forces.
+                mu=0.0,
+                tangential_deformation=jnp.zeros(3),
+            )[0]
+        )(position, velocity).flatten()
 
         # Get the solver options.
         solver_options = self.solver_options
@@ -498,8 +506,8 @@ class RelaxedRigidContacts(common.ContactModel):
     @staticmethod
     def _regularizers(
         model: js.model.JaxSimModel,
-        penetration: jtp.Array,
-        velocity: jtp.Array,
+        position_constraint: jtp.Vector,
+        velocity_constraint: jtp.Vector,
         parameters: RelaxedRigidContactsParams,
     ) -> tuple:
         """
@@ -507,12 +515,13 @@ class RelaxedRigidContacts(common.ContactModel):
 
         Args:
             model: The jaxsim model.
-            penetration: The penetration of the collidable points.
-            velocity: The velocity of the collidable points.
+            penetration: The point position in the constraint frame.
+            velocity: The point velocity in the constraint frame.
             parameters: The parameters of the relaxed rigid contacts model.
 
         Returns:
-            A tuple containing the reference acceleration, the regularization matrix, the stiffness, and the damping.
+            A tuple containing the reference acceleration, the regularization matrix,
+            the stiffness, and the damping.
         """
 
         # Extract the parameters of the contact model.
@@ -545,70 +554,79 @@ class RelaxedRigidContacts(common.ContactModel):
         M_L = js.model.link_spatial_inertia_matrices(model=model)
 
         def imp_aref(
-            penetration: jtp.Array, velocity: jtp.Array
-        ) -> tuple[jtp.Array, jtp.Array]:
+            pos: jtp.Vector,
+            vel: jtp.Vector,
+        ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector, jtp.Vector]:
             """
             Calculates impedance and offset acceleration in constraint frame.
 
             Args:
-                penetration: penetration in constraint frame
-                velocity: velocity in constraint frame
+                pos: position in constraint frame.
+                vel: velocity in constraint frame.
 
             Returns:
+                ξ: computed impedance
                 a_ref: offset acceleration in constraint frame
-                R: regularization matrix
                 K: computed stiffness
                 D: computed damping
             """
-            position = jnp.zeros(shape=(3,)).at[2].set(penetration)
 
-            imp_x = jnp.abs(position) / width
-            imp_a = (1.0 / jnp.power(mid, p - 1)) * jnp.power(imp_x, p)
+            imp_x = jnp.abs(pos) / width
 
+            imp_a = (1.0 / jnp.power(mid, p - 1)) * jnp.power(imp_x, p)
             imp_b = 1 - (1.0 / jnp.power(1 - mid, p - 1)) * jnp.power(1 - imp_x, p)
-
             imp_y = jnp.where(imp_x < mid, imp_a, imp_b)
 
-            imp = jnp.clip(ξ_min + imp_y * (ξ_max - ξ_min), ξ_min, ξ_max)
-            imp = jnp.atleast_1d(jnp.where(imp_x > 1.0, ξ_max, imp))
+            # Compute the impedance.
+            ξ = ξ_min + imp_y * (ξ_max - ξ_min)
+            ξ = jnp.clip(ξ, ξ_min, ξ_max)
+            ξ = jnp.where(imp_x > 1.0, ξ_max, ξ)
+
+            # Compute the spring and damper parameters during runtime from the
+            # impedance and other contact parameters.
+            K = 1 / (ξ_max * Ω * ζ) ** 2
+            D = 2 / (ξ_max * Ω)
 
-            # When passing negative values, K and D represent a spring and damper, respectively.
-            K_f = jnp.where(K < 0, -K / ξ_max**2, 1 / (ξ_max * Ω * ζ) ** 2)
-            D_f = jnp.where(D < 0, -D / ξ_max, 2 / (ξ_max * Ω))
+            # If the user specifies K and D and they are negative, the computed `a_ref`
+            # becomes something more similar to a classic Baumgarte regularization.
+            K = jnp.where(K < 0, -K / ξ_max**2, K)
+            D = jnp.where(D < 0, -D / ξ_max, D)
 
-            a_ref = -jnp.atleast_1d(D_f * velocity + K_f * imp * position)
+            # Compute the reference acceleration.
+            a_ref = -(D * vel + K * ξ * pos)
 
-            return imp, a_ref, jnp.atleast_1d(K_f), jnp.atleast_1d(D_f)
+            return ξ, a_ref, K, D
 
         def compute_row(
             *,
-            link_idx: jtp.Float,
-            penetration: jtp.Array,
-            velocity: jtp.Array,
-        ) -> tuple[jtp.Array, jtp.Array]:
+            link_idx: jtp.Int,
+            pos: jtp.Vector,
+            vel: jtp.Vector,
+        ) -> tuple[jtp.Vector, jtp.Matrix, jtp.Vector, jtp.Vector]:
 
             # Compute the reference acceleration.
-            ξ, a_ref, K, D = imp_aref(
-                penetration=penetration,
-                velocity=velocity,
-            )
+            ξ, a_ref, K, D = imp_aref(pos=pos, vel=vel)
 
-            # Compute the regularization terms.
+            # Compute the regularization term.
             R = (
                 (2 * μ**2 * (1 - ξ) / (ξ + 1e-12))
                 * (1 + μ**2)
                 @ jnp.linalg.inv(M_L[link_idx, :3, :3])
             )
 
-            return jax.tree.map(lambda x: x * (penetration < 0), (a_ref, R, K, D))
+            # Return the computed values, setting them to zero in case of no contact.
+            is_active = (pos.dot(pos) > 0).astype(float)
+            return jax.tree.map(
+                lambda x: jnp.atleast_1d(x) * is_active, (a_ref, R, K, D)
+            )
 
         a_ref, R, K, D = jax.tree.map(
             f=jnp.concatenate,
             tree=(
                 *jax.vmap(compute_row)(
                     link_idx=parent_link_idx_of_enabled_collidable_points,
-                    penetration=penetration,
-                    velocity=velocity,
+                    pos=position_constraint,
+                    vel=velocity_constraint,
                 ),
             ),
         )

{jaxsim-0.4.3.dev352.dist-info → jaxsim-0.4.3.dev362.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev352
+Version: 0.4.3.dev362
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.4.3.dev352.dist-info → jaxsim-0.4.3.dev362.dist-info}/RECORD

@@ -1,18 +1,18 @@
 jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=0_67mou8TrQcvuWEau9dM_rtURNW7KokTnrV6Fq6e_g,428
+jaxsim/_version.py,sha256=yPZj4x5C1RSfHbZrDhkcc8pxpqIJyV7TN-gB5XXkOBo,428
 jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 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=D6RucrH9gnoUFLdmAEYwLGrimU0wLmuoDeOONu4ni74,25658
+jaxsim/api/contact.py,sha256=cPE7FIUycGzmmNW4Zh9w7qsWd4cQYowo_Tg3mL2evL0,25657
 jaxsim/api/data.py,sha256=ThRpoBlbdwf1N3xs8SWrY5d8RbfdYRwFcmkdIPgtee4,29004
 jaxsim/api/frame.py,sha256=yPSgNygHkvWlln4wShNt7vZm_fFobVEm7phsklNNyH8,12922
 jaxsim/api/joint.py,sha256=8rCIxRMeAidsaBbw7kkGp6z3-UmBPtqmYmV_arHDQJ8,7365
 jaxsim/api/kin_dyn_parameters.py,sha256=Y9wnMshz83Zm4UEPOAOTINdtfkBZ86w853c8Yi2qaVs,29670
 jaxsim/api/link.py,sha256=nHjffhNdi_xGkteMsqdb_hC9mdV9rNw7k3pl89Uhw_8,12798
-jaxsim/api/model.py,sha256=A88AaBZpWvQ-L9blFyl1GHvTWI05rvVFKbSaHzD77_k,79563
+jaxsim/api/model.py,sha256=HVoZ8AtFe5XOq8GlPbtpbMy-gzo4o1gM5gPchH0tHGw,79562
 jaxsim/api/ode.py,sha256=_t18avoCJngQk6eMFTGpaeahbpchQP20qJnUOCPkz8s,15360
 jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
 jaxsim/api/references.py,sha256=eIOk3MAOc9LJSKfI8M4WA8gGD-meo50vRfhXdea4sNI,20539
@@ -54,8 +54,8 @@ jaxsim/rbda/jacobian.py,sha256=L6Vn4Kf9I6wj-MYcFY6o67mgIfLFaaW4i2wNQJ2PDL0,10981
 jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
 jaxsim/rbda/utils.py,sha256=GLt7XIl1ROkx0_fnBCKUHYdB9_IBF3Yi4OnkHSX3gxA,5365
 jaxsim/rbda/contacts/__init__.py,sha256=L5MM-2pv76YPGzxExdz2EErgGBATuAjYnNHlq5QOySs,503
-jaxsim/rbda/contacts/common.py,sha256=ai49HeLQOsnckG0H2tUKW2KQ0Au_v9jRuNdnqie-YBk,11234
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=tbyskONuUhC6BZnZSpNUnlCjkI7LR6mCtmU_HimOAVE,20893
+jaxsim/rbda/contacts/common.py,sha256=mjOS1MJkf9zRfcAKBwmYtO5Vrcrv-kLFM57FfFB8LgM,11244
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=u7WliwuKff2RjS85eIEtJXbDLilZMqlz-j46-Pv7QAw,21681
 jaxsim/rbda/contacts/rigid.py,sha256=MSzkU6SFbW6CryNlyyxQ7K0-U-8k6VROGKv_DQrwqiw,17156
 jaxsim/rbda/contacts/soft.py,sha256=t6bqBfGAtV1AWoevY82LAcXy2XW8w_uu7bNywcyxF0s,17001
 jaxsim/rbda/contacts/visco_elastic.py,sha256=vQkfMuqQ3Qu8nbDTPY4jWBZjV3U7qtoRK1Aya3O3oFA,41424
@@ -65,8 +65,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=eEY28MZW0Lm_jJNt1NkFqZz0ek01tvhR46OXZYCo7tc,532
 jaxsim/utils/wrappers.py,sha256=ZY7olSORzZRvSzkdeNLj8yjwUIAt9L0Douwl7wItjpk,4008
-jaxsim-0.4.3.dev352.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev352.dist-info/METADATA,sha256=ZNz9p67E-hIsIWmtnTKkQZakk-LI8Nt_mKX3ywW_gOU,17513
-jaxsim-0.4.3.dev352.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-jaxsim-0.4.3.dev352.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev352.dist-info/RECORD,,
+jaxsim-0.4.3.dev362.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev362.dist-info/METADATA,sha256=YNsB6rTQXbGbbhmBKHhGQQ54P8nOASpjyEfGsXt5l-A,17513
+jaxsim-0.4.3.dev362.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+jaxsim-0.4.3.dev362.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev362.dist-info/RECORD,,