diff --git a/CHANGELOG.md b/CHANGELOG.md index 532871d383..b725c3225a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). - Add `internal::matrix_product` in `math.hpp` - Add `internal::matrix_inversion` in `math.hpp` - Add `internal::matrix_inversion_code_generated` in `math.hpp` +- Add `examples/admm-constraint-solver.py`: how to use the constraints API to model a stack of cubes, how to use the Delassus operator, how to use the ADMM solver to solve the constraint problem. ### Changed - Clean delassus API: DelassusOperatorBase define the main delassus API and each method calls `derived().[name-of-method]Impl` diff --git a/bindings/python/algorithm/solvers/expose-admm-solver.cpp b/bindings/python/algorithm/solvers/expose-admm-solver.cpp index 6eff45ae60..16f66aa0f9 100644 --- a/bindings/python/algorithm/solvers/expose-admm-solver.cpp +++ b/bindings/python/algorithm/solvers/expose-admm-solver.cpp @@ -339,8 +339,13 @@ namespace pinocchio bp::class_ cl( "ADMMConstraintSolver", "Alternating Direction Method of Multipliers (ADMM) solver for contact dynamics.", + bp::init<>(bp::arg("self"), "Default constructor.")); + cl.def( bp::init( - bp::args("self", "problem_size"), "Constructor with problem dimension.")); + bp::args("self", "problem_size"), + "Constructor with problem dimension. Allows to pre-allocate data if problem_size is " + "known in advance. The solver will automatically resize its workspace and the result in " + "any case.")); cl // Base solver diff --git a/bindings/python/algorithm/solvers/expose-pgs-solver.cpp b/bindings/python/algorithm/solvers/expose-pgs-solver.cpp index fb11c6cef6..d95fb590c6 100644 --- a/bindings/python/algorithm/solvers/expose-pgs-solver.cpp +++ b/bindings/python/algorithm/solvers/expose-pgs-solver.cpp @@ -206,8 +206,13 @@ namespace pinocchio // Expose the solver itself bp::class_ cl( "PGSConstraintSolver", "Projected Gauss-Seidel (PGS) solver for contact dynamics.", + bp::init<>(bp::arg("self"), "Default constructor.")); + cl.def( bp::init( - bp::args("self", "problem_size"), "Constructor with problem dimension.")); + bp::args("self", "problem_size"), + "Constructor with problem dimension. Allows to pre-allocate data if problem_size is " + "known in advance. The solver will automatically resize its workspace and the result in " + "any case.")); cl // Base solver diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index a784333088..547d699d53 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -122,6 +122,7 @@ if(BUILD_PYTHON_INTERFACE) model-graph model-configuration-converter ellipsoid-joint-kinematics + admm-constraint-solver ) if(BUILD_WITH_URDF_SUPPORT) diff --git a/examples/admm-constraint-solver.py b/examples/admm-constraint-solver.py new file mode 100644 index 0000000000..ebe5f99193 --- /dev/null +++ b/examples/admm-constraint-solver.py @@ -0,0 +1,234 @@ +""" +Stack of two cubes solved with the ADMM constraint solver. + +This example demonstrates how to: + - Build a kinematic model with two free-floating cubes + - Manually define 8 PointContactConstraintModel constraints: + 4 for the floor-cube 1 interaction (bottom corners of cube 1) + 4 for the cube 1-cube 2 interaction (corners at the shared face) + - Build the Delassus operator via the Cholesky decomposition + - Compute the constraint drift vector g = J * v_free + - Solve the constrained problem with the ADMM solver + - Print convergence statistics +""" + +import sys + +import numpy as np +import pinocchio as pin + +# ─── 1. Kinematic / dynamic model ──────────────────────────────────────────── + +model = pin.Model() +# Default gravity is already (0, 0, -9.81, 0, 0, 0); shown here for clarity. +model.gravity = pin.Motion(np.array([0.0, 0.0, -9.81, 0.0, 0.0, 0.0])) + +box_size = 1.0 # edge length of each cube [m] +box_half = box_size / 2.0 +box_mass = 1.0 # mass of each cube [kg] + +box_inertia = pin.Inertia.FromBox(box_mass, box_size, box_size, box_size) + +# Cube 1 - free-flyer joint branching from the universe (joint 0) +joint1_id = model.addJoint( + 0, pin.JointModelFreeFlyer(), pin.SE3.Identity(), "box1_joint" +) +model.appendBodyToJoint(joint1_id, box_inertia, pin.SE3.Identity()) + +# Cube 2 - free-flyer joint, also branching from the universe +joint2_id = model.addJoint( + 0, pin.JointModelFreeFlyer(), pin.SE3.Identity(), "box2_joint" +) +model.appendBodyToJoint(joint2_id, box_inertia, pin.SE3.Identity()) + +# ─── 2. Initial configuration ───────────────────────────────────────────────── + +# Free-flyer q = [tx, ty, tz, qx, qy, qz, qw] +# Cube 1 centre at z = box_half (bottom face touching z = 0) +# Cube 2 centre at z = 3*box_half (bottom face touching the top of cube 1) +q_box1 = np.array([0.0, 0.0, box_half, 0.0, 0.0, 0.0, 1.0]) +q_box2 = np.array([0.0, 0.0, 3.0 * box_half, 0.0, 0.0, 0.0, 1.0]) +q0 = np.concatenate([q_box1, q_box2]) + +v0 = np.zeros(model.nv) # zero initial velocity +zero_torque = np.zeros(model.nv) # no external torques +dt = 1e-3 # time-step [s] + +# ─── 3. Build the 8 contact constraints ────────────────────────────────────── + +friction_coeff = 0.8 + + +def make_corner_constraints(model, joint1_id, joint2_id, z1, z2): + """ + Return 4 PointContactConstraintModel objects for the corners of a planar + contact interface between two bodies. + + Parameters + ---------- + joint1_id, joint2_id : joint indices of the two contacting bodies + z1 : z-coordinate of the contact plane in jid1's local frame + (e.g. +box_half for the top face, 0.0 for the world floor) + z2 : z-coordinate of the contact plane in jid2's local frame + (e.g. -box_half for the bottom face) + """ + corners_xy = [ + np.array([+box_half, +box_half]), + np.array([-box_half, +box_half]), + np.array([-box_half, -box_half]), + np.array([+box_half, -box_half]), + ] + cms = [] + for xy in corners_xy: + p1 = np.array([xy[0], xy[1], z1]) + p2 = np.array([xy[0], xy[1], z2]) + joint1_placement = pin.SE3(np.eye(3), p1) + joint2_placement = pin.SE3(np.eye(3), p2) + cm = pin.PointContactConstraintModel( + model, joint1_id, joint1_placement, joint2_id, joint2_placement + ) + cm.setFriction(friction_coeff) + cms.append(cm) + return cms + + +constraint_models = [] + +# 4 constraints: floor (universe, jid=0) ↔ cube 1 bottom face +# - in universe frame: contact points are on the floor plane z = 0 +# - in box1 local frame: contact points are at the bottom corners z = -box_half +for cm in make_corner_constraints(model, 0, joint1_id, 0.0, -box_half): + constraint_models.append(pin.ConstraintModel(cm)) + +# 4 constraints: cube 1 top face ↔ cube 2 bottom face +# - in box1 local frame: top corners at z = +box_half +# - in box2 local frame: bottom corners at z = -box_half +for cm in make_corner_constraints(model, joint1_id, joint2_id, +box_half, -box_half): + constraint_models.append(pin.ConstraintModel(cm)) + +total_residual_size = sum(cm.residualSize() for cm in constraint_models) +print(f"Number of constraints: {len(constraint_models)}") +print(f"Total constraint residual size: {total_residual_size}") + +# ─── 4. Setup and run the simulation loop ──────────────────────────────────── + +data = model.createData() +fext = [pin.Force.Zero() for _ in range(model.njoints)] + +# Initialise constraint data +constraint_datas = [] +for cmodel in constraint_models: + cdata = cmodel.createData() + constraint_datas.append(cdata) + +# Initialize constraint cholesky +chol = pin.ConstraintCholeskyDecomposition( + model, data, constraint_models, constraint_datas +) + +# Initialize constraint solver, its settings and result +solver = pin.ADMMConstraintSolver() +settings = pin.ADMMSolverSettings() +settings.max_iterations = 1000 +settings.absolute_feasibility_tol = 1e-10 +settings.relative_feasibility_tol = 1e-12 +settings.absolute_complementarity_tol = 1e-10 +settings.relative_complementarity_tol = 1e-12 +settings.admm_update_rule = pin.ADMMUpdateRule.SPECTRAL +settings.mu_prox = 1e-6 +settings.stat_record = ( + True # per-iteration statistics. Turn off for faster solves if not needed. +) +settings.solve_ncp = True +result = pin.ADMMSolverResult() + +# Simulate for a few time-steps, solving the constraint problem at each step. +horizon = 10 +q = q0.copy() +v = v0.copy() +for t in range(horizon): + # Data needs to be informed of the current state of the system for + # downstream computations. + data.q_in = q + data.v_in = v + data.tau_in = zero_torque + + # CRBA is required before building the Cholesky decomposition of the Delassus + # matrix G = J M⁻¹ Jᵀ. + # Note that other delassus operators may not necessarily require CRBA. + pin.crba(model, data, q, pin.Convention.WORLD) + + # Free acceleration: velocity the system would reach in one + # time-step without contacts. + v_free = v + dt * pin.aba(model, data, q, v, zero_torque, fext) + + # In theory, the constraint models' placements need to be updated at each time step + # of the simulation to reflect the current configuration q. + # We don't do it in this specific example since the cubes are stable and the contact + # positions don't change. + # In a more general case, you would need to do something like: + # for cmodel in constraint_models: + # cmodel.joint1_placement = ... # relative placement of the + # contact point in joint1's local frame. + # cmodel.joint2_placement = ... # relative placement of the + # contact point in joint2's local frame. + + # Run calc on constraint models + for cmodel, cdata in zip(constraint_models, constraint_datas): + cmodel.calc(model, data, cdata) + + # Cholesky decomposition of the Delassus matrix. + chol.compute(model, data, constraint_models, constraint_datas, 1e-10) + + # DelassusCholeskyExpression wraps the Cholesky factors for efficient solves. + delassus_expr = chol.getDelassusOperatorCholeskyExpression() + + # Constraint Jacobian and drift g = J v_free. + Jc = pin.getConstraintsJacobian(model, data, constraint_models, constraint_datas) + g = Jc @ v_free + + print(f"Delassus matrix size: {delassus_expr.matrix().shape}") + print(f"Drift vector ‖g‖: {np.linalg.norm(g):.4e}") + + # Solve the constraint problem with the ADMM solver. + has_converged = solver.solve( + delassus_expr, g, constraint_models, constraint_datas, settings, result + ) + + print() + print(f"time step: {t}") + print("── ADMM solver results ──────────────────────────────────────────────") + print(f" Converged: {result.converged}") + print(f" Iterations: {result.iterations}") + print(f" Primal feasibility: {result.primal_feasibility:.4e}") + print(f" Dual feasibility: {result.dual_feasibility:.4e}") + print(f" Complementarity: {result.complementarity:.4e}") + print(f" Final rho: {result.rho:.4e}") + print("─────────────────────────────────────────────────────────────────────") + + constraint_impulses = result.retrieveConstraintImpulses() + constraint_velocities = result.retrieveConstraintVelocities() + print(f"\nConstraint impulses ‖λ‖: {np.linalg.norm(constraint_impulses):.4e}") + print(f"Constraint velocities ‖v‖: {np.linalg.norm(constraint_velocities):.4e}") + + if not has_converged: + print("\nWarning: solver did not converge within the iteration budget.") + sys.exit(1) + + # Update configuration and velocity for the next time step by + # applying the constraint impulses. + constraint_forces = ( + 1.0 / dt + ) * constraint_impulses # convert impulses to forces/torques + tau_constraints = Jc.T @ constraint_forces # map to generalized torques + v_new = v + dt * pin.aba(model, data, q, v, zero_torque + tau_constraints, fext) + q_new = pin.integrate(model, q, v_new * dt) + + # The cubes should be stable so the configuration should not change much, + # and the velocity should be close to zero: + assert np.linalg.norm(v_new) < 1e-8 + assert np.linalg.norm(q_new - q) < 1e-8 + + # Update q and v for the next iteration. + q = q_new.copy() + v = v_new.copy()