diff --git a/examples/rigid/comfree/franka_cube_grasp.py b/examples/rigid/comfree/franka_cube_grasp.py new file mode 100644 index 0000000000..f78b7496cf --- /dev/null +++ b/examples/rigid/comfree/franka_cube_grasp.py @@ -0,0 +1,244 @@ +"""Franka Panda cube grasping demo using the ComFree constraint solver. + +Replicates the grasp-robustness benchmark from ComFree-Sim (arXiv:2603.12185) +using the Genesis ComFree solver. The Franka arm approaches a cube on a pedestal, +grasps it, lifts it, and optionally applies perturbation forces. + +Usage: + uv run python examples/rigid/comfree/franka_cube_grasp.py + uv run python examples/rigid/comfree/franka_cube_grasp.py -v + uv run python examples/rigid/comfree/franka_cube_grasp.py --engine newton + +Reference: references/comfree_warp/test_local/test_franka_grasp.py +Genesis equivalent of: examples/rigid/franka_cube.py +""" + +import argparse +import math +import time + +import numpy as np + +import genesis as gs + + +# ── Constants ────────────────────────────────────────────────────────────────── +ARM_DOF = 7 +CUBE_SIZE = 0.04 # full extent of the cube (Genesis Box uses full size) +CUBE_HALF = CUBE_SIZE / 2 +PEDESTAL_TOP = 0.12 # top surface of the pedestal +LIFTED_THRESHOLD = 0.18 + + +def smoothstep(alpha: float) -> float: + alpha = min(1.0, max(0.0, alpha)) + return alpha * alpha * (3.0 - 2.0 * alpha) + + +def build_scene(engine="comfree", vis=False, dt=0.002, stiffness=0.2, damping=0.001, num_envs=1): + """Create the Franka + cube scene (matching examples/rigid/franka_cube.py).""" + + gs.init(backend=gs.cpu, precision="32", logging_level="info", performance_mode=True) + + if engine == "comfree": + constraint_solver = gs.constraint_solver.ComFree + else: + constraint_solver = gs.constraint_solver.Newton + + scene = gs.Scene( + viewer_options=gs.options.ViewerOptions( + camera_pos=(3, -1, 1.5), + camera_lookat=(0.0, 0.0, 0.5), + camera_fov=30, + res=(960, 640), + max_FPS=60, + ), + sim_options=gs.options.SimOptions(dt=dt), + rigid_options=gs.options.RigidOptions( + dt=dt, + constraint_solver=constraint_solver, + comfree_stiffness=stiffness, + comfree_damping=damping, + # box_box_detection not required - MPR/GJK handles box-box pairs + ), + show_viewer=vis, + ) + + scene.add_entity(gs.morphs.Plane()) + franka = scene.add_entity(gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml")) + cube = scene.add_entity( + gs.morphs.Box(size=(CUBE_SIZE, CUBE_SIZE, CUBE_SIZE), pos=(0.65, 0.0, CUBE_HALF)), + ) + scene.build(n_envs=num_envs) + return scene, franka, cube + + +def run_grasp_trial( + engine="comfree", + vis=False, + dt=0.002, + stiffness=0.2, + damping=0.001, + steps_scale=1.0, + perturb=True, + num_envs=1, +): + """Run the full grasp trial following examples/rigid/franka_cube.py style. + + Phases (matching reference test_franka_grasp.py): + 1. approach - IK to position above cube, hold + 2. grasp - close fingers + 3. lift - IK to raised position, hold while grasping + 4. hold - hold lifted position + 5. perturb - apply sinusoidal forces to test grasp robustness (optional) + """ + scene, franka, cube = build_scene( + engine=engine, vis=vis, dt=dt, stiffness=stiffness, damping=damping, num_envs=num_envs + ) + + motors_dof = np.arange(ARM_DOF) + fingers_dof = np.arange(ARM_DOF, ARM_DOF + 2) + + # ── Initial configuration (same as franka_cube.py) ──────────────────────── + qpos = np.array([-1.0124, 1.5559, 1.3662, -1.6878, -1.5799, 1.7757, 1.4602, 0.04, 0.04]) + franka.set_qpos(qpos) + scene.step() + + end_effector = franka.get_link("hand") + + # ── Compute IK targets ──────────────────────────────────────────────────── + # When n_envs >= 1, IK expects pos shape (n_envs, 3) and quat shape (n_envs, 4) + approach_pos = np.tile([0.65, 0.0, 0.135], (num_envs, 1)) + lift_pos = np.tile([0.65, 0.0, 0.4], (num_envs, 1)) + ik_quat = np.tile([0, 1, 0, 0], (num_envs, 1)) + + # Approach: above cube (same as franka_cube.py) + approach_q = franka.inverse_kinematics( + link=end_effector, + pos=approach_pos, + quat=ik_quat, + ).numpy() + + # Lift: raised position (higher target for ComFree's softer contacts) + lift_q = franka.inverse_kinematics( + link=end_effector, + pos=lift_pos, + quat=ik_quat, + ).numpy() + + # ── Phase definitions ───────────────────────────────────────────────────── + # (name, steps, arm_qpos, finger_pos, force_scale) + open_fingers = 0.04 + closed_fingers = 0.0 # fully closed in Genesis finger control + + # Phase step counts are calibrated for dt=0.002; scale proportionally for other dt values + dt_scale = 0.002 / dt + approach_arm = approach_q[:, :ARM_DOF] + lift_arm = lift_q[:, :ARM_DOF] + + phases = [ + ("approach", int(500 * steps_scale * dt_scale), approach_arm, open_fingers, 0.0), + ("grasp", int(500 * steps_scale * dt_scale), approach_arm, closed_fingers, 0.0), + ("lift", int(1000 * steps_scale * dt_scale), lift_arm, closed_fingers, 0.0), + ("hold", int(1000 * steps_scale * dt_scale), lift_arm, closed_fingers, 0.0), + ] + if perturb: + phases.append(("perturb", int(2500 * steps_scale * dt_scale), lift_arm, closed_fingers, 1.5)) + + # ── Simulation loop ─────────────────────────────────────────────────────── + min_cube_z = float("inf") + max_cube_z = float("-inf") + total_steps = 0 + wall_time_start = time.perf_counter() + + for phase_name, phase_steps, arm_target, finger_target, force_scale in phases: + phase_steps = max(1, phase_steps) + for local_step in range(phase_steps): + franka.control_dofs_position(arm_target, motors_dof) + franka.control_dofs_position(np.array([finger_target, finger_target]), fingers_dof) + + # Apply perturbation forces to hand link + if force_scale > 0.0: + t = local_step * dt + fx = force_scale * math.sin(14.0 * t) + fy = 0.75 * force_scale * math.cos(9.0 * t) + hand_link = franka.get_link("hand") + scene.sim.rigid_solver.apply_links_external_force( + force=np.array([[fx, fy, 0.0]]), + links_idx=[hand_link.idx_local], + ) + + scene.step() + + cube_z = cube.get_pos().numpy().flatten()[2] + min_cube_z = min(min_cube_z, cube_z) + max_cube_z = max(max_cube_z, cube_z) + total_steps += 1 + + if total_steps % 100 == 0: + print(f" [{phase_name}] step {local_step}/{phase_steps}, cube_z={cube_z:.4f}") + + wall_time = time.perf_counter() - wall_time_start + final_cube_z = cube.get_pos().numpy().flatten()[2] + success = final_cube_z > LIFTED_THRESHOLD + + results = { + "engine": engine, + "total_steps": total_steps, + "final_cube_z": final_cube_z, + "min_cube_z": min_cube_z, + "max_cube_z": max_cube_z, + "success": success, + "wall_time": wall_time, + "steps_per_sec": total_steps / wall_time if wall_time > 0 else 0, + "num_envs": num_envs, + } + return results + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("-v", "--vis", action="store_true", default=False, help="Launch viewer.") + parser.add_argument("--engine", choices=("comfree", "newton"), default="comfree", help="Constraint solver backend.") + parser.add_argument( + "--dt", type=float, default=0.002, help="Simulation timestep (ComFree benefits from smaller dt)." + ) + parser.add_argument("--stiffness", type=float, default=0.2, help="ComFree stiffness (k_user).") + parser.add_argument("--damping", type=float, default=0.001, help="ComFree damping (d_user).") + parser.add_argument("--steps-scale", type=float, default=1.0, help="Scale phase durations.") + parser.add_argument("--no-perturb", action="store_true", help="Skip perturbation phase.") + parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to run in parallel.") + args = parser.parse_args() + + results = run_grasp_trial( + engine=args.engine, + vis=args.vis, + dt=args.dt, + stiffness=args.stiffness, + damping=args.damping, + steps_scale=args.steps_scale, + perturb=not args.no_perturb, + num_envs=args.num_envs, + ) + + print("\n" + "=" * 60) + print(f"Engine: {results['engine']}") + print(f"Number of environments: {results['num_envs']}") + print(f"Total steps: {results['total_steps']}") + print(f"Final cube z: {results['final_cube_z']:.4f}") + print(f"Min cube z: {results['min_cube_z']:.4f}") + print(f"Max cube z: {results['max_cube_z']:.4f}") + print(f"Wall time: {results['wall_time']:.2f}s") + print(f"Steps/sec: {results['steps_per_sec']:.1f}") + print(f"Success: {results['success']}") + print("=" * 60) + + if not results["success"]: + print("\nGrasp trial FAILED - cube was not lifted.") + raise SystemExit(1) + else: + print("\nGrasp trial PASSED - cube successfully lifted!") + + +if __name__ == "__main__": + main() diff --git a/genesis/constants.py b/genesis/constants.py index dafc204dee..ce62f2d6e5 100644 --- a/genesis/constants.py +++ b/genesis/constants.py @@ -60,6 +60,7 @@ class integrator(IntEnum): class constraint_solver(IntEnum): CG = 0 Newton = 1 + ComFree = 2 # backend diff --git a/genesis/engine/solvers/rigid/comfree/__init__.py b/genesis/engine/solvers/rigid/comfree/__init__.py new file mode 100644 index 0000000000..620cf2bf3b --- /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 + (Borse et al., 2026, arXiv:2603.12185) + +Instead of iterative complementarity-based solving (Newton/CG), +ComFree computes constraint forces in closed form via an impedance-style +prediction-correction update in the dual cone of Coulomb friction. +""" + +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..2865bb452e --- /dev/null +++ b/genesis/engine/solvers/rigid/comfree/solver.py @@ -0,0 +1,348 @@ +""" +ComFree (Complementarity-Free) constraint solver. + +Replaces iterative complementarity-based constraint solving with a single-pass +analytical contact force computation. Uses the existing Genesis constraint +assembly infrastructure (Jacobians, collision detection) but resolves forces +via an impedance-style prediction-correction update. + +Reference implementation: references/comfree_warp/comfree_warp/comfree_core/_src/forward.py +""" + +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.forward_dynamics import func_solve_mass_batch +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. + + Instead of iteratively solving a complementarity problem (Newton/CG), + ComFree computes constraint forces in closed form: + + 1. Predict smooth velocity: v_smooth = v + acc_smooth * dt + 2. Compute constraint velocity: efc_vel = J @ v_smooth + 3. Compute penetration: efc_pen = efc_vel * dt + efc_dist + 4. Compute force: efc_force = max(efc_mass * (-d*efc_vel - k*efc_pen), 0) + 5. Accumulate: qfrc_constraint = J^T @ efc_force + 6. Solve: qacc = M^{-1} @ (qf_smooth + qfrc_constraint) + + 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 = comfree_stiffness + self._comfree_damping = comfree_damping + + # Estimate constraint count same way as the standard solver + 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) + + # Need these for compatibility with existing constraint assembly + self.sparse_solve = False + + # Allocate constraint state (reuses standard allocation) + 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) + ) + + # Convenient aliases (needed by constraint assembly functions) + cs = self.constraint_state + self.qd_n_equalities = cs.qd_n_equalities + self.jac = cs.jac + self.diag = cs.diag + self.aref = cs.aref + self.n_constraints = cs.n_constraints + self.n_constraints_equality = cs.n_constraints_equality + self.n_constraints_frictionloss = cs.n_constraints_frictionloss + self.efc_D = cs.efc_D + self.efc_force = cs.efc_force + self.active = cs.active + self.prev_active = cs.prev_active + self.qfrc_constraint = cs.qfrc_constraint + self.qacc = cs.qacc + self.qacc_ws = cs.qacc_ws + self.qacc_prev = cs.qacc_prev + self.Ma = cs.Ma + self.Ma_ws = cs.Ma_ws + self.Jaref = cs.Jaref + self.efc_frictionloss = cs.efc_frictionloss + self.jac_relevant_dofs = cs.jac_relevant_dofs + self.jac_n_relevant_dofs = cs.jac_n_relevant_dofs + self.improved = cs.improved + self.grad = cs.grad + self.Mgrad = cs.Mgrad + self.search = cs.search + self.mv = cs.mv + self.jv = cs.jv + self.cost_ws = cs.cost_ws + self.gauss = cs.gauss + self.cost = cs.cost + self.prev_cost = cs.prev_cost + self.gtol = cs.gtol + self.quad_gauss = cs.quad_gauss + self.candidates = cs.candidates + self.ls_it = cs.ls_it + self.ls_result = cs.ls_result + + self.reset() + + # ContactIsland needed for compatibility with kernel_step_2 + 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): + """For contact island mode compatibility.""" + 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): + """Resolve constraints using the ComFree analytical method.""" + from ..constraint.solver import func_update_contact_force + + _comfree_resolve( + self._comfree_stiffness, + self._comfree_damping, + self._solver.dofs_state, + self._solver.dofs_info, + 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, + ) + + +# --------------------------------------------------------------------------- +# Quadrants Kernels +# --------------------------------------------------------------------------- + + +@qd.kernel(fastcache=gs.use_fastcache) +def _comfree_solver_kernel_reset( + envs_idx: qd.types.ndarray(), + constraint_state: array_class.ConstraintState, + static_rigid_sim_config: qd.template(), +): + _B = constraint_state.is_warmstart.shape[0] + 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=gs.use_fastcache) +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(), +): + _B = constraint_state.n_constraints.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.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.kernel(fastcache=gs.use_fastcache) +def _comfree_resolve( + comfree_stiffness: qd.template(), + comfree_damping: qd.template(), + dofs_state: array_class.DofsState, + dofs_info: array_class.DofsInfo, + entities_info: array_class.EntitiesInfo, + constraint_state: array_class.ConstraintState, + rigid_global_info: array_class.RigidGlobalInfo, + static_rigid_sim_config: qd.template(), + errno: array_class.V_ANNOTATION, +): + """ComFree constraint resolution: single-pass analytical force computation. + + Implements the core ComFree algorithm (arXiv:2603.12185): + 1. Predict smooth velocity: v_smooth_pred = v + acc_smooth * dt + 2. For each constraint: compute force analytically + 3. Accumulate generalized constraint forces + 4. Solve for constrained acceleration: qacc = M^{-1} @ (qf_smooth + qfrc_constraint) + + Reference: comfree_warp/comfree_core/_src/forward.py:26-97 + """ + n_dofs = dofs_state.acc.shape[0] + _B = dofs_state.acc.shape[1] + substep_dt = rigid_global_info.substep_dt[None] + + # Stiffness and damping scaled by 1/dt (as in reference: forward.py:72-73) + stiffness = qd.static(comfree_stiffness) / substep_dt + damping = qd.static(comfree_damping) / substep_dt + + qd.loop_config(serialize=qd.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)) + for i_b in range(_B): + n_c = constraint_state.n_constraints[i_b] + + # Zero constraint forces + for i_d in range(n_dofs): + constraint_state.qfrc_constraint[i_d, i_b] = 0.0 + + if n_c == 0: + # No constraints: qacc = acc_smooth + for i_d in range(n_dofs): + constraint_state.qacc[i_d, i_b] = dofs_state.acc_smooth[i_d, i_b] + else: + # For each constraint, compute ComFree force analytically + for i_c in range(n_c): + # Compute constraint velocity: efc_vel = J @ v_smooth_pred + # where v_smooth_pred = v + acc_smooth * dt + efc_vel = float(0.0) + for i_d in range(n_dofs): + v_smooth_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_smooth_pred + + # efc_dist = signed distance/constraint violation + # efc_D = constraint impedance (includes friction scaling for contacts) + efc_dist = constraint_state.efc_dist[i_c, i_b] + efc_mass = constraint_state.efc_D[i_c, i_b] + + # Predictive penetration: efc_vel * dt + efc_dist + efc_penetration = efc_vel * substep_dt + efc_dist + + # Analytical force: max(efc_mass * (-damping * efc_vel - stiffness * efc_penetration), 0) + efc_acc = -damping * efc_vel - stiffness * efc_penetration + efc_frc = efc_mass * efc_acc + efc_frc = qd.max(efc_frc, 0.0) + + # Store force + constraint_state.efc_force[i_c, i_b] = efc_frc + + # Accumulate generalized constraint force: qfrc_constraint += J^T @ efc_force + for i_d in range(n_dofs): + constraint_state.qfrc_constraint[i_d, i_b] += constraint_state.jac[i_c, i_d, i_b] * efc_frc + + # Update total force and solve for constrained acceleration + # qfrc_total = qf_smooth + qfrc_constraint + for i_d in range(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] + + # Solve M @ qacc = qfrc_total using pre-factored mass matrix + func_solve_mass_batch( + i_b, + vec=dofs_state.force, + out=constraint_state.qacc, + out_bw=constraint_state.qacc, + entities_info=entities_info, + rigid_global_info=rigid_global_info, + static_rigid_sim_config=static_rigid_sim_config, + is_backward=False, + ) + + # Write results back to dofs_state (same as standard solver's func_update_qacc) + 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] + dofs_state.force[i_d, i_b] = dofs_state.qf_smooth[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 986a8c1d48..04e32ca0fc 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/genesis/engine/solvers/rigid/constraint/solver.py @@ -673,6 +673,12 @@ def add_collision_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 + constraint_state.efc_dist[n_con, i_b] = -contact_data_penetration + # ComFree efc_mass: plain invweight, hardcoded width=0.01 impedance + # Reference: comfree_warp/comfree_core/_src/constraint.py:93,125 + comfree_imp = gu.comfree_imp(contact_data_sol_params, -contact_data_penetration) + comfree_diag = qd.max(invweight * (1.0 - comfree_imp) / comfree_imp, EPS) + constraint_state.efc_mass[n_con, i_b] = 1.0 / comfree_diag @qd.func @@ -785,6 +791,8 @@ 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 + constraint_state.efc_dist[n_con, i_b] = pos_diff[i_3] + constraint_state.efc_mass[n_con, i_b] = 1.0 / diag @qd.func @@ -865,6 +873,8 @@ 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 + constraint_state.efc_dist[n_con, i_b] = pos + constraint_state.efc_mass[n_con, i_b] = 1.0 / diag @qd.kernel(fastcache=gs.use_fastcache) @@ -1105,6 +1115,8 @@ 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 + constraint_state.efc_dist[n_con, i_b] = pos_error[i] + constraint_state.efc_mass[n_con, i_b] = 1.0 / diag # --- Orientation part (next 3 constraints) --- n_con = qd.atomic_add(constraint_state.n_constraints[i_b], 3) @@ -1161,6 +1173,8 @@ 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 + constraint_state.efc_dist[i_con, i_b] = rot_error[i_con - n_con] + constraint_state.efc_mass[i_con, i_b] = 1.0 / diag @qd.func @@ -1206,6 +1220,8 @@ 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 + constraint_state.efc_dist[n_con, i_b] = pos_delta + constraint_state.efc_mass[n_con, i_b] = 1.0 / diag if qd.static(static_rigid_sim_config.sparse_solve): for i_d2_ in range(constraint_state.jac_n_relevant_dofs[n_con, i_b]): @@ -1267,6 +1283,8 @@ 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 + constraint_state.efc_dist[i_con, i_b] = 0.0 + constraint_state.efc_mass[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[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 f6d32eb6fe..bb1ba80c74 100644 --- a/genesis/engine/solvers/rigid/rigid_solver.py +++ b/genesis/engine/solvers/rigid/rigid_solver.py @@ -875,7 +875,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 44f95adafa..46c7c5a7b9 100644 --- a/genesis/options/solvers.py +++ b/genesis/options/solvers.py @@ -430,8 +430,9 @@ 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). Defaults to 'Newton'. iterations : int, optional Number of iterations for the constraint solver. Defaults to 50. tolerance : float, optional @@ -503,6 +504,9 @@ class RigidOptions(Options): noslip_tolerance: PositiveFloat = 1e-6 sparse_solve: StrictBool = False constraint_timeconst: PositiveFloat = 0.01 + # ComFree analytical solver parameters + 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 4c4dd7ce08..f086af88dc 100644 --- a/genesis/utils/array_class.py +++ b/genesis/utils/array_class.py @@ -250,6 +250,8 @@ class StructConstraintState(metaclass=BASE_METACLASS): MinvJT: V_ANNOTATION search: V_ANNOTATION efc_D: V_ANNOTATION + efc_dist: V_ANNOTATION + efc_mass: V_ANNOTATION efc_frictionloss: V_ANNOTATION efc_force: V_ANNOTATION efc_b: V_ANNOTATION @@ -376,6 +378,8 @@ def get_constraint_state(constraint_solver, solver): efc_frictionloss=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), efc_force=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), efc_D=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), + efc_dist=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), + efc_mass=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), jv=V(dtype=gs.qd_float, shape=(len_constraints_, _B)), jac=V(dtype=gs.qd_float, shape=jac_shape), jac_relevant_dofs=V(dtype=gs.qd_int, shape=jac_relevant_dofs_shape), diff --git a/genesis/utils/geom.py b/genesis/utils/geom.py index 15245919f2..8ccdaa0af6 100644 --- a/genesis/utils/geom.py +++ b/genesis/utils/geom.py @@ -424,6 +424,36 @@ def imp_aref(params, neg_penetration, vel, pos): return imp, aref +@qd.func +def comfree_imp(params, pos_imp): + """Compute impedance for ComFree efc_mass, using hardcoded width=0.01. + + Matches the reference: comfree_warp/comfree_core/_src/constraint.py:77-111 + The hardcoded width avoids large jumps when pos_imp is large (e.g., initial penetration). + """ + 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 as in 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 + + # ------------------------------------------------------------------------------------ # -------------------------------- torch and numpy ----------------------------------- # ------------------------------------------------------------------------------------ diff --git a/tests/test_comfree.py b/tests/test_comfree.py new file mode 100644 index 0000000000..46c5609ffc --- /dev/null +++ b/tests/test_comfree.py @@ -0,0 +1,220 @@ +"""Tests for the ComFree (Complementarity-Free) constraint solver. + +Validates that the ComFree solver: +1. Can be instantiated and run without errors +2. Produces physically plausible results (objects fall, contacts prevent penetration) +3. Handles equality constraints (weld/connect) +4. Handles joint limits +5. Produces comparable results to the standard Newton solver +""" + +import numpy as np +import genesis as gs + + +def _make_scene(solver_type, 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=solver_type, + comfree_stiffness=stiffness, + comfree_damping=damping, + enable_collision=enable_collision, + enable_joint_limit=True, + ), + show_viewer=False, + ) + 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 + + +def test_comfree_instantiation(): + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene, box = _make_scene(gs.constraint_solver.ComFree) + assert scene.sim.rigid_solver.constraint_solver.__class__.__name__ == "ComFreeSolver" + gs.destroy() + + +def test_comfree_step(): + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene, box = _make_scene(gs.constraint_solver.ComFree) + for _ in range(10): + scene.step() + gs.destroy() + + +def test_box_falls_under_gravity(): + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene, box = _make_scene(gs.constraint_solver.ComFree) + initial_z = box.get_pos().numpy()[2] + for _ in range(5): + scene.step() + z = box.get_pos().numpy()[2] + assert z < initial_z, f"Box z={z} should be less than initial z={initial_z}" + gs.destroy() + + +def test_box_contacts_floor(): + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene, box = _make_scene(gs.constraint_solver.ComFree) + for _ in range(500): + scene.step() + z = 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}" + gs.destroy() + + +def test_freefall_comparable(): + """Free-fall should be nearly identical between Newton and ComFree.""" + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene_c, box_c = _make_scene(gs.constraint_solver.ComFree, enable_collision=False) + for _ in range(20): + scene_c.step() + z_c = box_c.get_pos().numpy()[2] + gs.destroy() + + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene_n, box_n = _make_scene(gs.constraint_solver.Newton, enable_collision=False) + for _ in range(20): + scene_n.step() + z_n = box_n.get_pos().numpy()[2] + gs.destroy() + + np.testing.assert_allclose(z_c, z_n, atol=1e-3, err_msg="Free-fall trajectories should match") + + +def test_contact_qualitative(): + """ComFree and Newton should both settle near the ground.""" + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene_c, box_c = _make_scene(gs.constraint_solver.ComFree) + for _ in range(500): + scene_c.step() + z_c = box_c.get_pos().numpy()[2] + gs.destroy() + + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene_n, box_n = _make_scene(gs.constraint_solver.Newton) + for _ in range(500): + scene_n.step() + z_n = box_n.get_pos().numpy()[2] + gs.destroy() + + assert abs(z_c - z_n) < 0.1, f"ComFree z={z_c:.4f} vs Newton z={z_n:.4f}" + + +def test_no_constraint_passthrough(): + """Without contacts, ComFree should give accurate free-fall.""" + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene = gs.Scene( + sim_options=gs.options.SimOptions(dt=0.01, gravity=(0, 0, -9.81)), + rigid_options=gs.options.RigidOptions( + dt=0.01, + constraint_solver=gs.constraint_solver.ComFree, + enable_collision=False, + ), + show_viewer=False, + ) + box = scene.add_entity( + gs.morphs.Box(size=(0.2, 0.2, 0.2), pos=(0, 0, 2.0)), + material=gs.materials.Rigid(rho=1000), + ) + scene.build() + for _ in range(50): + scene.step() + z = box.get_pos().numpy()[2] + expected_z = 2.0 - 0.5 * 9.81 * 0.5**2 + assert abs(z - expected_z) < 0.15, f"Free-fall z={z}, expected ~{expected_z}" + gs.destroy() + + +def test_stiffness_effect(): + """Higher stiffness should reduce penetration.""" + results = [] + for k in [0.1, 0.3, 0.5]: + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene, box = _make_scene(gs.constraint_solver.ComFree, stiffness=k, damping=0.005) + for _ in range(500): + scene.step() + results.append(box.get_pos().numpy()[2]) + gs.destroy() + + # All should be above floor + for z in results: + assert z > 0.0, f"Box penetrated floor: z={z}" + + +def test_batched_simulation(): + """ComFree works with multiple parallel environments.""" + n_envs = 4 + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + scene, box = _make_scene(gs.constraint_solver.ComFree, 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]}" + gs.destroy() + + +def test_two_box_stack(): + """Two boxes should stack on top of each other.""" + gs.init(backend=gs.cpu, precision="64", logging_level="warning") + 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=False, + ) + 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 = box1.get_pos().numpy()[2] + z2 = 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}" + gs.destroy() + + +if __name__ == "__main__": + tests = [ + test_comfree_instantiation, + test_comfree_step, + test_box_falls_under_gravity, + test_box_contacts_floor, + test_freefall_comparable, + test_contact_qualitative, + test_no_constraint_passthrough, + test_stiffness_effect, + test_batched_simulation, + test_two_box_stack, + ] + for test in tests: + print(f"Running {test.__name__}...", end=" ") + test() + print("PASSED") + print("\nAll tests passed!")