[BUG FIX] Fix constraint solver mujoco consistency. (#1509)

Author: duburcqa

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -1229,14 +1229,6 @@ def func_solve(
                 if constraint_state.improved[i_b] < 1:
                     break
 
-                gradient = gs.ti_float(0.0)
-                for i_d in range(n_dofs):
-                    gradient += constraint_state.grad[i_d, i_b] * constraint_state.grad[i_d, i_b]
-                gradient = ti.sqrt(gradient)
-                improvement = constraint_state.prev_cost[i_b] - constraint_state.cost[i_b]
-                if gradient < tol_scaled or improvement < tol_scaled:
-                    break
-
 
 @ti.func
 def func_ls_init(
@@ -1301,10 +1293,10 @@ def func_ls_point_fn(
     alpha,
     constraint_state: array_class.ConstraintState,
 ):
-    tmp_quad_total0, tmp_quad_total1, tmp_quad_total2 = gs.ti_float(0.0), gs.ti_float(0.0), gs.ti_float(0.0)
-    tmp_quad_total0 = constraint_state.quad_gauss[0, i_b]
-    tmp_quad_total1 = constraint_state.quad_gauss[1, i_b]
-    tmp_quad_total2 = constraint_state.quad_gauss[2, i_b]
+    tmp_quad_total_0, tmp_quad_total_1, tmp_quad_total_2 = gs.ti_float(0.0), gs.ti_float(0.0), gs.ti_float(0.0)
+    tmp_quad_total_0 = constraint_state.quad_gauss[0, i_b]
+    tmp_quad_total_1 = constraint_state.quad_gauss[1, i_b]
+    tmp_quad_total_2 = constraint_state.quad_gauss[2, i_b]
     for i_c in range(constraint_state.n_constraints[i_b]):
         x = constraint_state.Jaref[i_c, i_b] + alpha * constraint_state.jv[i_c, i_b]
         active = 1
@@ -1331,14 +1323,16 @@ def func_ls_point_fn(
         elif i_c >= nef:
             active = x < 0
 
-        tmp_quad_total0 += qf_0 * active
-        tmp_quad_total1 += qf_1 * active
-        tmp_quad_total2 += qf_2 * active
+        tmp_quad_total_0 += qf_0 * active
+        tmp_quad_total_1 += qf_1 * active
+        tmp_quad_total_2 += qf_2 * active
 
-    cost = alpha * alpha * tmp_quad_total2 + alpha * tmp_quad_total1 + tmp_quad_total0
+    cost = alpha * alpha * tmp_quad_total_2 + alpha * tmp_quad_total_1 + tmp_quad_total_0
 
-    deriv_0 = 2 * alpha * tmp_quad_total2 + tmp_quad_total1
-    deriv_1 = 2 * tmp_quad_total2 + gs.EPS * (ti.abs(tmp_quad_total2) < gs.EPS)
+    deriv_0 = 2 * alpha * tmp_quad_total_2 + tmp_quad_total_1
+    deriv_1 = 2 * tmp_quad_total_2
+    if deriv_1 <= 0.0:
+        deriv_1 = gs.EPS
 
     constraint_state.ls_its[i_b] = constraint_state.ls_its[i_b] + 1
 
@@ -1615,7 +1609,6 @@ def func_solve_body(
     if ti.abs(alpha) < gs.EPS:
         constraint_state.improved[i_b] = 0
     else:
-        constraint_state.improved[i_b] = 1
         for i_d in range(n_dofs):
             constraint_state.qacc[i_d, i_b] = (
                 constraint_state.qacc[i_d, i_b] + constraint_state.search[i_d, i_b] * alpha
@@ -1640,56 +1633,58 @@ def func_solve_body(
             static_rigid_sim_config=static_rigid_sim_config,
         )
 
-        if ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.CG):
-            func_update_gradient(
+        if ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
+            func_nt_hessian_incremental(
                 i_b,
-                dofs_state=dofs_state,
                 entities_info=entities_info,
-                rigid_global_info=rigid_global_info,
                 constraint_state=constraint_state,
+                rigid_global_info=rigid_global_info,
                 static_rigid_sim_config=static_rigid_sim_config,
             )
 
-            constraint_state.cg_beta[i_b] = gs.ti_float(0.0)
-            constraint_state.cg_pg_dot_pMg[i_b] = gs.ti_float(0.0)
+        func_update_gradient(
+            i_b,
+            dofs_state=dofs_state,
+            entities_info=entities_info,
+            rigid_global_info=rigid_global_info,
+            constraint_state=constraint_state,
+            static_rigid_sim_config=static_rigid_sim_config,
+        )
 
-            for i_d in range(n_dofs):
-                constraint_state.cg_beta[i_b] += constraint_state.grad[i_d, i_b] * (
-                    constraint_state.Mgrad[i_d, i_b] - constraint_state.cg_prev_Mgrad[i_d, i_b]
-                )
-                constraint_state.cg_pg_dot_pMg[i_b] += (
-                    constraint_state.cg_prev_Mgrad[i_d, i_b] * constraint_state.cg_prev_grad[i_d, i_b]
-                )
+        tol_scaled = (rigid_global_info.meaninertia[i_b] * ti.max(1, n_dofs)) * static_rigid_sim_config.tolerance
+        improvement = constraint_state.prev_cost[i_b] - constraint_state.cost[i_b]
+        gradient = gs.ti_float(0.0)
+        for i_d in range(n_dofs):
+            gradient += constraint_state.grad[i_d, i_b] * constraint_state.grad[i_d, i_b]
+        gradient = ti.sqrt(gradient)
+        if gradient < tol_scaled or improvement < tol_scaled:
+            constraint_state.improved[i_b] = 0
+        else:
+            constraint_state.improved[i_b] = 1
 
-            constraint_state.cg_beta[i_b] = ti.max(
-                0.0, constraint_state.cg_beta[i_b] / ti.max(gs.EPS, constraint_state.cg_pg_dot_pMg[i_b])
-            )
-            for i_d in range(n_dofs):
-                constraint_state.search[i_d, i_b] = (
-                    -constraint_state.Mgrad[i_d, i_b]
-                    + constraint_state.cg_beta[i_b] * constraint_state.search[i_d, i_b]
-                )
+            if ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
+                for i_d in range(n_dofs):
+                    constraint_state.search[i_d, i_b] = -constraint_state.Mgrad[i_d, i_b]
+            else:
+                constraint_state.cg_beta[i_b] = gs.ti_float(0.0)
+                constraint_state.cg_pg_dot_pMg[i_b] = gs.ti_float(0.0)
 
-        elif ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
-            improvement = constraint_state.prev_cost[i_b] - constraint_state.cost[i_b]
-            if improvement > 0:
-                func_nt_hessian_incremental(
-                    i_b,
-                    entities_info=entities_info,
-                    constraint_state=constraint_state,
-                    rigid_global_info=rigid_global_info,
-                    static_rigid_sim_config=static_rigid_sim_config,
-                )
-                func_update_gradient(
-                    i_b,
-                    dofs_state=dofs_state,
-                    entities_info=entities_info,
-                    rigid_global_info=rigid_global_info,
-                    constraint_state=constraint_state,
-                    static_rigid_sim_config=static_rigid_sim_config,
+                for i_d in range(n_dofs):
+                    constraint_state.cg_beta[i_b] += constraint_state.grad[i_d, i_b] * (
+                        constraint_state.Mgrad[i_d, i_b] - constraint_state.cg_prev_Mgrad[i_d, i_b]
+                    )
+                    constraint_state.cg_pg_dot_pMg[i_b] += (
+                        constraint_state.cg_prev_Mgrad[i_d, i_b] * constraint_state.cg_prev_grad[i_d, i_b]
+                    )
+
+                constraint_state.cg_beta[i_b] = ti.max(
+                    0.0, constraint_state.cg_beta[i_b] / ti.max(gs.EPS, constraint_state.cg_pg_dot_pMg[i_b])
                 )
                 for i_d in range(n_dofs):
-                    constraint_state.search[i_d, i_b] = -constraint_state.Mgrad[i_d, i_b]
+                    constraint_state.search[i_d, i_b] = (
+                        -constraint_state.Mgrad[i_d, i_b]
+                        + constraint_state.cg_beta[i_b] * constraint_state.search[i_d, i_b]
+                    )
 
 
 @ti.func
@@ -1792,10 +1787,7 @@ def func_update_gradient(
         )
 
     elif ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
-        func_nt_chol_solve(
-            i_b,
-            constraint_state=constraint_state,
-        )
+        func_nt_chol_solve(i_b, constraint_state=constraint_state)
 
 
 @ti.func

tests/test_rigid_physics.py

@@ -357,7 +357,7 @@ def test_simple_kinematic_chain(gs_sim, mj_sim, tol):
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_frictionloss(gs_sim, mj_sim, tol):
     qvel = np.array([0.7, -0.9])
-    simulate_and_check_mujoco_consistency(gs_sim, mj_sim, qvel=qvel, num_steps=400, tol=1e-7)
+    simulate_and_check_mujoco_consistency(gs_sim, mj_sim, qvel=qvel, num_steps=2000, tol=tol)
 
     # Check that final velocity is almost zero
     gs_qvel = gs_sim.rigid_solver.dofs_state.vel.to_numpy()
@@ -394,14 +394,12 @@ def test_walker(gs_sim, mj_sim, gjk_collision, tol):
 @pytest.mark.parametrize("gs_solver", [gs.constraint_solver.CG, gs.constraint_solver.Newton])
 @pytest.mark.parametrize("gs_integrator", [gs.integrator.implicitfast, gs.integrator.Euler])
 @pytest.mark.parametrize("backend", [gs.cpu])
-def test_equality_joint(gs_sim, mj_sim, gs_solver):
+def test_equality_joint(gs_sim, mj_sim, gs_solver, tol):
     # there is an equality constraint
     assert gs_sim.rigid_solver.n_equalities == 1
 
     qpos = np.array((0.0, -1.0))
     qvel = np.array((1.0, -0.3))
-    # Note that it is impossible to be more accurate than this because of the inherent stiffness of the problem.
-    tol = 2e-8 if gs_solver == gs.constraint_solver.Newton else 1e-8
     simulate_and_check_mujoco_consistency(gs_sim, mj_sim, qpos, qvel, num_steps=300, tol=tol)
 
     # check if the two joints are equal
@@ -435,7 +433,7 @@ def test_equality_weld(gs_sim, mj_sim, gs_solver):
     # apply transform internally) is about 1e-15. This is fine and not surprising as it is consistent with machine
     # precision. These rounding errors are then amplified by 1e8 when computing the forces resulting from the kinematic
     # constraints. The constraints could be made softer by changing its impede parameters.
-    tol = 1e-7 if gs_solver == gs.constraint_solver.Newton else 2e-5
+    tol = 1e-7 if gs_solver == gs.constraint_solver.Newton else 5e-6
     simulate_and_check_mujoco_consistency(gs_sim, mj_sim, qpos, num_steps=300, tol=tol)