[BUG FIX] Fix joint friction loss. (Genesis-Embodied-AI#1555)

Author: duburcqa

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -1035,22 +1035,22 @@ def add_frictionloss_constraints(
                 for i_d in range(joints_info.dof_start[I_j], joints_info.dof_end[I_j]):
                     I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d
 
-                    if dofs_info.frictionloss[I_d] > 0:
+                    if dofs_info.frictionloss[I_d] > gs.EPS:
                         jac = 1.0
                         jac_qvel = jac * dofs_state.vel[i_d, i_b]
                         imp, aref = gu.imp_aref(joints_info.sol_params[I_j], 0.0, jac_qvel, 0.0)
-                        diag = ti.max(dofs_info.invweight[I_d] * (1 - imp) / imp, gs.EPS)
+                        diag = ti.max(dofs_info.invweight[I_d] * (1.0 - imp) / imp, gs.EPS)
 
-                        n_con = constraint_state.n_constraints[i_b]
+                        i_con = ti.atomic_add(constraint_state.n_constraints[i_b], 1)
                         ti.atomic_add(constraint_state.n_constraints_frictionloss[i_b], 1)
-                        constraint_state.n_constraints[i_b] = n_con + 1
-                        constraint_state.diag[n_con, i_b] = diag
-                        constraint_state.aref[n_con, i_b] = aref
-                        constraint_state.efc_D[n_con, i_b] = 1 / diag
-                        constraint_state.efc_frictionloss[n_con, i_b] = dofs_info.frictionloss[I_d]
+
+                        constraint_state.diag[i_con, i_b] = diag
+                        constraint_state.aref[i_con, i_b] = aref
+                        constraint_state.efc_D[i_con, i_b] = 1.0 / diag
+                        constraint_state.efc_frictionloss[i_con, i_b] = dofs_info.frictionloss[I_d]
                         for i_d2 in range(n_dofs):
-                            constraint_state.jac[n_con, i_d2, i_b] = gs.ti_float(0.0)
-                        constraint_state.jac[n_con, i_d, i_b] = jac
+                            constraint_state.jac[i_con, i_d2, i_b] = gs.ti_float(0.0)
+                        constraint_state.jac[i_con, i_d, i_b] = jac
 
 
 @ti.func
@@ -1430,19 +1430,21 @@ def func_ls_point_fn(
     alpha,
     constraint_state: array_class.ConstraintState,
 ):
+    ne = constraint_state.n_constraints_equality[i_b]
+    nef = ne + constraint_state.n_constraints_frictionloss[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
-        ne = constraint_state.n_constraints_equality[i_b]
-        nef = ne + constraint_state.n_constraints_frictionloss[i_b]
         qf_0 = constraint_state.quad[i_c, 0, i_b]
         qf_1 = constraint_state.quad[i_c, 1, i_b]
         qf_2 = constraint_state.quad[i_c, 2, i_b]
-        if i_c >= ne and i_c < nef:
+
+        active = gs.ti_bool(True)  # Equality constraints
+        if ne <= i_c and i_c < nef:  # Friction constraints
             f = constraint_state.efc_frictionloss[i_c, i_b]
             r = 1.0 / ti.max(constraint_state.efc_D[i_c, i_b], gs.EPS)
             rf = r * f
@@ -1457,7 +1459,7 @@ def func_ls_point_fn(
                     f * constraint_state.jv[i_c, i_b]
                 )
                 qf_2 = 0.0
-        elif i_c >= nef:
+        elif nef <= i_c:  # Contact constraints
             active = x < 0
 
         tmp_quad_total_0 += qf_0 * active
@@ -1835,6 +1837,8 @@ def func_update_constraint(
     static_rigid_sim_config: ti.template(),
 ):
     n_dofs = constraint_state.qfrc_constraint.shape[0]
+    ne = constraint_state.n_constraints_equality[i_b]
+    nef = ne + constraint_state.n_constraints_frictionloss[i_b]
 
     constraint_state.prev_cost[i_b] = cost[i_b]
     cost[i_b] = gs.ti_float(0.0)
@@ -1843,11 +1847,10 @@ def func_update_constraint(
     for i_c in range(constraint_state.n_constraints[i_b]):
         if ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
             constraint_state.prev_active[i_c, i_b] = constraint_state.active[i_c, i_b]
-        constraint_state.active[i_c, i_b] = 1
+        constraint_state.active[i_c, i_b] = True
+
         floss_force = gs.ti_float(0.0)
-        ne = constraint_state.n_constraints_equality[i_b]
-        nef = ne + constraint_state.n_constraints_frictionloss[i_b]
-        if i_c >= ne and i_c < nef:
+        if ne <= i_c and i_c < nef:  # Friction constraints
             f = constraint_state.efc_frictionloss[i_c, i_b]
             r = 1.0 / ti.max(constraint_state.efc_D[i_c, i_b], gs.EPS)
             rf = r * f
@@ -1858,7 +1861,7 @@ def func_update_constraint(
             floss_cost_local = linear_neg * f * (-0.5 * rf - constraint_state.Jaref[i_c, i_b])
             floss_cost_local += linear_pos * f * (-0.5 * rf + constraint_state.Jaref[i_c, i_b])
             cost[i_b] = cost[i_b] + floss_cost_local
-        elif i_c >= nef:
+        elif nef <= i_c:  # Contact constraints
             constraint_state.active[i_c, i_b] = constraint_state.Jaref[i_c, i_b] < 0
 
         constraint_state.efc_force[i_c, i_b] = floss_force + (

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -952,11 +952,8 @@ def _func_constraint_force(self):
         # from genesis.utils.tools import create_timer
 
         # timer = create_timer(name="constraint_force", level=2, ti_sync=True, skip_first_call=True)
-        if self._enable_collision or self._enable_joint_limit or self.n_equalities > 0:
-            self.constraint_solver.constraint_state.n_constraints.fill(0)
-            self.constraint_solver.constraint_state.n_constraints_equality.fill(0)
-            self._func_constraint_clear()
-            # timer.stamp("constraint_solver.clear")
+        self._func_constraint_clear()
+        # timer.stamp("constraint_solver.clear")
 
         if self._enable_collision:
             self.collider.detection()
@@ -969,6 +966,7 @@ def _func_constraint_force(self):
     def _func_constraint_clear(self):
         self.constraint_solver.constraint_state.n_constraints.fill(0)
         self.constraint_solver.constraint_state.n_constraints_equality.fill(0)
+        self.constraint_solver.constraint_state.n_constraints_frictionloss.fill(0)
         self.collider._collider_state.n_contacts.fill(0)
 
     def _func_forward_dynamics(self):

genesis/utils/array_class.py

@@ -184,7 +184,7 @@ def get_constraint_state(constraint_solver, solver):
         "efc_D": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
         "efc_frictionloss": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
         "efc_force": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
-        "active": V(dtype=gs.ti_int, shape=solver._batch_shape(len_constraints_)),
+        "active": V(dtype=gs.ti_bool, shape=solver._batch_shape(len_constraints_)),
         "prev_active": V(dtype=gs.ti_int, shape=solver._batch_shape(len_constraints_)),
         "qfrc_constraint": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
         "qacc": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),

tests/test_render.py

@@ -118,7 +118,7 @@ def test_segmentation(segmentation_level, particle_mode):
 
 
 @pytest.mark.required
-@pytest.mark.flaky(reruns=3, condition=(sys.platform == "darwin"))
+@pytest.mark.flaky(reruns=3, condition=(sys.platform == "darwin"), reason="Flaky on MacOS with CPU-based OpenGL")
 def test_batched_offscreen_rendering(tmp_path, show_viewer, tol):
     scene = gs.Scene(
         vis_options=gs.options.VisOptions(
@@ -275,6 +275,7 @@ def test_batched_offscreen_rendering(tmp_path, show_viewer, tol):
 
 
 @pytest.mark.required
+@pytest.mark.flaky(reruns=3, condition=(sys.platform == "darwin"), reason="Flaky on MacOS with CPU-based OpenGL")
 def test_render_api(show_viewer):
     scene = gs.Scene(
         show_viewer=show_viewer,

tests/test_rigid_physics.py

@@ -1654,6 +1654,37 @@ def test_mass_mat(show_viewer, tol):
     assert_allclose(mass_mat, mass_mat_1, tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_frictionloss_advanced(show_viewer, tol):
+    scene = gs.Scene(
+        show_viewer=show_viewer,
+        show_FPS=False,
+    )
+    scene.add_entity(gs.morphs.Plane())
+    asset_path = get_hf_dataset(pattern="SO101/*")
+    robot = scene.add_entity(
+        morph=gs.morphs.MJCF(
+            file=f"{asset_path}/SO101/so101_new_calib.xml",
+        ),
+    )
+    box = scene.add_entity(
+        gs.morphs.Box(
+            size=(0.025, 0.025, 0.025),
+        ),
+    )
+    scene.build(n_envs=0)
+
+    scene.reset()
+    box.set_pos(torch.tensor((0.1, 0.0, 1.0), dtype=gs.tc_float, device=gs.device))
+    for _ in range(200):
+        scene.step()
+
+    assert_allclose(robot.get_contacts()["position"][:, 2].min(), 0.0, tol=1e-4)
+    # assert_allclose(torch.stack([geom.get_AABB() for geom in robot.geoms])[:, :, 2].min(), 0.0, tol=1e-3)
+    assert_allclose(box.get_dofs_velocity(), 0.0, tol=tol)
+
+
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_nonconvex_collision(show_viewer):
     scene = gs.Scene(