Source code for state_space_setup_assistant.equilibrium

"""Automatic operating-point (equilibrium) finder.

An equilibrium for the actuation set S is a configuration q where the
gravity/bias torque on every *unactuated* joint vanishes:

    rnea(q, 0, 0)[unactuated dofs] = 0

Actuated joints can hold any position (u_eq = gravity compensation), so their
coordinates are held at the user's seed values and the root-finding runs only
over the unactuated coordinates -- a square problem.

Indexing note: the residual and the decision variables are mapped through
Pinocchio's own joint.idx_q / joint.idx_v, never through the position of a
name in a Python list; Pinocchio's q ordering is not guaranteed to match URDF
declaration order.
"""

from typing import Dict, List, Mapping, Optional

import numpy as np

import pinocchio as pin

from urdf_state_space.state_space import _load_model


def _joint_index_maps(model) -> Dict[str, Dict]:
    """name -> {idx_q, nq, idx_v, nv} for every movable joint."""
    out = {}
    for i in range(1, model.njoints):
        j = model.joints[i]
        out[model.names[i]] = {
            'idx_q': j.idx_q, 'nq': j.nq, 'idx_v': j.idx_v, 'nv': j.nv}
    return out


[docs] def find_equilibrium( urdf_xml: str, actuated_joints: List[str], q_seed: Optional[Mapping[str, float]] = None, floating_base: bool = False, tol: float = 1e-9, ) -> Dict: """Find q_eq near ``q_seed`` such that unactuated joints need no torque. Returns {q_eq: {name: value}, u_eq: {name: value}, residual_norm, converged, iterations, free_joints}. """ model = _load_model(urdf_xml, floating_base) data = model.createData() idx = _joint_index_maps(model) unknown = [n for n in actuated_joints if n not in idx] if unknown: raise ValueError(f'Unknown actuated joints {unknown}; ' f'movable joints are {sorted(idx)}') free = [n for n in sorted(idx) if n not in actuated_joints] multi_dof = [n for n in free if idx[n]['nq'] != 1] if multi_dof: raise ValueError( f'auto-equilibrium does not support multi-DoF/continuous ' f'unactuated joints in this version: {multi_dof}') q = pin.neutral(model) for name, val in (q_seed or {}).items(): if name not in idx: raise ValueError(f'Unknown joint {name!r} in q_seed') if idx[name]['nq'] == 1: q[idx[name]['idx_q']] = float(val) free_iq = np.array([idx[n]['idx_q'] for n in free], dtype=int) free_iv = np.array([idx[n]['idx_v'] for n in free], dtype=int) def gravity_torque(qfull) -> np.ndarray: return pin.rnea(model, data, qfull, np.zeros(model.nv), np.zeros(model.nv)) iterations = 0 if free: lower = np.asarray(model.lowerPositionLimit)[free_iq] upper = np.asarray(model.upperPositionLimit)[free_iq] # Pinocchio uses +/-inf-ish sentinels when the URDF has no limits. bad = ~np.isfinite(lower) | ~np.isfinite(upper) | (lower > upper) lower = np.where(bad, -np.inf, lower) upper = np.where(bad, np.inf, upper) def residual(x): qfull = q.copy() qfull[free_iq] = x return gravity_torque(qfull)[free_iv] from scipy.optimize import least_squares sol = least_squares(residual, q[free_iq], bounds=(lower, upper), xtol=1e-12, ftol=1e-12, gtol=1e-12) q[free_iq] = sol.x iterations = int(sol.nfev) tau = gravity_torque(q) res_norm = float(np.linalg.norm(tau[free_iv], np.inf)) if free else 0.0 converged = res_norm < max(tol, 1e-6) q_eq = {n: float(q[idx[n]['idx_q']]) for n in idx if idx[n]['nq'] == 1} u_eq = {n: float(tau[idx[n]['idx_v']]) for n in actuated_joints if idx[n]['nv'] == 1} return { 'q_eq': q_eq, 'u_eq': u_eq, 'residual_norm': res_norm, 'converged': bool(converged), 'iterations': iterations, 'free_joints': free, }