Use JaxOPT PGD instead of L-BFGS

src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -373,7 +373,7 @@ def compute_contact_forces(
 
         # Create the objective function to minimize as a lambda computing the cost
         # from the optimized variables x.
-        objective = lambda x, A, b: jnp.sum(jnp.square(A @ x + b))
+        objective = lambda x: jnp.sum(jnp.square(A @ x + b))
 
         # ========================================
         # Helper function to run the L-BFGS solver
@@ -466,13 +466,27 @@ def continuing_criterion(carry: OptimizationCarry) -> jtp.Bool:
         maxiter = solver_options.pop("maxiter")
 
         # Compute the 3D linear force in C[W] frame.
-        solution, _ = run_optimization(
-            init_params=init_params,
+        # solution, _ = run_optimization(
+        #     init_params=init_params,
+        #     fun=objective,
+        #     opt=optax.lbfgs(**solver_options),
+        #     tol=tol,
+        #     maxiter=maxiter,
+        # )
+
+        import jaxopt
+
+        opt = jaxopt.ProjectedGradient(
             fun=objective,
-            opt=optax.lbfgs(**solver_options),
+            projection=jaxopt.projection.projection_non_negative,
+            # maxiter=maxiter,
             tol=tol,
-            maxiter=maxiter,
+            maxls=30,
+            implicit_diff=False,
+            # history_size=10,
+            # max_stepsize=100.0,
         )
+        solution = opt.run(init_params=init_params).params
 
         # Reshape the optimized solution to be a matrix of 3D contact forces.
         CW_fl_C = solution.reshape(-1, 3)