jaxsim 0.4.3.dev133__py3-none-any.whl → 0.4.3.dev139__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.dev133'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev133')
+__version__ = version = '0.4.3.dev139'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev139')

jaxsim/integrators/variable_step.py

@@ -262,7 +262,9 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             **kwargs,
         )
 
-    def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
+    def __call__(
+        self, x0: State, t0: Time, dt: TimeStep, **kwargs
+    ) -> tuple[NextState, dict[str, Any]]:
 
         # This method is called differently in three stages:
         #
@@ -294,14 +296,17 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         # In Stage 3, dt0 is taken from the previous step. If the integrator supports
         # FSAL, dxdt0 is taken from the previous step. Otherwise, it is computed by
         # evaluating the dynamics.
-        self.params["dt0"], self.params["dxdt0"] = jax.lax.cond(
+        self.params["dt0"], self.params["dxdt0"], aux_dict = jax.lax.cond(
             pred=jnp.logical_or("dt0" not in self.params, integrator_first_step),
-            true_fun=lambda params: estimate_step_size(
-                x0=x0, t0=t0, f=f, order=p, atol=self.atol, rtol=self.rtol
+            true_fun=lambda params: (
+                *estimate_step_size(
+                    x0=x0, t0=t0, f=f, order=p, atol=self.atol, rtol=self.rtol
+                ),
+                self.params.get("dxdt0", f(x0, t0))[1],
             ),
             false_fun=lambda params: (
                 params.get("dt0", jnp.array(0).astype(float)),
-                self.params.get("dxdt0", f(x0, t0)[0]),
+                *self.params.get("dxdt0", f(x0, t0)),
             ),
             operand=self.params,
         )
@@ -355,7 +360,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             # The output z contains multiple solutions (depending on the rows of b.T).
             with self.editable(validate=True) as integrator:
                 integrator.params = params
-                z = integrator._compute_next_state(x0=x0, t0=t0, dt=Δt0, **kwargs)
+                z, _ = integrator._compute_next_state(x0=x0, t0=t0, dt=Δt0, **kwargs)
                 params_next = integrator.params
 
             # Extract the high-order solution xf and the low-order estimate x̂f.
@@ -481,7 +486,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         with self.mutable_context(mutability=Mutability.MUTABLE):
             self.params = params_tf
 
-        return xf
+        return xf, aux_dict
 
     @property
     def order_of_solution(self) -> int:

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import dataclasses
+from collections.abc import Callable
 from typing import Any
 
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxopt
+import optax
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
@@ -297,24 +298,71 @@ class RelaxedRigidContacts(ContactModel):
         A = G + R
         b = CW_al_free_WC - a_ref
 
-        objective = lambda x: jnp.sum(jnp.square(A @ x + b))
+        objective = lambda x, A, b: jnp.sum(jnp.square(A @ x + b))
 
-        # Compute the 3D linear force in C[W] frame
-        opt = jaxopt.LBFGS(
-            fun=objective,
-            maxiter=self.parameters.max_iterations,
-            tol=self.parameters.tolerance,
-            maxls=30,
-            history_size=10,
-            max_stepsize=100.0,
-        )
+        def run_optimization(
+            init_params: jtp.Array,
+            fun: Callable,
+            opt: optax.GradientTransformation,
+            maxiter: jtp.Int,
+            tol: jtp.Float,
+            **kwargs,
+        ):
+            value_and_grad_fn = optax.value_and_grad_from_state(fun)
+
+            def step(carry):
+                params, state = carry
+                value, grad = value_and_grad_fn(
+                    params,
+                    state=state,
+                    A=A,
+                    b=b,
+                )
+                updates, state = opt.update(
+                    updates=grad,
+                    state=state,
+                    params=params,
+                    value=value,
+                    grad=grad,
+                    value_fn=fun,
+                    A=A,
+                    b=b,
+                )
+                params = optax.apply_updates(params, updates)
+                return params, state
+
+            def continuing_criterion(carry):
+                _, state = carry
+                iter_num = optax.tree_utils.tree_get(state, "count")
+                grad = optax.tree_utils.tree_get(state, "grad")
+                err = optax.tree_utils.tree_l2_norm(grad)
+                return (iter_num == 0) | ((iter_num < maxiter) & (err >= tol))
+
+            init_carry = (init_params, opt.init(init_params))
+            final_params, final_state = jax.lax.while_loop(
+                continuing_criterion, step, init_carry
+            )
+            return final_params, final_state
 
         init_params = (
             K[:, jnp.newaxis] * jnp.zeros_like(position).at[:, 2].set(δ)
             + D[:, jnp.newaxis] * velocity
         ).flatten()
 
-        CW_f_Ci = opt.run(init_params=init_params).params.reshape(-1, 3)
+        # Compute the 3D linear force in C[W] frame
+        CW_f_Ci, _ = run_optimization(
+            init_params=init_params,
+            A=A,
+            b=b,
+            maxiter=self.parameters.max_iterations,
+            opt=optax.lbfgs(
+                memory_size=10,
+            ),
+            fun=objective,
+            tol=self.parameters.tolerance,
+        )
+
+        CW_f_Ci = CW_f_Ci.reshape((-1, 3))
 
         def mixed_to_inertial(W_H_C: jax.Array, CW_fl: jax.Array) -> jax.Array:
             W_Xf_CW = Adjoint.from_transform(

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev133
+Version: 0.4.3.dev139
 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>
@@ -61,11 +61,11 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
 Requires-Dist: jax>=0.4.26
-Requires-Dist: jaxopt>=0.8.0
 Requires-Dist: jaxlib>=0.4.26
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax-dataclasses>=1.4.0
 Requires-Dist: pptree
+Requires-Dist: optax>=0.2.3
 Requires-Dist: qpax
 Requires-Dist: rod>=0.3.3
 Requires-Dist: typing-extensions; python_version < "3.12"

{jaxsim-0.4.3.dev133.dist-info → jaxsim-0.4.3.dev139.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=CaqH0tx8hRg2Qr8jQUmxmXjeOYVWMraQneScxC-v37c,428
+jaxsim/_version.py,sha256=3AziyYsOidZySopkrxUFcVGUEMK5LFxmckSLgq6qtds,428
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -19,7 +19,7 @@ jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,2077
 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
-jaxsim/integrators/variable_step.py,sha256=cJD98q5BaiSKvp_KY_1KN3PZpAUJR3L8YRmLX5WPJJo,21114
+jaxsim/integrators/variable_step.py,sha256=1VoSU3GeFcGEuP2dgZQ83sTkI5Xe-IThqKlRoVtwGSE,21270
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
 jaxsim/math/adjoint.py,sha256=o1FCipkGwPtMbN2gFNIyUV8ADF3TX5fxElpTEXK0bIs,4377
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
@@ -54,7 +54,7 @@ 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=aXWgChShcnJFZkE-1NIqaFQ2O2qC-xJt8iCZHft54Sk,14106
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=8kytUPYUmYXVrPEoHbCduFp5KOmOFPK4Vmqv3KhDqy8,15738
 jaxsim/rbda/contacts/rigid.py,sha256=6cU8kM8LMjEFbt8dtSg5nnz_uh4aD50sKw_svCzYUms,15633
 jaxsim/rbda/contacts/soft.py,sha256=NzzCYw5rvK8Fx_qH3fiMzPgey-KoxmRe9xkF3fluidE,18866
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
@@ -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.dev133.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev133.dist-info/METADATA,sha256=-YnzXwFdB6rR132PeB4w_BdqA4K_ul22Uv8uB3EzmQM,17277
-jaxsim-0.4.3.dev133.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev133.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev133.dist-info/RECORD,,
+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,,