diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp.py b/genesis/engine/solvers/rigid/constraint_solver_decomp.py index 7f07658808..430e4c42b1 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp.py +++ b/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 + ( diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index b8cd5ed187..5d4ed9695d 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/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): diff --git a/genesis/utils/array_class.py b/genesis/utils/array_class.py index d6b11fab1a..35872b91b9 100644 --- a/genesis/utils/array_class.py +++ b/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_)), diff --git a/tests/test_render.py b/tests/test_render.py index c93ba04996..1ad008f371 100644 --- a/tests/test_render.py +++ b/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, diff --git a/tests/test_rigid_physics.py b/tests/test_rigid_physics.py index 9260b8d63e..662af3a545 100644 --- a/tests/test_rigid_physics.py +++ b/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(