diff --git a/genesis/constants.py b/genesis/constants.py index bebf4307c6..7a9f1d8a84 100644 --- a/genesis/constants.py +++ b/genesis/constants.py @@ -59,6 +59,7 @@ class integrator(IntEnum): class constraint_solver(IntEnum): CG = 0 Newton = 1 + ComFree = 2 # rigid solver broadphase traversal strategy diff --git a/genesis/engine/solvers/rigid/comfree/__init__.py b/genesis/engine/solvers/rigid/comfree/__init__.py new file mode 100644 index 0000000000..f41deeafac --- /dev/null +++ b/genesis/engine/solvers/rigid/comfree/__init__.py @@ -0,0 +1,14 @@ +""" +ComFree (Complementarity-Free) constraint solver for rigid body simulation. + +Implements the analytical contact resolution method from: + ComFree-Sim: A GPU-Parallelized Analytical Contact Physics Engine + for Scalable Contact-Rich Robotics Simulation and Control + (arXiv:2603.12185, reference impl: https://github.com/asu-iris/comfree_warp) + +Instead of iterative complementarity-based solving (Newton/CG), ComFree computes +constraint forces in closed form via an impedance-style prediction-correction +update, then solves a single mass-weighted linear system for the acceleration. +""" + +from .solver import ComFreeSolver diff --git a/genesis/engine/solvers/rigid/comfree/solver.py b/genesis/engine/solvers/rigid/comfree/solver.py new file mode 100644 index 0000000000..59f789dbf3 --- /dev/null +++ b/genesis/engine/solvers/rigid/comfree/solver.py @@ -0,0 +1,353 @@ +""" +ComFree (Complementarity-Free) constraint solver. + +Replaces iterative complementarity-based constraint solving (Newton/CG) with a +single-pass analytical contact-force computation. It reuses the existing Genesis +constraint-assembly infrastructure (Jacobians, collision detection, equality / +joint-limit / frictionloss rows) but resolves forces with an impedance-style +prediction-correction update instead of an optimization loop. + +Reference implementation: + comfree_warp/comfree_core/_src/forward.py (compute_qfrc_total / forward_comfree) + comfree_warp/comfree_core/_src/constraint.py (_efc_row -> efc_dist / efc_mass) +""" + +from typing import TYPE_CHECKING + +import numpy as np +import quadrants as qd + +import genesis as gs +import genesis.utils.array_class as array_class + +from ..abd import func_solve_mass +from ..collider.contact_island import ContactIsland + +if TYPE_CHECKING: + from genesis.engine.solvers.rigid.rigid_solver import RigidSolver + + +class ComFreeSolver: + """Complementarity-free analytical constraint solver. + + For each constraint row ``c`` with Jacobian ``J_c``, signed position ``efc_dist`` + and effective mass ``efc_mass`` the per-step update is (reference forward.py:67-85):: + + v_pred = v + acc_smooth * dt # smooth predicted velocity + efc_vel = J_c @ v_pred + efc_pen = efc_vel * dt + efc_dist # predictive penetration + force = max(efc_mass * (-d*efc_vel - k*efc_pen), 0) + qfrc += J_c^T * force + + then ``qacc = M^{-1} (qf_smooth + qfrc)``. ``k`` and ``d`` are the user stiffness + and damping scaled by ``1/dt``. + + Parameters + ---------- + rigid_solver : RigidSolver + The parent rigid body solver. + comfree_stiffness : float + Global stiffness parameter ``k_user`` (default 0.2). + comfree_damping : float + Global damping parameter ``d_user`` (default 0.001). + """ + + def __init__(self, rigid_solver: "RigidSolver", comfree_stiffness: float = 0.2, comfree_damping: float = 0.001): + self._solver = rigid_solver + self._collider = rigid_solver.collider + self._B = rigid_solver._B + + self._comfree_stiffness = float(comfree_stiffness) + self._comfree_damping = float(comfree_damping) + + # Over-estimate the constraint count exactly like the standard solver: + # * 4 constraints per contact + # * 1 constraint per 1-DoF joint limit + # * 1 constraint per DoF frictionloss + # * up to 6 constraints per equality (weld) + self.len_constraints = int( + 4 * rigid_solver.collider._collider_info.max_contact_pairs[None] + + sum(joint.type in (gs.JOINT_TYPE.REVOLUTE, gs.JOINT_TYPE.PRISMATIC) for joint in self._solver.joints) + + self._solver.n_dofs + + self._solver.n_candidate_equalities_ * 6 + ) + self.len_constraints_ = max(1, self.len_constraints) + + # ComFree never builds a sparse Jacobian; needed by get_constraint_state allocation. + self.sparse_solve = False + + # Reuse the standard constraint-state allocation (jac, efc_dist, efc_mass, efc_force, ...). + self.constraint_state = array_class.get_constraint_state(self, self._solver) + self.constraint_state.qd_n_equalities.from_numpy( + np.full((self._solver._B,), self._solver.n_equalities, dtype=gs.np_int) + ) + + self.reset() + + # ContactIsland is required as a parameter by some kernels (e.g. kernel_forward_dynamics); + # it is inert when hibernation / contact-island mode is disabled. + self.contact_island = ContactIsland(self._collider) + + def reset(self, envs_idx=None): + envs_idx = self._solver._scene._sanitize_envs_idx(envs_idx) + _comfree_solver_kernel_reset(envs_idx, self.constraint_state, self._solver._static_rigid_sim_config) + + def clear(self, envs_idx=None): + self.reset(envs_idx) + envs_idx = self._solver._scene._sanitize_envs_idx(envs_idx) + _comfree_solver_kernel_clear( + envs_idx, + self.constraint_state, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + ) + + def add_equality_constraints(self): + from ..constraint.solver import add_equality_constraints + + add_equality_constraints( + self._solver.links_info, + self._solver.links_state, + self._solver.dofs_state, + self._solver.dofs_info, + self._solver.joints_info, + self._solver.equalities_info, + self.constraint_state, + self._collider._collider_state, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + ) + + def add_inequality_constraints(self): + from ..constraint.solver import add_inequality_constraints + + add_inequality_constraints( + self._solver.links_info, + self._solver.links_state, + self._solver.dofs_state, + self._solver.dofs_info, + self._solver.joints_info, + self.constraint_state, + self._collider._collider_state, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + ) + + def add_constraints(self): + """Contact-island assembly path (kept for interface parity with ConstraintSolverIsland).""" + from ..constraint.solver_island import add_constraints + + add_constraints( + self._solver.links_info, + self._solver.links_state, + self._solver.dofs_state, + self._solver.dofs_info, + self._solver.joints_info, + self._solver.equalities_info, + self.constraint_state, + self._collider._collider_state, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self.contact_island.contact_island_state, + ) + + def resolve(self, entities_info=None, rigid_global_info=None): + """Resolve constraints with the ComFree analytical update, then publish contact forces.""" + from ..constraint.solver import func_update_contact_force + + _comfree_resolve( + self._comfree_stiffness, + self._comfree_damping, + self._solver.dofs_state, + self._solver.entities_info, + self.constraint_state, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._solver._errno, + ) + + func_update_contact_force( + self._solver.links_state, + self._collider._collider_state, + self.constraint_state, + self._solver._static_rigid_sim_config, + ) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- + + +@qd.kernel(fastcache=True) +def _comfree_solver_kernel_reset( + envs_idx: qd.types.ndarray(), + constraint_state: array_class.ConstraintState, + static_rigid_sim_config: qd.template(), +): + n_dofs = constraint_state.qacc_ws.shape[0] + + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b_ in range(envs_idx.shape[0]): + i_b = envs_idx[i_b_] + constraint_state.is_warmstart[i_b] = False + for i_d in range(n_dofs): + constraint_state.qacc_ws[i_d, i_b] = 0.0 + + +@qd.kernel(fastcache=True) +def _comfree_solver_kernel_clear( + envs_idx: qd.types.ndarray(), + constraint_state: array_class.ConstraintState, + rigid_global_info: array_class.RigidGlobalInfo, + static_rigid_sim_config: qd.template(), +): + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b_ in range(envs_idx.shape[0]): + i_b = envs_idx[i_b_] + constraint_state.n_constraints[i_b] = 0 + constraint_state.n_constraints_equality[i_b] = 0 + constraint_state.n_constraints_frictionloss[i_b] = 0 + constraint_state.qd_n_equalities[i_b] = rigid_global_info.n_equalities[None] + + +@qd.func +def _comfree_efc_force_body( + i_c, + i_b, + n_dofs, + substep_dt, + stiffness, + damping, + dofs_state: array_class.DofsState, + constraint_state: array_class.ConstraintState, +): + """Analytical one-sided force for a single constraint row (reference forward.py:67-85). + + efc_vel = J_c @ (v + acc_smooth * dt) + force = max(efc_mass * (-d*efc_vel - k*(efc_vel*dt + efc_dist)), 0) + """ + efc_vel = gs.qd_float(0.0) + for i_d in range(n_dofs): + v_pred = dofs_state.vel[i_d, i_b] + dofs_state.acc_smooth[i_d, i_b] * substep_dt + efc_vel += constraint_state.jac[i_c, i_d, i_b] * v_pred + + efc_penetration = efc_vel * substep_dt + constraint_state.efc_dist[i_c, i_b] + efc_acc = -damping * efc_vel - stiffness * efc_penetration + constraint_state.efc_force[i_c, i_b] = qd.max(constraint_state.efc_mass[i_c, i_b] * efc_acc, 0.0) + + +@qd.func +def _comfree_qfrc_body( + i_d, + i_b, + constraint_state: array_class.ConstraintState, +): + """Accumulate the generalized constraint force for one DoF: qfrc_constraint[i_d] = (J^T @ force)[i_d].""" + qfrc = gs.qd_float(0.0) + for i_c in range(constraint_state.n_constraints[i_b]): + qfrc += constraint_state.jac[i_c, i_d, i_b] * constraint_state.efc_force[i_c, i_b] + constraint_state.qfrc_constraint[i_d, i_b] = qfrc + + +@qd.kernel(fastcache=True) +def _comfree_resolve( + comfree_stiffness: qd.template(), + comfree_damping: qd.template(), + dofs_state: array_class.DofsState, + entities_info: array_class.EntitiesInfo, + constraint_state: array_class.ConstraintState, + rigid_global_info: array_class.RigidGlobalInfo, + static_rigid_sim_config: qd.template(), + errno: qd.Tensor, +): + """Single-pass analytical constraint resolution (reference forward.py:43-97 + solve_m). + + Unlike Newton/CG this needs no iteration, so the whole resolution is a fixed sequence of + fully data-parallel passes (rather than one serial thread per env, which leaves the GPU idle): + + 1. ``efc_force`` : per-constraint analytical force, parallel over (constraint, env). + 2. ``qfrc_constraint`` : J^T @ efc_force, parallel over (dof, env). + 3. ``force`` : qf_smooth + qfrc_constraint, parallel over (dof, env). + 4. ``qacc`` : M^{-1} @ force via the pre-factored mass matrix, parallel over (entity, env). + 5. publish + warmstart : parallel over (dof, env) / (env,). + + The no-constraint case needs no special handling: with zero constraint force step 4 reduces to + ``qacc = M^{-1} qf_smooth = acc_smooth`` (the same ``func_solve_mass`` that produced acc_smooth). + """ + n_dofs = dofs_state.acc.shape[0] + _B = dofs_state.acc.shape[1] + len_constraints = constraint_state.efc_force.shape[0] + substep_dt = rigid_global_info.substep_dt[None] + + # Stiffness / damping scaled by 1/dt (reference forward.py:72-73). + stiffness = qd.static(comfree_stiffness) / substep_dt + damping = qd.static(comfree_damping) / substep_dt + + # --- 1. Per-constraint analytical force, parallel over (constraint, env). --- + # The jac flip makes i_c the stride-1 axis under the transposed layout; pick the ndrange order so adjacent + # lanes vary along the coalesced axis (i_c when flipped, i_b otherwise), mirroring _initialize_Jaref_parallel. + if qd.static(static_rigid_sim_config.constraint_layout_transposed): + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b, i_c in qd.ndrange(_B, len_constraints): + if i_c < constraint_state.n_constraints[i_b]: + _comfree_efc_force_body(i_c, i_b, n_dofs, substep_dt, stiffness, damping, dofs_state, constraint_state) + else: + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_c, i_b in qd.ndrange(len_constraints, _B): + if i_c < constraint_state.n_constraints[i_b]: + _comfree_efc_force_body(i_c, i_b, n_dofs, substep_dt, stiffness, damping, dofs_state, constraint_state) + + # --- 2. qfrc_constraint = J^T @ efc_force, parallel over (dof, env). --- + # qfrc_constraint is a dof-vec tensor (i_d stride-1 when flipped); keep i_d innermost under the transposed layout. + if qd.static(static_rigid_sim_config.constraint_layout_transposed): + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b, i_d in qd.ndrange(_B, n_dofs): + _comfree_qfrc_body(i_d, i_b, constraint_state) + else: + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_d, i_b in qd.ndrange(n_dofs, _B): + _comfree_qfrc_body(i_d, i_b, constraint_state) + + # --- 3. qfrc_total = qf_smooth + qfrc_constraint, parallel over (dof, env). --- + if qd.static(static_rigid_sim_config.constraint_layout_transposed): + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b, i_d in qd.ndrange(_B, n_dofs): + dofs_state.force[i_d, i_b] = dofs_state.qf_smooth[i_d, i_b] + constraint_state.qfrc_constraint[i_d, i_b] + else: + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_d, i_b in qd.ndrange(n_dofs, _B): + dofs_state.force[i_d, i_b] = dofs_state.qf_smooth[i_d, i_b] + constraint_state.qfrc_constraint[i_d, i_b] + + # --- 4. qacc = M^{-1} @ qfrc_total, parallel over (entity, env). --- + func_solve_mass( + vec=dofs_state.force, + out=constraint_state.qacc, + out_bw=None, + entities_info=entities_info, + rigid_global_info=rigid_global_info, + static_rigid_sim_config=static_rigid_sim_config, + is_backward=False, + ) + + # --- 5. Publish results back to dofs_state (mirrors func_update_qacc), parallel over (dof, env). --- + if qd.static(static_rigid_sim_config.constraint_layout_transposed): + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b, i_d in qd.ndrange(_B, n_dofs): + dofs_state.acc[i_d, i_b] = constraint_state.qacc[i_d, i_b] + dofs_state.qf_constraint[i_d, i_b] = constraint_state.qfrc_constraint[i_d, i_b] + constraint_state.qacc_ws[i_d, i_b] = constraint_state.qacc[i_d, i_b] + if qd.math.isnan(constraint_state.qacc[i_d, i_b]): + errno[i_b] = errno[i_b] | array_class.ErrorCode.INVALID_FORCE_NAN + else: + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_d, i_b in qd.ndrange(n_dofs, _B): + dofs_state.acc[i_d, i_b] = constraint_state.qacc[i_d, i_b] + dofs_state.qf_constraint[i_d, i_b] = constraint_state.qfrc_constraint[i_d, i_b] + constraint_state.qacc_ws[i_d, i_b] = constraint_state.qacc[i_d, i_b] + if qd.math.isnan(constraint_state.qacc[i_d, i_b]): + errno[i_b] = errno[i_b] | array_class.ErrorCode.INVALID_FORCE_NAN + + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b in range(_B): + constraint_state.is_warmstart[i_b] = True diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index 0283f7c206..e9a8e72842 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/genesis/engine/solvers/rigid/constraint/solver.py @@ -701,6 +701,13 @@ def _add_friction_constraint( 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 + # ComFree contact scalars: signed distance (-penetration) and effective mass from the plain + # invweight with a hardcoded-width impedance (no friction-cone scaling). + # Reference: comfree_warp/comfree_core/_src/constraint.py:124-125 + constraint_state.efc_dist[n_con, i_b] = -contact_data_penetration + constraint_state.efc_mass[n_con, i_b] = gu.comfree_efc_mass( + contact_data_sol_params, -contact_data_penetration, invweight, EPS + ) @qd.func @@ -848,6 +855,13 @@ def _add_collision_constraints_per_contact( 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 + # ComFree contact scalars: signed distance (-penetration) and effective mass from the + # plain invweight with a hardcoded-width impedance (no friction-cone scaling). + # Reference: comfree_warp/comfree_core/_src/constraint.py:124-125 + constraint_state.efc_dist[n_con, i_b] = -contact_data_penetration + constraint_state.efc_mass[n_con, i_b] = gu.comfree_efc_mass( + contact_data_sol_params, -contact_data_penetration, invweight, EPS + ) @qd.func @@ -1001,6 +1015,9 @@ def func_equality_connect( 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.0 / diag + # ComFree: signed constraint position is the per-axis anchor error; mass uses hardcoded-width impedance. + constraint_state.efc_dist[n_con, i_b] = pos_diff[i_3] + constraint_state.efc_mass[n_con, i_b] = gu.comfree_efc_mass(sol_params, pos_diff[i_3], invweight, EPS) @qd.func @@ -1081,6 +1098,9 @@ def func_equality_joint( 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.0 / diag + # ComFree: signed constraint position is the joint coupling residual. + constraint_state.efc_dist[n_con, i_b] = pos + constraint_state.efc_mass[n_con, i_b] = gu.comfree_efc_mass(sol_params, pos, invweight, EPS) # Populate jac_relevant_dofs for this joint-equality constraint. # Without this, sparse iterations see 0 relevant DOFs and produce @@ -1335,6 +1355,9 @@ def func_equality_weld( 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.0 / diag + # ComFree: signed constraint position is the per-axis weld translation error. + constraint_state.efc_dist[n_con, i_b] = pos_error[i] + constraint_state.efc_mass[n_con, i_b] = gu.comfree_efc_mass(sol_params, pos_error[i], invweight[0], EPS) # --- Orientation part (next 3 constraints) --- n_con = qd.atomic_add(constraint_state.n_constraints[i_b], 3) @@ -1392,6 +1415,11 @@ def func_equality_weld( 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 + # ComFree: signed constraint position is the per-axis weld rotation error. + constraint_state.efc_dist[i_con, i_b] = rot_error[i_con - n_con] + constraint_state.efc_mass[i_con, i_b] = gu.comfree_efc_mass( + sol_params, rot_error[i_con - n_con], invweight[1], EPS + ) @qd.func @@ -1439,6 +1467,11 @@ def add_joint_limit_constraints( 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 + # ComFree: signed constraint position is the (negative) limit violation. + constraint_state.efc_dist[n_con, i_b] = pos_delta + constraint_state.efc_mass[n_con, i_b] = gu.comfree_efc_mass( + joints_info.sol_params[I_j], pos_delta, dofs_info.invweight[I_d], EPS + ) if qd.static(static_rigid_sim_config.sparse_solve): for i_d2_ in range(constraint_state.jac_n_relevant_dofs[n_con, i_b]): @@ -1501,6 +1534,11 @@ def add_frictionloss_constraints( 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 + # ComFree: frictionloss has no positional violation (pos_imp == 0). + constraint_state.efc_dist[i_con, i_b] = 0.0 + constraint_state.efc_mass[i_con, i_b] = gu.comfree_efc_mass( + joints_info.sol_params[I_j], 0.0, dofs_info.invweight[I_d], EPS + ) constraint_state.efc_frictionloss[i_con, i_b] = dofs_info.frictionloss[I_d] for i_d2 in range(n_dofs): constraint_state.jac[i_con, i_d2, i_b] = gs.qd_float(0.0) diff --git a/genesis/engine/solvers/rigid/rigid_solver.py b/genesis/engine/solvers/rigid/rigid_solver.py index cf61c94ef2..93f4a61df9 100644 --- a/genesis/engine/solvers/rigid/rigid_solver.py +++ b/genesis/engine/solvers/rigid/rigid_solver.py @@ -985,7 +985,15 @@ def _init_collider(self): self.terrain_xyz_maxmin.from_numpy(xyz_maxmin) def _init_constraint_solver(self): - if self._use_contact_island: + if self._options.constraint_solver == gs.constraint_solver.ComFree: + from .comfree import ComFreeSolver + + self.constraint_solver = ComFreeSolver( + self, + comfree_stiffness=self._options.comfree_stiffness, + comfree_damping=self._options.comfree_damping, + ) + elif self._use_contact_island: self.constraint_solver = ConstraintSolverIsland(self) else: self.constraint_solver = ConstraintSolver(self) diff --git a/genesis/options/solvers.py b/genesis/options/solvers.py index 67ebdf29f1..c36f753d06 100644 --- a/genesis/options/solvers.py +++ b/genesis/options/solvers.py @@ -430,8 +430,15 @@ class RigidOptions(Options): Maximum number of IK targets. Increasing this doesn't affect IK solving speed, but will increase memory usage. Defaults to 6. constraint_solver : gs.constraint_solver, optional - Constraint solver type. Current supported constraint solvers are 'gs.constraint_solver.CG' (conjugate gradient) - and 'gs.constraint_solver.Newton' (Newton's method). Defaults to 'Newton'. + Constraint solver type. Supported solvers are 'gs.constraint_solver.CG' (conjugate gradient), + 'gs.constraint_solver.Newton' (Newton's method), and 'gs.constraint_solver.ComFree' (complementarity-free + analytical solver, arXiv:2603.12185). Defaults to 'Newton'. + comfree_stiffness : float, optional + Global stiffness parameter (k_user) for the ComFree solver. Only used when + `constraint_solver=gs.constraint_solver.ComFree`. Defaults to 0.2. + comfree_damping : float, optional + Global damping parameter (d_user) for the ComFree solver. Only used when + `constraint_solver=gs.constraint_solver.ComFree`. Defaults to 0.001. iterations : int, optional Number of iterations for the constraint solver. Defaults to 50. tolerance : float, optional @@ -510,6 +517,9 @@ class RigidOptions(Options): contact_pruning_tolerance: PositiveFloat | None = 0.02 sparse_solve: StrictBool = False constraint_timeconst: PositiveFloat = 0.01 + # ComFree analytical solver parameters (only used when constraint_solver == ComFree) + comfree_stiffness: PositiveFloat = 0.2 + comfree_damping: PositiveFloat = 0.001 use_contact_island: StrictBool = False box_box_detection: StrictBool = False diff --git a/genesis/utils/array_class.py b/genesis/utils/array_class.py index 5baad880da..09573dda8c 100644 --- a/genesis/utils/array_class.py +++ b/genesis/utils/array_class.py @@ -253,6 +253,10 @@ class ConstraintState: MinvJT: qd.Tensor search: qd.Tensor efc_D: qd.Tensor + # ComFree-only constraint scalars: signed constraint position (efc_dist) and constraint-space + # effective mass (efc_mass). Allocated for all solvers but only populated/read by ComFreeSolver. + efc_dist: qd.Tensor + efc_mass: qd.Tensor efc_frictionloss: qd.Tensor efc_force: qd.Tensor efc_b: qd.Tensor @@ -412,6 +416,8 @@ def get_constraint_state(constraint_solver, solver): efc_frictionloss=V(dtype=gs.qd_float, shape=(len_constraints_, _B), layout=con_layout), efc_force=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), efc_D=V(dtype=gs.qd_float, shape=(len_constraints_, _B), layout=con_layout), + efc_dist=V(dtype=gs.qd_float, shape=(len_constraints_, _B), layout=con_layout), + efc_mass=V(dtype=gs.qd_float, shape=(len_constraints_, _B), layout=con_layout), jv=V(dtype=gs.qd_float, shape=(len_constraints_, _B), layout=con_layout), jac=V(dtype=gs.qd_float, shape=jac_shape, layout=jac_layout), jac_relevant_dofs=V( diff --git a/genesis/utils/geom.py b/genesis/utils/geom.py index 1c6d3d71bc..d9936d9b11 100644 --- a/genesis/utils/geom.py +++ b/genesis/utils/geom.py @@ -424,6 +424,57 @@ def imp_aref(params, neg_penetration, vel, pos): return imp, aref +@qd.func +def comfree_imp(params, pos_imp): + """Compute the constraint impedance used for the ComFree effective mass (``efc_mass``). + + This mirrors the standard MuJoCo impedance curve (see ``imp_aref``) but pins the impedance + ``width`` to a constant 0.01 instead of using ``solimp[2]``. The reference ComFree + implementation hardcodes the width so that ``efc_mass`` does not jump abruptly when the + constraint violation ``pos_imp`` is large (e.g. deep initial penetration). + + Reference: comfree_warp/comfree_core/_src/constraint.py:77-111 (``_efc_row``). + """ + timeconst, dampratio, dmin, dmax, width, mid, power = params + + MJ_MINIMP = 0.0001 + MJ_MAXIMP = 0.9999 + + dmin = qd.math.clamp(dmin, MJ_MINIMP, MJ_MAXIMP) + dmax = qd.math.clamp(dmax, MJ_MINIMP, MJ_MAXIMP) + width = 0.01 # hardcoded, matching reference constraint.py:93 + mid = qd.math.clamp(mid, MJ_MINIMP, MJ_MAXIMP) + power = qd.max(1.0, power) + + imp_x = qd.abs(pos_imp) / width + imp_a = (1.0 / mid ** (power - 1.0)) * imp_x**power + imp_b = 1.0 - (1.0 / (1.0 - mid) ** (power - 1.0)) * (1.0 - imp_x) ** power + imp_y = imp_a if imp_x < mid else imp_b + + imp = dmin + imp_y * (dmax - dmin) + imp = qd.math.clamp(imp, dmin, dmax) + imp = dmax if imp_x > 1.0 else imp + + return imp + + +@qd.func +def comfree_efc_mass(params, pos_imp, invweight, eps): + """Constraint-space effective mass used by the ComFree solver (``efc_mass``). + + Matches the single ``efc_mass`` formula shared by every constraint type in the reference + (``comfree_warp/comfree_core/_src/constraint.py:125``):: + + efc_mass = 1 / max(invweight * (1 - imp) / imp, eps) + + where ``imp`` comes from :func:`comfree_imp` (hardcoded width). Unlike the standard contact + ``diag`` this uses the plain ``invweight`` (no friction-cone scaling), exactly as the reference. + """ + imp = comfree_imp(params, pos_imp) + diag = qd.max(invweight * (1.0 - imp) / imp, eps) + return 1.0 / diag + + # ------------------------------------------------------------------------------------ # -------------------------------- torch and numpy ----------------------------------- # ------------------------------------------------------------------------------------ diff --git a/tests/test_comfree.py b/tests/test_comfree.py new file mode 100644 index 0000000000..7cdea7ce67 --- /dev/null +++ b/tests/test_comfree.py @@ -0,0 +1,143 @@ +"""Tests for the ComFree (Complementarity-Free) constraint solver. + +Validates that the ComFree solver: +1. Can be instantiated and stepped without errors +2. Produces physically plausible results (objects fall, contacts prevent penetration) +3. Reproduces analytic free-fall when there are no contacts +4. Works across multiple parallel environments and simple stacking + +These tests rely on the session-scoped ``initialize_genesis`` autouse fixture +(see ``tests/conftest.py``) for ``gs.init`` / ``gs.destroy`` and the parametrized +backend / precision, so they must not initialize Genesis themselves. +""" + +import numpy as np +import pytest + +import genesis as gs + + +def _make_scene(show_viewer, dt=0.002, stiffness=0.3, damping=0.005, n_envs=0, enable_collision=True): + scene = gs.Scene( + sim_options=gs.options.SimOptions(dt=dt, gravity=(0, 0, -9.81)), + rigid_options=gs.options.RigidOptions( + dt=dt, + constraint_solver=gs.constraint_solver.ComFree, + comfree_stiffness=stiffness, + comfree_damping=damping, + enable_collision=enable_collision, + enable_joint_limit=True, + ), + show_viewer=show_viewer, + ) + scene.add_entity(gs.morphs.Plane()) + box = scene.add_entity( + gs.morphs.Box(size=(0.2, 0.2, 0.2), pos=(0, 0, 0.5)), + material=gs.materials.Rigid(rho=1000), + ) + scene.build(n_envs=n_envs) + return scene, box + + +@pytest.mark.required +def test_comfree_instantiation(show_viewer): + scene, _ = _make_scene(show_viewer) + assert scene.sim.rigid_solver.constraint_solver.__class__.__name__ == "ComFreeSolver" + + +@pytest.mark.required +def test_comfree_step(show_viewer): + scene, _ = _make_scene(show_viewer) + for _ in range(10): + scene.step() + + +def test_box_falls_under_gravity(show_viewer): + scene, box = _make_scene(show_viewer) + initial_z = float(box.get_pos().numpy()[2]) + for _ in range(5): + scene.step() + z = float(box.get_pos().numpy()[2]) + assert z < initial_z, f"Box z={z} should be less than initial z={initial_z}" + + +def test_box_contacts_floor(show_viewer): + scene, box = _make_scene(show_viewer) + for _ in range(500): + scene.step() + z = float(box.get_pos().numpy()[2]) + assert z > -0.05, f"Box fell through floor: z={z}" + assert z < 0.2, f"Box didn't settle: z={z}" + + +def test_freefall_matches_analytic(show_viewer): + """Without contacts ComFree must integrate gravity exactly like the smooth dynamics.""" + dt = 0.01 + z0 = 2.0 + n_steps = 50 + scene = gs.Scene( + sim_options=gs.options.SimOptions(dt=dt, gravity=(0, 0, -9.81)), + rigid_options=gs.options.RigidOptions( + dt=dt, + constraint_solver=gs.constraint_solver.ComFree, + enable_collision=False, + ), + show_viewer=show_viewer, + ) + box = scene.add_entity( + gs.morphs.Box(size=(0.2, 0.2, 0.2), pos=(0, 0, z0)), + material=gs.materials.Rigid(rho=1000), + ) + scene.build() + for _ in range(n_steps): + scene.step() + z = float(box.get_pos().numpy()[2]) + t = n_steps * dt + expected_z = z0 - 0.5 * 9.81 * t * t + assert abs(z - expected_z) < 0.15, f"Free-fall z={z}, expected ~{expected_z}" + + +def test_batched_simulation(show_viewer): + """ComFree works with multiple parallel environments.""" + n_envs = 4 + scene, box = _make_scene(show_viewer, n_envs=n_envs) + for _ in range(300): + scene.step() + pos = box.get_pos().numpy() + assert pos.shape == (n_envs, 3) + for i in range(n_envs): + assert pos[i, 2] > -0.05, f"Env {i}: z={pos[i, 2]}" + assert pos[i, 2] < 0.2, f"Env {i}: z={pos[i, 2]}" + + +def test_two_box_stack(show_viewer): + """Two boxes should stack on top of each other without penetrating the floor.""" + scene = gs.Scene( + sim_options=gs.options.SimOptions(dt=0.002, gravity=(0, 0, -9.81)), + rigid_options=gs.options.RigidOptions( + dt=0.002, + constraint_solver=gs.constraint_solver.ComFree, + comfree_stiffness=0.3, + comfree_damping=0.005, + enable_collision=True, + ), + show_viewer=show_viewer, + ) + scene.add_entity(gs.morphs.Plane()) + box1 = scene.add_entity( + gs.morphs.Box(size=(0.2, 0.2, 0.2), pos=(0, 0, 0.3)), + material=gs.materials.Rigid(rho=1000), + ) + box2 = scene.add_entity( + gs.morphs.Box(size=(0.2, 0.2, 0.2), pos=(0, 0, 0.7)), + material=gs.materials.Rigid(rho=1000), + ) + scene.build() + for _ in range(800): + scene.step() + + z1 = float(box1.get_pos().numpy()[2]) + z2 = float(box2.get_pos().numpy()[2]) + assert z1 > 0.0, f"Box1 penetrated: z={z1}" + assert z2 > z1, f"Box2 below box1: z2={z2}, z1={z1}" + assert z2 < 0.5, f"Box2 too high: z={z2}"