from pathlib import Path
import glob
import importlib
import os
import re
from dataclasses import dataclass, field
import numpy as np
from .common import _require_pin, _pin_available, initialize_warm_start
from .rowkinds import MECHS, KIND_CONTACT_POS, is_vector_kind
from .config import SolverParams, COST_FIELDS
from .linsys_autotune import resolve_linsys
VARIANTS = ("default", "fc", "eh") # module ABI variants: base / contact forces / exact Hessian
# plant names are identifiers (underscores allowed); the variant suffix is split off lazily
_MODULE_RE = re.compile(r"bsqpN(\d+)_([A-Za-z0-9_]+?)(?:_(fc|eh))?\.")
[docs]
def module_name(plant, N, variant=None):
"""Import name of a built solver module: bsqpN{N}_{plant}[_{variant}]."""
variant = variant or "default"
if variant not in VARIANTS:
raise ValueError(f"variant must be one of {VARIANTS}, got {variant!r}")
return f"bsqpN{N}_{plant}" + ("" if variant == "default" else f"_{variant}")
[docs]
def available(variant="default"):
"""Discover built solver modules in this package.
variant="default" (the plain solver), "fc" (contact-force controls) or
"eh" (exact Hessian) -> {(plant, N): filename} for that variant;
variant="all" -> {(plant, N, variant): filename}."""
here = os.path.dirname(os.path.abspath(__file__))
found = {}
entries = []
for so in sorted(glob.glob(os.path.join(here, "bsqpN*_*.so"))):
m = _MODULE_RE.match(os.path.basename(so))
if m:
entries.append((m.group(2), int(m.group(1)), m.group(3) or "default", os.path.basename(so)))
# modules compiled through gato.register_robot (per-host cache); package modules win
from . import _cache
try:
cached = _cache.available_cached()
except Exception: # no cache dir / unreadable manifest: nothing registered
cached = {}
entries += [(p, n, v, os.path.basename(f)) for (p, n, v), f in cached.items()]
for plant, N, var, fname in entries:
if variant == "all":
found.setdefault((plant, N, var), fname)
elif var == variant:
found.setdefault((plant, N), fname)
return found
[docs]
def robot_info(plant_type):
"""Registry metadata for a plant ({nq, nv, ee_frame, urdf, floating_base,
contact_frames, collision, codegen_key}), or {} if unregistered.
The registry (_registry.json) is written by gato.build / tools/regen_grid.py;
robots registered with gato.register_robot are read from the per-host cache."""
from .builder import load_registry
info = load_registry().get(plant_type)
if info is None:
from . import _cache
entry = _cache.robot_entry(plant_type)
info = dict(entry["meta"], codegen_key=entry["codegen_key"]) if entry else {}
return info
[docs]
def robot_urdf(plant_type):
"""Absolute path of a registered robot's URDF: the checkout's copy for the
vendored robots (examples/<robot>_description, with meshes), the wheel's
bundled copy (URDF only — the headers are pre-generated, so no meshes are
needed), or the path recorded by gato.register_robot."""
info = robot_info(plant_type)
if not info:
raise ValueError(f"unknown robot {plant_type!r}: not vendored and not registered with gato.register_robot")
path = Path(info["urdf"])
if path.is_absolute():
return path
from ._cache import native_root
return native_root() / path
# S·λ = γ linear-system paths (see gato/bsqp/kernels/{pcg,bdsv}.cuh):
# pcg = iterative (warm-start friendly), bdsv = direct block-Cholesky (exact,
# iteration-count free), bdsv_first = direct on SQP iteration 0 then pcg.
LINSYS_MODES = {"pcg": 0, "bdsv": 1, "bdsv_first": 2}
_LINSYS_NAMES = {v: k for k, v in LINSYS_MODES.items()}
[docs]
@dataclass(frozen=True)
class SolverStats:
"""Per-solve solver statistics (batch-shaped numpy arrays)."""
solve_time_us: int
sqp_iters: np.ndarray # (B,) int32
kkt_converged: np.ndarray # (B,) int32
final_merit: np.ndarray # (B,) float32 — merit of the returned trajectory
initial_merit: np.ndarray # (B,) float32 — merit before the first iteration
ls_num_iters: int # SQP iterations that reached the line search
pcg_iters: np.ndarray # (sqp_iters_run, B) int32
pcg_times_us: np.ndarray # (sqp_iters_run,) float32
min_merit: np.ndarray # (ls_num_iters, B) float32 — accepted merit per iter
step_size: np.ndarray # (ls_num_iters, B) float32 — -1 marks a line-search failure
# linear-system path used for this solve. NOTE pcg_iters semantics on the
# bdsv path: 0 still means "converged at start" (same guard as pcg), a
# direct solve reports 1, and 2 marks a SKIPPED update (f32 Cholesky hit a
# non-PD pivot — barely-regularized costs; λ kept its warm start and rho
# adaptation retries). Only == 0 carries meaning downstream.
linsys: str = "pcg" # "pcg" | "bdsv" | "bdsv_first"
budget_hit: bool = False # the SQP loop stopped on solve(budget_s=...), not on convergence
# constraint row-group telemetry (None until a row group exists — any
# enable_limit_* mode or an appended group): (n_groups, B) true violation of
# the RETURNED trajectory per group, in get_row_groups() order (the three
# limit boxes first, then the appended groups)
row_max_violation: np.ndarray = None
row_sum_violation: np.ndarray = None
# (B,) last-iteration ADMM residuals (None unless enable_limit_admm())
admm_r_prim: np.ndarray = None
admm_r_dual: np.ndarray = None
# (B,) ADMM inner iterations run this solve (summed over SQP iterations) and
# the inner linear-system work they took: PCG steps on the "pcg" arm, one
# per back-solve on the "bdsv_factor" arm (None unless ADMM is active)
admm_iters: np.ndarray = None
admm_linsys_iters: np.ndarray = None
[docs]
@dataclass(frozen=True)
class SolveResult:
"""Result of one batched solve.
``xu`` is the flat batch of trajectories, one row per solve:
``[x_0, u_0, x_1, u_1, ..., x_{N-1}]`` with row length ``N*(nx+nu) - nu``.
"""
xu: np.ndarray # (B, N*(nx+nu)-nu) float32
solve_time_us: int
stats: SolverStats
nx: int
nu: int # full control width (== n_actuated + n_fc)
N: int
# On GATO_CONTACT_FORCES builds the control is [tau; fc] and only the first
# n_actuated entries are torques an MPC applies; fc is the solver's contact
# explanation. 0 = legacy default (all of nu is actuated).
n_actuated: int = 0
@property
def batch_size(self):
return self.xu.shape[0]
@property
def _na(self):
return self.n_actuated or self.nu
@property
def n_fc(self):
return self.nu - self._na
[docs]
def xu_b(self, b=0):
"""Trajectory of batch entry b (flat)."""
return self.xu[b]
[docs]
def u0(self, b=0):
"""First ACTUATED control of batch entry b — what an MPC applies
(contact-force slots, if any, are excluded)."""
return self.xu[b, self.nx:self.nx + self._na]
[docs]
def control_at(self, k, b=0):
"""Actuated control at knot k (clamped to the last control knot) of batch
entry b (contact-force slots, if any, are excluded)."""
k = min(int(k), self.N - 2) # last knot has no control
start = self.nx + (self.nx + self.nu) * k
return self.xu[b, start:start + self._na]
[docs]
def fc_at(self, k, b=0):
"""Contact-wrench slots at knot k of batch entry b ((n_fc,); empty on
non-contact-force builds)."""
k = min(int(k), self.N - 2)
start = self.nx + (self.nx + self.nu) * k + self._na
return self.xu[b, start:start + self.n_fc]
[docs]
def fc_traj(self, b=0):
"""(N-1, n_fc) contact-wrench trajectory of batch entry b."""
return np.stack([self.fc_at(k, b) for k in range(self.N - 1)])
@property
def diverged(self):
"""(B,) bool — non-finite merit or trajectory."""
bad_merit = ~np.isfinite(self.stats.final_merit)
bad_xu = ~np.isfinite(self.xu).all(axis=1)
return bad_merit | bad_xu
[docs]
class BSQP:
"""Batched SQP trajectory optimizer over one compiled solver module.
Args:
model_path: robot URDF (must be the one the module was codegen'd from).
batch_size: number of parallel solves B (runtime; any B >= 1).
N: horizon length — selects the module bsqpN{N}_{plant}[_{variant}].
dt: knot spacing [s].
params: gato.SolverParams (defaults = the paper/MPC set); keyword
overrides of its fields may be passed directly (``BSQP(..., mu=5.0)``).
plant_type: plant name; None auto-detects from model_path against the
built modules / registry.
variant: "default" | "fc" (contact-force controls appended to u) |
"eh" (exact-Hessian SO-SQP). Variant modules are separate ABIs.
rho_batch / mu_batch / pcg_tol_batch: optional per-solve (B,) overrides.
Construction needs numpy and the module only. pinocchio is loaded lazily
by the FK helpers (ee_pos) and the sim/estimator layers (``[test]`` extra).
Warm start: ``solve(x, ref, xu_warm=None)`` is STATELESS w.r.t. the
trajectory — None seeds a hold at x; pass the previous ``SolveResult.xu``
(or let MPCController own the shifted warm start) to warm-start. The
solver IS stateful across solves in rho adaptation, duals and ADMM state
(reset_rho / reset_dual / reset).
Cost precedence: SolverParams scalars < set_cost_weights (scalars) <
set_cost_weights_per_knot (N,3) [ee, qd, u]; per-joint vectors
(set_q_pos_cost / set_u_cost_vec) and the fc cost are orthogonal terms.
Reference layout: ``ref`` is (B, 6N): per knot [x, y, z, 0, 0, 0] — the
EE position target (orientation slots are reserved, currently unused).
Trajectory layout: ``xu`` is (B, N*(nx+nu) - nu) =
[x_0, u_0, x_1, u_1, ..., x_{N-1}] with x = [q; qd].
"""
def __init__(self, model_path, batch_size, N, dt, params=None, *,
plant_type=None, variant=None,
rho_batch=None, mu_batch=None, pcg_tol_batch=None, **overrides):
params = SolverParams.from_mapping(params)
if overrides:
params = params.replace(**overrides) # unknown names raise here
self.params = params
model_path = str(model_path)
# Auto-detect plant type from model_path if not explicitly specified.
# Unknown robots are a hard error (a wrong plant silently runs the wrong
# dynamics with mismatched state size).
if plant_type is None:
from .builder import load_registry
plants = sorted({p for p, _ in available(variant or "default")} | set(load_registry()))
low = model_path.lower()
# match the plant name or its alpha prefix (iiwa14 -> "iiwa") in the path
matches = [p for p in plants
if p in low or re.sub(r"\d+$", "", p) in os.path.basename(low)]
if len(matches) != 1:
raise ValueError(
f"Could not auto-detect plant from model_path={model_path!r}; "
f"pass plant_type explicitly. Built plants: {plants or 'none'}"
)
plant_type = matches[0]
self.variant = variant or "default"
mod = module_name(plant_type, N, self.variant)
try:
base = importlib.import_module("gato." + mod)
except ImportError as e:
# the per-host cache (gato.register_robot): compiles once on first use
from . import _cache
base = _cache.load_module(plant_type, N, self.variant)
if base is None:
flag = {"fc": " contact_forces=True", "eh": " exact_hessian=True", "default": ""}[self.variant]
raise ValueError(
f"No compiled module for plant={plant_type!r}, N={N}, variant={self.variant!r} "
f"(could not import gato.{mod}): {e}\n"
f"Built modules: {sorted(available('all')) or 'none'}. Register the robot once —\n"
f" python -c \"import gato; gato.register_robot('{plant_type}', '{model_path}', N=[{N}])\"\n"
f"— and the first solver constructed for it compiles the module (cached afterwards);\n"
f"from a source checkout: gato.build('{model_path}', name='{plant_type}', N=[{N}],{flag})"
)
# batch_size is a runtime constructor argument (one class per precision)
class_name = "BSQP_float"
if not hasattr(base, class_name):
raise ValueError(
f"Module {mod} does not export {class_name} — rebuild the "
f"solver modules (old per-batch-size builds are incompatible)"
)
self.lib = base
self.solver_class = getattr(base, class_name)
self.plant_type = plant_type
# Body-major external-force buffer width is 6*NUM_BODIES per solve (the GPU
# d_f_ext_batch_ buffer + set_f_ext_batch upload are sized to this).
self.n_bodies = int(getattr(base, "NUM_BODIES", 0))
self._cost_weights = {k: float(getattr(params, k)) for k in COST_FIELDS}
self.solver = self.solver_class(
batch_size,
int(params.max_sqp_iters),
int(params.max_pcg_iters), float(params.pcg_tol), float(params.solve_ratio), float(params.mu),
float(params.q_cost), float(params.qd_cost), float(params.u_cost), float(params.N_cost),
float(params.q_lim_cost), float(params.vel_lim_cost), float(params.ctrl_lim_cost),
float(params.rho),
)
# Dimensions come from the MODULE (NQ/NV/CONTROL_SIZE/... are baked at
# compile time) — pinocchio is NOT needed to construct or run a solver.
# The pinocchio model is built lazily (self.model / self.data) by the
# FK helpers (ee_pos), the sim/estimator layers and the tests; when it
# IS available the module/URDF dimension cross-check below runs.
self.floating_base = bool(getattr(base, "FLOATING_BASE", False))
self.model_path = model_path
self._model = None
self._data = None
self._ee_frame_id = None
self.nq, self.nv = int(base.NQ), int(base.NV)
if _pin_available():
if (self.nq, self.nv) != (self.model.nq, self.model.nv):
raise ValueError(
f"module {mod} has (nq, nv) = ({self.nq}, {self.nv}) but the "
f"pinocchio model from {model_path!r} has ({self.model.nq}, {self.model.nv}) — "
f"wrong URDF or a fixed/floating mismatch")
# The solver/grid.cuh optimizes the EE-position cost in the frame the module
# was codegen'd with (fixed_target_name; recorded in the registry, "EE" for
# the vendored robots). The frame id is resolved lazily with the model.
self.ee_frame = robot_info(plant_type).get("ee_frame", "EE")
self.batch_size = batch_size
self.N = N
self.dt = dt
self.f_ext_B = np.zeros((self.batch_size, 6), dtype=np.float32)
self.set_f_ext_B(self.f_ext_B)
self.nx = self.nq + self.nv
# Control width comes from the MODULE: on fc variants CONTROL_SIZE =
# ACTUATED_SIZE + FC_SIZE (contact-wrench slots appended after the
# torques), so nu > nv and every xu stride follows it.
self.n_actuated = int(getattr(base, "ACTUATED_SIZE", self.nv))
self.n_fc = int(getattr(base, "FC_SIZE", 0))
self.nu = int(getattr(base, "CONTROL_SIZE", self.nv))
self.xu_size = self.N * (self.nx + self.nu) - self.nu
# The baked contact frames (registry, codegen order): fc slot block i
# (fc_slots(i)) is the world-aligned wrench [n; f] at contact_frames[i].
# Empty when the module was generated without contact frames.
self.contact_frames = list(robot_info(plant_type).get("contact_frames") or [])
if self.n_fc and len(self.contact_frames) * 6 != self.n_fc:
raise ValueError(
f"module {mod} has FC_SIZE={self.n_fc} but the registry lists "
f"{len(self.contact_frames)} contact frames for {plant_type!r} — stale registry/headers, regen")
# Optional batched hyperparameters
if rho_batch is not None:
rho_batch = np.asarray(rho_batch, dtype=np.float32).reshape(self.batch_size)
self.solver.set_rho_penalty_batch(rho_batch, True)
self.solver.set_rho_adaptation(bool(params.adapt_rho))
if mu_batch is not None:
mu_batch = np.asarray(mu_batch, dtype=np.float32).reshape(self.batch_size)
self.solver.set_mu_batch(mu_batch)
if pcg_tol_batch is not None:
pcg_tol_batch = np.asarray(pcg_tol_batch, dtype=np.float32).reshape(self.batch_size)
self.solver.set_pcg_tol_batch(pcg_tol_batch)
self.max_pcg_iters = int(params.max_pcg_iters)
self.linsys = "pcg" # the C++ default; set_linsys only calls into the module on change
# ONE resolver for both the raw solver and the controller
# (linsys_autotune.resolve_linsys): the solver takes the STATIC arm —
# "auto" (a per-step controller policy) maps to its warm-body path, pcg.
mode, _ = resolve_linsys(self.floating_base, params.linsys, None,
plant=plant_type, N=N, task_tag=None)
self.set_linsys("pcg" if mode == "auto" else mode)
self._row_mech = None # active enable_limit_* mode (add_lin_u_rows mech=None default)
self._n_appended_groups = 0
self.exact_hessian = False # the C++ default
self._time_budget_us = 0.0 # solve(budget_s=...) state mirrored in the C++ solver
self._admm_linsys = "pcg" # set_admm_linsys; the C++ default
self._admm_exit = (0.0, 0.0, 1) # set_admm_exit (eps_abs, eps_rel, check_every); the C++ default
if params.exact_hessian:
self.set_exact_hessian(True)
[docs]
def set_linsys(self, mode):
"""Pick the S·λ = γ path for subsequent solves: "pcg" | "bdsv" | "bdsv_first".
Host-side and zero-cost — an MPC loop can switch it per step.
"""
if mode not in LINSYS_MODES:
raise ValueError(f"linsys must be one of {sorted(LINSYS_MODES)}, got {mode!r}")
if mode != self.linsys:
self.solver.set_linsys_mode(LINSYS_MODES[mode])
self.linsys = mode
self.params = self.params.replace(linsys=mode)
[docs]
def set_admm_linsys(self, mode):
"""ADMM inner-loop linear solver: "pcg" (default) | "bdsv_factor".
Measured defaults, rulings and provenance: docs/constraints.md (`set_admm_linsys`)."""
if mode not in ("bdsv_factor", "pcg"):
raise ValueError(f'admm_linsys must be "bdsv_factor" or "pcg", got {mode!r}')
self.solver.set_admm_linsys_pcg(mode == "pcg")
self._admm_linsys = mode
@property
def admm_linsys(self):
"""The ADMM inner-loop linear solver currently set ("pcg" | "bdsv_factor")."""
return self._admm_linsys
@property
def admm_active(self):
"""True while any row group runs the ADMM mechanism (limits or appended)."""
return bool(self.solver.admm_active())
[docs]
def set_admm_exit(self, eps_abs=0.0, eps_rel=0.0, check_every=1):
"""Residual early exit of the ADMM inner loop. Per solve, after each inner
iteration, the loop stops once every ADMM row group g has
``r_prim_g <= eps_abs + eps_rel*max|z_g|`` and
``r_dual_g <= eps_abs + eps_rel*rho_g*max|z_g|`` on the iterate just
produced (rows of a group share units; the test is per group so one
group's scale cannot excuse another's) AND that iterate's linear solve converged (PCG below its
cap, or an OK direct factor): an unconverged inner solve cannot read as
a fixed point. ``eps_abs = eps_rel = 0`` (default) is the fixed budget of
``enable_limit_admm(iters=...)``, bit for bit; with a tolerance ``iters``
is the cap. ``check_every``: inner iterations between host checks that
break the launch loop once every batch entry has exited (0 = never; exited
entries are skipped on-device regardless). ``stats.admm_iters`` and
``stats.admm_linsys_iters`` report what ran.
Measured defaults, rulings and provenance: docs/constraints.md (`set_admm_exit`)."""
if eps_abs < 0.0 or eps_rel < 0.0 or int(check_every) < 0:
raise ValueError(f"set_admm_exit: eps_abs, eps_rel and check_every must be >= 0, "
f"got {eps_abs!r}, {eps_rel!r}, {check_every!r}")
self.solver.set_admm_exit(float(eps_abs), float(eps_rel), int(check_every))
self._admm_exit = (float(eps_abs), float(eps_rel), int(check_every))
@property
def admm_eps_abs(self):
"""Absolute tolerance of the ADMM residual exit (0 with the exit off)."""
return self._admm_exit[0]
@property
def admm_eps_rel(self):
"""Relative tolerance of the ADMM residual exit (0 with the exit off)."""
return self._admm_exit[1]
@property
def admm_check_every(self):
"""Host break period of the ADMM residual exit, in inner iterations."""
return self._admm_exit[2]
def _apply_budget(self, budget_s):
"""Mirror solve(budget_s=...) into the C++ solver (a host-side switch; only
touched when the value changes, so the unbudgeted path makes no extra call)."""
if budget_s is None:
us = 0.0
else:
us = float(budget_s) * 1e6
if not us > 0.0:
raise ValueError(f"budget_s must be positive (None = no budget), got {budget_s!r}")
if us != self._time_budget_us:
self.solver.set_time_budget_us(us)
self._time_budget_us = us
[docs]
def set_collect_stats(self, on):
"""Collect per-iteration solver stats (PCG timings via cudaEvents,
line-search merits). On by default; MPC loops that do not read
SolverStats can switch it off (two event records per SQP iteration)."""
self.solver.set_collect_stats(bool(on))
[docs]
def set_drho_batch(self, drho_batch, set_as_reset_default=True):
"""Per-solve trust-region rho ADAPTATION factor (B,) — the multiplier the
line search applies on reject/accept (reset_rho restores the default)."""
d = np.ascontiguousarray(np.asarray(drho_batch, dtype=np.float32).reshape(self.batch_size))
self.solver.set_drho_batch(d, bool(set_as_reset_default))
[docs]
def get_lambda(self):
"""(B, N+2, nx) lagged KKT multipliers λ from the last solve (padded layout)."""
return np.asarray(self.solver.get_lambda()).reshape(self.batch_size, self.N + 2, self.nx)
[docs]
def debug_setup_kkt(self, xu_B, xcur_B, eepos_goals_B):
"""Run setup_kkt ONLY on a trajectory and return the assembled KKT blocks
{Q, R, q, r, A, B, c} (B, ...) — the way the anchor/exact-Hessian gates
pin cost claims directly instead of arguing from solve outcomes."""
xu = np.ascontiguousarray(np.asarray(xu_B, dtype=np.float32).reshape(self.batch_size, self.xu_size))
x = np.ascontiguousarray(np.asarray(xcur_B, dtype=np.float32).reshape(self.batch_size, self.nx))
g = np.ascontiguousarray(np.asarray(eepos_goals_B, dtype=np.float32).reshape(self.batch_size, -1))
return self.solver.debug_setup_kkt(xu, self.dt, x, g)
[docs]
def exact_hessian_available(self):
"""True if the loaded module was compiled with -DGATO_EXACT_HESSIAN=ON."""
return bool(getattr(self.lib, "EXACT_HESSIAN_AVAILABLE", False))
[docs]
def set_exact_hessian(self, on):
"""Toggle the SO-SQP stage-Hessian PSD projection for subsequent solves.
Measured defaults, rulings and provenance: docs/constraints.md (`set_exact_hessian`)."""
on = bool(on)
if on and not self.exact_hessian_available():
raise RuntimeError(
f"module {self.lib.__name__} compiled without USE_EXACT_HESSIAN — "
"rebuild with cmake -DGATO_EXACT_HESSIAN=ON"
)
if on != self.exact_hessian:
self.solver.set_exact_hessian(on)
self.exact_hessian = on
self.params = self.params.replace(exact_hessian=bool(on))
def _heal_floating_warm_start(self, XU_B, xcur_B):
"""Replace warm-start knots whose base quaternion is degenerate by the
measured state (controls kept). The constructed XU_B is all-zero, and a
zero quaternion is not a point on the manifold — the device solve
cannot step off it (state_difference/retract are undefined there), so
a cold floating solve would silently return the zero tail unchanged.
Valid warm starts pass through untouched (bit-identical)."""
step = self.nx + self.nu
for k in range(self.N):
quat = XU_B[:, k * step + 3:k * step + 7]
bad = np.linalg.norm(quat, axis=1) < 0.5
if bad.any():
XU_B[bad, k * step:k * step + self.nx] = xcur_B[bad]
def _solve_device(self, xcur_B, eepos_goals_B, xu_warm):
"""solve() on CUDA inputs: validate, hand device pointers to the binding."""
from . import _arrays
B = self.batch_size
ref_size = int(getattr(self.lib, "REFERENCE_TRAJ_SIZE", 6 * self.N))
x_ptr, keep_x = _arrays.device_f32(xcur_B, B * self.nx, "xcur_B")
g_ptr, keep_g = _arrays.device_f32(eepos_goals_B, B * ref_size, "eepos_goals_B")
if xu_warm is None:
xu_ptr, keep_xu = 0, None
else:
if not _arrays.is_cuda(xu_warm):
raise TypeError("xu_warm must live on the same device as xcur_B for the device path")
xu_ptr, keep_xu = _arrays.device_f32(xu_warm, B * self.xu_size, "xu_warm")
return self.solver.solve_device(int(xu_ptr), self.dt, int(x_ptr), int(g_ptr))
[docs]
def solve(self, xcur_B, eepos_goals_B, xu_warm=None, *, budget_s=None):
"""Solve the batch from states ``xcur_B`` (B, nx) toward the EE reference
``eepos_goals_B`` (B, 6N) — see the class docstring for both layouts.
``budget_s``: wall-clock budget for the SQP loop, in seconds. The loop
runs the next iteration only while the time spent so far plus that
iteration's predicted cost (its measured cost in the previous budgeted
solve; the most expensive iteration of this solve when there is none)
still fits, and returns the current trajectory (the line search
accepts descent steps only, so it is the best plan so far); the first
iteration always runs. The cut is at the iteration boundary, never
inside a kernel, so a budget is exceeded by at most one iteration. ``SolverStats.budget_hit`` reports a
budget stop and ``sqp_iters`` the iterations run. None (default): no
budget, ``max_sqp_iters`` iterations, the path unchanged by this option.
``xu_warm`` (B, N*(nx+nu)-nu): the warm-start trajectory. None seeds a
hold at ``xcur_B`` (a cold start; NOT the previous solution — the solver
keeps no trajectory state; pass the last ``SolveResult.xu`` to
warm-start, or use MPCController). The caller's array is never
modified; ``SolveResult.xu`` is a fresh array.
Arrays may be numpy, torch or JAX; the result's ``xu`` is of the same
kind. CUDA tensors (torch on a CUDA device, JAX on a GPU) take the
device path: inputs are read in place through their device pointer and
``xu`` comes back on the device through DLPack, with no host round
trip; a device ``xu_warm`` is used as given (its knot-0 states are
replaced by ``xcur_B``; the floating-base quaternion healing of the
host path is not applied). CPU tensors go through the host path.
"""
from . import _arrays
self._apply_budget(budget_s)
in_kind = _arrays.kind(xcur_B)
if _arrays.is_cuda(xcur_B):
raw = self._solve_device(xcur_B, eepos_goals_B, xu_warm)
xu = _arrays.from_device(raw["XU"], in_kind)
else:
xcur_np = _arrays.to_numpy(xcur_B).reshape(self.batch_size, self.nx)
goals_np = _arrays.to_numpy(eepos_goals_B).reshape(self.batch_size, -1)
if xu_warm is None:
XU_B = np.stack([initialize_warm_start(xcur_np[b], self.N, self.nx, self.nu)
for b in range(self.batch_size)]).astype(np.float32)
else:
XU_B = _arrays.to_numpy(xu_warm).reshape(self.batch_size, self.xu_size).copy()
XU_B[:, : self.nx] = xcur_np
if self.floating_base:
self._heal_floating_warm_start(XU_B, xcur_np)
raw = self.solver.solve(XU_B, self.dt, xcur_np, goals_np)
xu = np.array(raw["XU"], dtype=np.float32, copy=True)
if in_kind != "numpy":
xu = _arrays.like(xu, in_kind, device=getattr(xcur_B, "device", None))
return self._result(raw, xu)
# ---- device-resident warm start: MPCController's device path ----------------
# The solver's trajectory buffer doubles as a warm-start buffer that stays on
# the device between calls: set_warm / hold_warm seed it, solve_warm solves
# from it (knot-0 states := x) and leaves the solution in it, shift_warm
# advances it one knot for the next tick. x and the reference are CUDA
# tensors (torch or JAX) the solver only reads; ``xu`` comes back on the
# device. solve() never touches this buffer's contents between calls: it
# overwrites them with its own warm start.
[docs]
def set_warm(self, xu):
"""Seed the resident warm start: (xu_size,) broadcast to the batch or
(B, xu_size); numpy (uploaded once) or a CUDA tensor (device copy)."""
from . import _arrays
if _arrays.is_cuda(xu):
ptr, _keep, batched = _arrays.device_rows(xu, self.xu_size, self.batch_size, "xu")
self.solver.resident_set_device(int(ptr), batched)
return
a = _arrays.to_numpy(xu)
if a.shape not in ((self.xu_size,), (self.batch_size, self.xu_size)):
raise ValueError(f"xu must have shape ({self.xu_size},) or ({self.batch_size}, {self.xu_size}), got {a.shape}")
self.solver.resident_set_host(a)
[docs]
def get_warm(self):
"""The resident warm start as a fresh (B, xu_size) numpy array."""
return np.array(self.solver.resident_get_host(), dtype=np.float32, copy=True)
[docs]
def hold_warm(self, x):
"""Seed the resident warm start with hold-at-x (every knot's state = x,
zero controls): x a CUDA tensor, (nx,) broadcast or (B, nx)."""
from . import _arrays
ptr, _keep, batched = _arrays.device_rows(x, self.nx, self.batch_size, "x")
self.solver.resident_hold(int(ptr), batched)
[docs]
def pred_err_warm(self, x):
"""‖x − knot 0 of the resident row 0‖ for a CUDA state x (nx,) (the
controller's warm-startedness signal; floating base with a unit base
quaternion: the tangent norm). Returns a Python float."""
from . import _arrays
ptr, _keep, _batched = _arrays.device_rows(x, self.nx, self.batch_size, "x")
return float(self.solver.resident_pred_err(int(ptr)))
[docs]
def solve_warm(self, x, ref, *, budget_s=None):
"""Solve from the resident warm start (its knot-0 states replaced by x).
x: CUDA (nx,) broadcast or (B, nx); ref: CUDA (6N,) broadcast or (B, 6N);
``budget_s`` as in solve(). The solution stays resident and comes back
as ``SolveResult.xu`` on the device in the caller's kind."""
from . import _arrays
if not (_arrays.is_cuda(x) and _arrays.is_cuda(ref)):
raise TypeError("solve_warm needs CUDA tensors for x and ref (torch on a CUDA device or JAX on a GPU)")
self._apply_budget(budget_s)
ref_size = int(getattr(self.lib, "REFERENCE_TRAJ_SIZE", 6 * self.N))
x_ptr, _kx, xb = _arrays.device_rows(x, self.nx, self.batch_size, "x")
r_ptr, _kr, rb = _arrays.device_rows(ref, ref_size, self.batch_size, "ref")
raw = self.solver.solve_resident(int(x_ptr), xb, int(r_ptr), rb, self.dt)
return self._result(raw, _arrays.from_device(raw["XU"], _arrays.kind(x)))
[docs]
def shift_warm(self, best_id=0, hold=False):
"""Next tick's warm start from resident row ``best_id`` into every row:
the plan advanced one knot with the last knot repeated (``hold=False``,
the controller's "shift") or the plan as solved (``hold=True``)."""
self.solver.resident_shift(int(best_id), bool(hold))
def _result(self, raw, xu):
B = self.batch_size
stats = SolverStats(
solve_time_us=int(raw["sqp_time_us"]),
sqp_iters=np.asarray(raw["sqp_iters"], dtype=np.int32).reshape(B),
kkt_converged=np.asarray(raw["kkt_converged"], dtype=np.int32).reshape(B),
final_merit=np.asarray(raw["final_merit"], dtype=np.float32).reshape(B),
initial_merit=np.asarray(raw["initial_merit"], dtype=np.float32).reshape(B),
ls_num_iters=int(raw["ls_num_iters"]),
pcg_iters=np.asarray(raw["pcg_iters"], dtype=np.int32).reshape(-1, B),
pcg_times_us=np.asarray(raw["pcg_times_us"], dtype=np.float32).reshape(-1),
min_merit=np.asarray(raw["ls_min_merit"], dtype=np.float32).reshape(-1, B),
step_size=np.asarray(raw["ls_step_size"], dtype=np.float32).reshape(-1, B),
linsys=_LINSYS_NAMES[int(raw.get("linsys_mode", 0))],
budget_hit=bool(raw["budget_hit"]),
row_max_violation=(np.asarray(raw["row_max_violation"], dtype=np.float32)
if "row_max_violation" in raw else None),
row_sum_violation=(np.asarray(raw["row_sum_violation"], dtype=np.float32)
if "row_sum_violation" in raw else None),
admm_r_prim=(np.asarray(raw["admm_r_prim"], dtype=np.float32)
if "admm_r_prim" in raw else None),
admm_r_dual=(np.asarray(raw["admm_r_dual"], dtype=np.float32)
if "admm_r_dual" in raw else None),
admm_iters=(np.asarray(raw["admm_iters"], dtype=np.int32).reshape(B)
if "admm_iters" in raw else None),
admm_linsys_iters=(np.asarray(raw["admm_linsys_iters"], dtype=np.int32).reshape(B)
if "admm_linsys_iters" in raw else None),
)
return SolveResult(xu=xu, solve_time_us=stats.solve_time_us,
stats=stats, nx=self.nx, nu=self.nu, N=self.N,
n_actuated=self.n_actuated)
# rows::Mechanism enum values (rowgroups.cuh)
_MECHS = MECHS
def _resolve_mech_rho(self, mech, rho, admm_rho):
"""(mech name, rho): the active enable_limit_* mode by default; per-mechanism
rho defaults, with the ADMM one passed by the caller (0.01 for control-block
rows whose Hessian scale is the control cost, 1.0 for rows that fold onto
the state block)."""
if mech is None:
mech = self._row_mech or "telemetry"
if mech not in self._MECHS:
raise ValueError(f"mech must be one of {sorted(self._MECHS)}, got {mech!r}")
if rho is None:
rho = {"telemetry": 0.0, "barrier": 3e-3, "admm": admm_rho, "al": 1.0}[mech]
return mech, float(rho)
def _require_fc(self, what):
if self.n_fc == 0:
raise RuntimeError(f"{what} needs a GATO_CONTACT_FORCES build (this module has no fc slots)")
@staticmethod
def _bcast(v, n):
"""scalar -> (n,) float32, or an array reshaped to (n,)"""
return np.ascontiguousarray((np.full(n, v) if np.ndim(v) == 0 else np.asarray(v)).astype(np.float32).reshape(n))
@property
def row_mech(self):
"""The active enable_limit_* mechanism ("telemetry" | "barrier" | "admm" |
"al"), or None when no limit mechanism is installed."""
return self._row_mech
[docs]
def enable_limit_telemetry(self):
"""Install the canonical limit row-groups (URDF position/velocity/torque boxes)
in TELEMETRY mode: solves report each group's true violation in
``stats.row_{max,sum}_violation``; trajectories stay bit-identical.
Measured defaults, rulings and provenance: docs/constraints.md (`enable_limit_telemetry`)."""
self._check_no_appended_groups("enable_limit_telemetry")
self.solver.enable_limit_telemetry()
self._row_mech = "telemetry"
[docs]
def enable_limit_barrier(self, mu=3e-3, delta=0.05):
"""Bind the limit row-groups to the RELAXED log-barrier mechanism (soft prior;
C² barrier with bounded Hessian, quadratic extension within ``delta``).
Measured defaults, rulings and provenance: docs/constraints.md (`enable_limit_barrier`)."""
self._check_no_appended_groups("enable_limit_barrier")
self.solver.enable_limit_barrier(float(mu), float(delta))
self._row_mech = "barrier"
[docs]
def enable_limit_admm(self, rho=0.01, iters=10, *, rho_q=None, rho_qd=None, rho_u=None,
eps_abs=None, eps_rel=None, check_every=1):
"""Bind the limit row-groups to the ADMM-projection mechanism: an OSQP-style
inner loop per SQP iteration (``iters`` iterations, warm PCG or the reused
direct factor, see set_admm_linsys) — the layer's "approximately hard"
mode; ``rho`` is the ADMM penalty, and ``rho_q`` / ``rho_qd`` / ``rho_u``
override it per box (position, velocity, torque rows; see
set_row_group_rho: the proximal fold belongs at each block's cost scale).
``eps_abs`` / ``eps_rel`` (either given) turn on the residual early exit of
set_admm_exit, with ``iters`` as the cap.
Measured defaults, rulings and provenance: docs/constraints.md (`enable_limit_admm`)."""
self._check_no_appended_groups("enable_limit_admm")
self.solver.enable_limit_admm(float(rho), int(iters))
self._row_mech = "admm"
self._set_limit_rhos(rho_q, rho_qd, rho_u)
if eps_abs is not None or eps_rel is not None:
self.set_admm_exit(eps_abs or 0.0, eps_rel or 0.0, check_every)
[docs]
def enable_limit_al(self, rho=1.0, *, rho_q=None, rho_qd=None, rho_u=None):
"""Bind the limit row-groups to the PHR augmented-Lagrangian mechanism: the
outer dual update runs once per solve on the final trajectory, so
warm-started repeat solves are the outer loop (``rho`` fixed per enable,
``rho_q`` / ``rho_qd`` / ``rho_u`` override it per box; duals persist
across solves, reset_dual() zeroes them).
Measured defaults, rulings and provenance: docs/constraints.md (`enable_limit_al`)."""
self._check_no_appended_groups("enable_limit_al")
self.solver.enable_limit_al(float(rho))
self._row_mech = "al"
self._set_limit_rhos(rho_q, rho_qd, rho_u)
def _set_limit_rhos(self, rho_q, rho_qd, rho_u):
"""Per-box mechanism weights for the three canonical limit groups (None = keep)."""
from .rowkinds import KIND_BOX_Q, KIND_BOX_QD, KIND_BOX_U
wanted = {KIND_BOX_Q: rho_q, KIND_BOX_QD: rho_qd, KIND_BOX_U: rho_u}
for g, grp in enumerate(self.solver.get_row_groups()):
r = wanted.get(grp["kind"])
if r is not None:
self.set_row_group_rho(g, r)
[docs]
def set_row_group_rho(self, g, rho):
"""Mechanism weight of row group ``g`` (barrier mu, ADMM rho or AL rho,
whichever mechanism the group runs; ``get_row_groups()[g]["rho"]`` reads
it). The enable calls give the three limit boxes one weight; the ADMM
proximal fold rho*G'G belongs at each block's cost scale (measured
October 16, 2026: rho 0.01 on the torque rows against the 2e-6 control
cost over-regularizes the controls and slows the arm where no limit
binds). ADMM re-initializes (z, y) on the next solve.
Measured defaults, rulings and provenance: docs/constraints.md (`enable_limit_admm`)."""
if not float(rho) >= 0.0:
raise ValueError(f"set_row_group_rho: rho must be >= 0, got {rho!r}")
self.solver.set_row_group_rho(int(g), float(rho))
[docs]
def enable_ee_terminal_equality(self, target, rho=10.0):
"""Append an EE terminal-position equality row-group (final-knot EE position ==
``target``, on-device FK in the solver's EE frame). The mechanism follows the
active enable_limit_* mode; call it AFTER the mechanism enable.
Measured defaults, rulings and provenance: docs/constraints.md (`enable_ee_terminal_equality`)."""
self.solver.enable_ee_terminal_equality(
np.asarray(target, dtype=np.float32).reshape(3), float(rho))
self._n_appended_groups += 1
[docs]
def disable_row_groups(self):
"""Remove all constraint row-groups (stats lose the row_* fields)."""
self.solver.disable_row_groups()
self._row_mech = None
self._n_appended_groups = 0
def _check_no_appended_groups(self, what):
# Mechanism enables reinstall the canonical 3 limit groups and would
# silently DROP appended groups (EE rows, LIN_U rows, collision) — the
# ordering contract is enforced instead of documented: enables first.
if self._n_appended_groups:
raise RuntimeError(
f"{what}() after {self._n_appended_groups} appended row-group(s) would drop them; "
"call enable_limit_* BEFORE add_lin_u_rows / enable_ee_terminal_equality / "
"enable_collision (or disable_row_groups() first)")
[docs]
def get_row_groups(self):
"""List of installed row-group descriptors (dicts with kind/block/mech,
knot mask, and per-row lo/hi bounds)."""
return self.solver.get_row_groups()
[docs]
def get_row_duals(self):
"""AL multipliers dict {lam_hi, lam_lo}, each shaped
(B, MAX_ROW_GROUPS, N, MAX_ROWS_PER_GROUP) in the dense row-state
layout (group gi's active slots are [:, gi, knot_lo:knot_hi, :n_rows]).
Equality rows carry their signed multiplier in the lam_hi slot."""
return self.solver.get_row_duals()
[docs]
def get_admm_state(self):
"""ADMM state dict {z, y} (auxiliary/dual), same dense layout as
get_row_duals(). y is the interval-constraint multiplier estimate."""
return self.solver.get_admm_state()
[docs]
def set_row_group_bounds(self, g, lo, hi):
"""Override group ``g``'s interval bounds (n_rows each); ``lo == hi`` rows are
always-active equalities under AL, ADMM's auxiliary state re-clips next solve.
Measured defaults, rulings and provenance: docs/constraints.md (`set_row_group_bounds`)."""
self.solver.set_row_group_bounds(int(g),
np.asarray(lo, dtype=np.float32),
np.asarray(hi, dtype=np.float32))
[docs]
def set_row_group_mask(self, g, mask):
"""Per-knot row-activity mask for row group ``g``.
``mask`` is an (N, n_rows) boolean array (row i live at knot k), or (N,)
for vector kinds (SOC cones, collision groups — gated per knot). The
group's knot window still applies first. All-True is the default
behaviour, bitwise. An inactive (knot, row) is treated exactly like a
knot below the group's ``knot_lo``: ADMM keeps its proximal fold but no
z/y update, gradient, merit or residual; AL/barrier/telemetry contribute
nothing and the AL dual of that slot is reset. ADMM re-inits z/y on the
next solve after every call. ``None`` restores all-True."""
groups = self.solver.get_row_groups()
if g < 0 or g >= len(groups):
raise ValueError(f"row group {g} out of range (have {len(groups)})")
grp = groups[g]
n_rows = int(grp["n_rows"])
vector_kind = is_vector_kind(grp["kind"], grp["cone"]) # SOC cone / COLLISION: bit 0 gates the knot
if mask is None:
words = np.full(self.N, np.uint64(2**64 - 1), dtype=np.uint64)
else:
m = np.asarray(mask, dtype=bool)
if vector_kind:
m = m.reshape(self.N, -1)[:, 0] if m.ndim == 2 else m.reshape(self.N)
words = np.where(m, np.uint64(1), np.uint64(0)).astype(np.uint64)
else:
if m.shape != (self.N, n_rows):
raise ValueError(f"mask must be (N, n_rows) = ({self.N}, {n_rows}); got {m.shape}")
words = np.zeros(self.N, dtype=np.uint64)
for i in range(n_rows):
words |= (m[:, i].astype(np.uint64) << np.uint64(i))
self.solver.set_row_group_mask(int(g), words)
[docs]
def set_row_group_soft(self, g, sigma):
"""Soft/slack toggle for group ``g``: sigma > 0 makes its rows ELASTIC (AL: the
multiplier saturates at sigma; ADMM: smoothed projection); 0 = hard.
Measured defaults, rulings and provenance: docs/constraints.md (`set_row_group_soft`)."""
self.solver.set_row_group_soft(int(g), float(sigma))
[docs]
def set_admm_merit(self, on=True):
"""Ablation toggle: include the AL-form ADMM constraint value
y'(g - z) + (rho/2)|g - z|^2 (current row state) in the line-search
merit. v1 ADMM's merit is tracking-only, so the line search rejects
steps that trade tracking for feasibility (measured: closed-loop MPC
parks in conservative basins). Off by default — the exact v1
semantics; only read while ADMM mode is active.
Measured defaults, rulings and provenance: docs/constraints.md (`set_admm_merit`)."""
self.solver.set_admm_merit(bool(on))
[docs]
def set_admm_rho_adaptation(self, on=True):
"""OSQP-style ADMM rho adaptation (opt-in; default OFF = bitwise
pre-adaptation path). One per-solve SCALAR multiplier on top of every
ADMM group's rho baseline (the bound per-group ratios — cone u-block
0.01 vs collision Q-block 1.0+ — are preserved), updated once per SQP
iteration from the inner loop's final residuals: adapt when
r_prim/r_dual is imbalanced by >5x, step by sqrt(ratio), clamp to
[1e-2, 1e2] (OSQP's rule). The dual form is unscaled, so rho changes
need no y-rescaling; the rho*G'G fold refreshes each SQP iteration.
The scale persists across solves (warm rho, like the (z, y) dual warm
start); toggling resets it to 1. Telemetry: get_admm_rho_scale().
Measured defaults, rulings and provenance: docs/constraints.md (`set_admm_rho_adaptation`)."""
self.solver.set_admm_rho_adaptation(bool(on))
[docs]
def get_admm_rho_scale(self):
"""Per-solve adapted rho scale, shape (B,). Effective ADMM rho of a
group = its rho baseline * this scale (1.0 = unadapted)."""
return np.asarray(self.solver.get_admm_rho_scale())
[docs]
def add_lin_u_rows(self, C, d=None, lo=None, hi=None, mech=None, rho=None,
delta=0.05, sigma=0.0, cone=False, knot_lo=0,
knot_hi=None, admm_iters=0, equilibrate=False,
normalize=True):
"""Append a LIN_U row-group: m rows ``g = C @ u + d`` on the control
block (C shape (m, nu), FROZEN at a host-chosen configuration — the
cross-term audit's contact-frame rule for config-dependent maps).
Measured defaults, rulings and provenance: docs/constraints.md (`add_lin_u_rows`)."""
C = np.ascontiguousarray(np.asarray(C, dtype=np.float32))
if C.ndim != 2 or C.shape[1] != self.nu:
raise ValueError(f"C must be (m, {self.nu}); got {C.shape}")
m = C.shape[0]
mech, rho = self._resolve_mech_rho(mech, rho, admm_rho=0.01)
# ascontiguousarray: the binding consumes raw .ptr buffers — a strided view
# (e.g. np.broadcast_to) must be densified here (belt; the binding also
# forces c_style since 2026-08-02)
d = np.ascontiguousarray(np.asarray([] if d is None else d,
dtype=np.float32).reshape(-1))
if cone:
if normalize:
s = float(np.linalg.norm(C.astype(np.float64), 2))
if s > 0.0:
C = np.ascontiguousarray((C.astype(np.float64) / s).astype(np.float32))
d = np.ascontiguousarray((d.astype(np.float64) / s).astype(np.float32))
lo_a = np.asarray([], dtype=np.float32)
hi_a = np.asarray([], dtype=np.float32)
else:
if lo is None or hi is None:
raise ValueError("interval LIN_U rows need lo and hi (length m)")
lo_a = np.ascontiguousarray(np.asarray(lo, dtype=np.float32).reshape(m))
hi_a = np.ascontiguousarray(np.asarray(hi, dtype=np.float32).reshape(m))
if knot_hi is None:
knot_hi = self.N - 1 # no terminal control
self.solver.add_lin_u_group(self._MECHS[mech], C, d, lo_a, hi_a,
bool(cone), float(rho), float(delta),
float(sigma), int(knot_lo), int(knot_hi),
int(admm_iters), bool(equilibrate))
self._n_appended_groups += 1
return len(self.get_row_groups()) - 1 # the new group's index (set_row_group_mask/bounds/soft)
[docs]
def fc_slots(self, frame, part=None):
"""fc slot indices of one contact frame (index or baked frame name):
the 6-wide block [n(3); f(3)] (world-aligned, about the frame origin), or
just its ``"n"`` (moment) / ``"f"`` (force) half. Feeds add_fc_box /
set_fc_ref / SolveResult.fc_at slicing on multi-contact (fc-on-feet) plants."""
self._require_fc("fc_slots")
i = self.contact_frames.index(frame) if isinstance(frame, str) else int(frame)
if i < 0 or i >= len(self.contact_frames):
raise ValueError(f"contact frame {frame!r} not in {self.contact_frames}")
base = 6 * i
if part is None:
return list(range(base, base + 6))
if part == "n":
return list(range(base, base + 3))
if part == "f":
return list(range(base + 3, base + 6))
raise ValueError(f"part must be None, 'n' or 'f', got {part!r}")
[docs]
def add_fc_box(self, lo, hi, slots=None, **kw):
"""Box rows on contact-force slots (GATO_CONTACT_FORCES builds only):
selection LIN_U rows on control columns n_actuated+slots. ``slots``
indexes into the fc block (default: all n_fc slots); lo/hi broadcast.
Pin the wrench torque rows of a point contact with
``add_fc_box(0, 0, slots=range(3))`` (wrench layout is [n; f]).
Extra kwargs go to add_lin_u_rows (mech/rho/knot range/...)."""
self._require_fc("add_fc_box")
slots = list(range(self.n_fc)) if slots is None else list(slots)
if any(s < 0 or s >= self.n_fc for s in slots):
raise ValueError(f"fc slots must be in [0, {self.n_fc}); got {slots}")
m = len(slots)
C = np.zeros((m, self.nu), dtype=np.float32)
for i, s in enumerate(slots):
C[i, self.n_actuated + s] = 1.0
return self.add_lin_u_rows(C, lo=self._bcast(lo, m), hi=self._bcast(hi, m), **kw)
[docs]
def add_fc_cone(self, frame, mu, mech=None, rho=None, form="soc", **kw):
"""Friction cone on one contact frame's fc FORCE slots (fc builds):
rows [mu*f_z; f_x; f_y] on ``fc_slots(frame, "f")`` — the wrench is
world-aligned, so the map is CONSTANT (no freezing at q, unlike the arm
EE cone) and holds at every knot. Mask it to stance knots with
``set_row_group_mask``. Returns the group index."""
self._require_fc("add_fc_cone")
if not mu > 0:
raise ValueError("mu must be > 0")
fx, fy, fz = self.fc_slots(frame, "f")
C = np.zeros((3, self.nu), dtype=np.float64)
C[0, self.n_actuated + fz] = float(mu)
C[1, self.n_actuated + fx] = 1.0
C[2, self.n_actuated + fy] = 1.0
return self.enable_u_cone(C, mech=mech, rho=rho, form=form, **kw)
# ---- contact-frame POSITION rows ------------------------------
@property
def n_contact_rows(self):
"""3 rows per baked contact frame (the CONTACT_POS group width)."""
return 3 * len(self.contact_frames)
def _contact_targets(self, targets):
n = self.n_contact_rows
if targets is None:
return np.zeros((self.N, n), dtype=np.float32)
t = np.asarray(targets, dtype=np.float32)
if t.shape == (n,) or t.shape == (len(self.contact_frames), 3):
t = np.tile(t.reshape(1, n), (self.N, 1))
return np.ascontiguousarray(t.reshape(self.N, n))
[docs]
def set_row_group_targets(self, g, targets):
"""Per-knot target table of CONTACT_POS group ``g``: (N, n_rows),
(n_frames, 3) or (n_rows,) (broadcast over knots). ADMM z/y re-init
on the next solve (the residual moved)."""
groups = self.get_row_groups()
if g < 0 or g >= len(groups) or int(groups[g]["kind"]) != KIND_CONTACT_POS:
raise ValueError(f"row group {g} is not a CONTACT_POS group")
self.solver.set_row_group_targets(int(g), self._contact_targets(targets))
[docs]
def enable_u_cone(self, C, d=None, mech=None, rho=None, form="soc",
facets=8, facet_scale="inscribed", **kw):
"""Cone constraint on a mapped control quantity g = C @ u + d
.T), rows [mu*f_n; f_t1; f_t2], frozen at q).
Measured defaults, rulings and provenance: docs/constraints.md (`enable_u_cone`)."""
C = np.asarray(C, dtype=np.float64)
m = C.shape[0]
d = np.zeros(m) if d is None else np.asarray(d, dtype=np.float64).reshape(m)
gi = len(self.get_row_groups())
if form == "soc":
self.add_lin_u_rows(C, d, mech=mech, rho=rho, cone=True, **kw)
elif form == "pyramid":
if m != 3:
raise ValueError("pyramid form supports 3-row cones (t, x, y)")
s = np.cos(np.pi / facets) if facet_scale == "inscribed" else 1.0
th = 2.0 * np.pi * np.arange(facets) / facets
F = np.stack([np.cos(th), np.sin(th), -s * np.ones(facets)], axis=1) # rows: [c, s, -s0] on (g1, g2, g0)
P = F @ C[[1, 2, 0], :] # facet map on u
pd = F @ d[[1, 2, 0]] # facet offsets
self.add_lin_u_rows(P, pd, lo=np.full(facets, -np.inf),
hi=np.zeros(facets), mech=mech, rho=rho, **kw)
else:
raise ValueError(f"form must be 'soc' or 'pyramid', got {form!r}")
return gi
[docs]
def set_collision_environment(self, spheres=None, capsules=None,
cuboids=None, planes=None):
"""Upload the runtime obstacle set for the COLLISION clearance rows
. Lists of tuples, one per obstacle (all in world frame, meters):
Measured defaults, rulings and provenance: docs/constraints.md (`set_collision_environment`)."""
def arr(x, w):
a = np.ascontiguousarray(np.asarray([] if x is None else x, dtype=np.float32))
return a.reshape(-1, w) if a.size else np.zeros((0, w), dtype=np.float32)
self.solver.set_collision_environment(arr(spheres, 4), arr(capsules, 7),
arr(cuboids, 15), arr(planes, 4))
[docs]
def enable_collision(self, mech=None, margin=0.02, rho=None, delta=0.05,
sigma=0.0, knot_lo=1, admm_iters=0):
"""Append THE collision clearance group (one max): per-sphere rows
d_i(q_k) >= margin over knots [knot_lo, N] — d_i = signed clearance of
collision sphere i (baked at codegen, ``collision_res``) to the
nearest obstacle from set_collision_environment (call that FIRST).
The covering spheres are already conservative (inflated by the
spherizer), so margin is extra safety on top.
Measured defaults, rulings and provenance: docs/constraints.md (`enable_collision`)."""
mech, rho = self._resolve_mech_rho(mech, rho, admm_rho=1.0) # admm 1.0: clearance rows fold onto the Q block
self.solver.enable_collision(self._MECHS[mech], float(margin), float(rho),
float(delta), float(sigma), int(knot_lo),
int(admm_iters))
self._n_appended_groups += 1
[docs]
def get_collision_row_duals(self):
"""COLLISION-band AL duals: dict of (B, N, n_spheres) arrays
(lam_hi unused — the rows are one-sided; lam_lo is the hinge dual)."""
return self.solver.get_collision_row_duals()
[docs]
def get_collision_admm_state(self):
"""COLLISION-band ADMM (z, y): dict of (B, N, n_spheres) arrays."""
return self.solver.get_collision_admm_state()
[docs]
def set_cost_weights(self, q_cost=None, qd_cost=None, u_cost=None, N_cost=None,
q_lim_cost=None, vel_lim_cost=None, ctrl_lim_cost=None):
"""Update scalar cost weights at runtime (None keeps the current value)."""
w = self._cost_weights
for name, val in (("q_cost", q_cost), ("qd_cost", qd_cost), ("u_cost", u_cost),
("N_cost", N_cost), ("q_lim_cost", q_lim_cost),
("vel_lim_cost", vel_lim_cost), ("ctrl_lim_cost", ctrl_lim_cost)):
if val is not None:
w[name] = float(val)
self.solver.set_cost_weights(w["q_cost"], w["qd_cost"], w["u_cost"], w["N_cost"],
w["q_lim_cost"], w["vel_lim_cost"], w["ctrl_lim_cost"])
self.params = self.params.replace(**w) # self.params stays truthful
[docs]
def set_q_pos_cost(self, weight):
"""Joint-posture nullspace anchor: adds 0.5*weight*||q - q_nom||^2 as a RUNNING
cost on the q-block (default 0 = the historic EE-only cost, bitwise-off).
The EE-position cost leaves a nullspace on redundant arms (iiwa14: wrist roll
j7 has a ~zero EE-position Jacobian column); by default only the q-limit
barrier anchors it. Where the operating posture sits near/outside the
margin-shrunk limits (barrier limits = URDF limits - 0.1 rad, plant.cuh
JOINT_LIMIT_MARGIN), set q_lim_cost=0 and use a small q_pos_cost toward
set_q_nom(x0[:nq]) instead — anchors without the barrier field.
``weight`` may be a scalar or a length-nq array (per-joint anchor
stiffness — tune light joints independently; see set_u_cost_vec for the
matching effort-side knob and the closed-loop rate-limit story).
"""
w = np.asarray(weight, dtype=np.float32)
if w.ndim == 0:
self.solver.set_q_pos_cost_vec(np.empty(0, dtype=np.float32)) # back to scalar
self.solver.set_q_pos_cost(float(w))
else:
self.solver.set_q_pos_cost_vec(np.ascontiguousarray(w.reshape(-1)))
[docs]
def set_u_cost_vec(self, weights=None):
"""Per-joint control effort weights (length n_actuated; None resets to the
scalar u_cost). The per-joint knob that makes a joint's commanded
correction respect a discrete control loop: with effort nearly free
(u_cost=1e-6 default) the optimal posture/velocity correction is
deadbeat-aggressive, and on a near-massless joint (iiwa14 j7,
J_eff ~ 1e-3 kg m^2) a plant-model inertia mismatch turns the 100 Hz
loop into a growing Nyquist oscillation (PDDP round-5, 2026-08-02).
Raising ONLY that joint's effort weight (e.g. [1e-6]*6 + [3e-3]) softens
its channel without degrading the arm's tracking."""
if weights is None:
self.solver.set_u_cost_vec(np.empty(0, dtype=np.float32))
else:
w = np.ascontiguousarray(np.asarray(weights, dtype=np.float32).reshape(-1))
self.solver.set_u_cost_vec(w)
[docs]
def set_q_nom(self, q_nom=None):
"""Posture target for set_q_pos_cost (length-nq array; None resets to zeros)."""
if q_nom is None:
self.solver.set_q_nom(np.empty(0, dtype=np.float32))
else:
q = np.ascontiguousarray(np.asarray(q_nom, dtype=np.float32).ravel())
self.solver.set_q_nom(q)
[docs]
def set_fc_cost(self, weight):
"""Contact-force regularization 0.5*weight*||f_c||^2 per knot (GATO_CONTACT_FORCES
modules only — the fc slots appended to every control; inert on default builds).
Default on fc builds is 1e-2, NOT 0: unregularized fc is a free 6-DoF wrench
actuator and the solve destabilizes (1e-3 is already marginal — measured
2026-08-02). Use large weights (~1e6) to effectively pin fc to zero, or 0
only together with fc box rows (add_fc_box)."""
self.solver.set_fc_cost(float(weight))
[docs]
def set_fc_ref(self, ref=None):
"""Contact-wrench reference for the fc regularization (GATO_CONTACT_FORCES
modules only): the fc cost becomes 0.5*fc_cost*||f_c - ref||^2 per running
knot. ``ref`` has n_fc entries (wrench layout [n; f] per contact frame,
world-aligned at the baked frame), shared across knots and batch rows.
None/empty resets to zeros — bitwise the historic pure regularization.
This is how a force SETPOINT enters the solve: ref = the desired contact
wrench (e.g. [0,0,0, 0,0,+F] for an F-newton press reaction), fc_cost =
the force-tracking weight. Pair with cone rows on the fc columns
(add_lin_u_rows / enable_u_cone with a selection C) for friction limits."""
self._require_fc("set_fc_ref")
if ref is None:
self.solver.set_fc_ref(np.empty(0, dtype=np.float32))
return
r = np.asarray(ref, dtype=np.float32)
if r.ndim == 2 and r.shape == (self.N, self.n_fc): # per-knot reference (gait scheduling)
self.solver.set_fc_ref(np.ascontiguousarray(r).ravel())
return
r = np.ascontiguousarray(r.ravel())
if r.size != self.n_fc:
raise ValueError(f"fc_ref must have n_fc = {self.n_fc} entries or shape (N, n_fc) = "
f"({self.N}, {self.n_fc}); got {np.asarray(ref).shape}")
self.solver.set_fc_ref(r)
[docs]
def set_cost_weights_per_knot(self, knot_weights):
"""Per-knot [ee, qd, u] weight triples, shape (N, 3): overrides the scalar
q/qd/u/N weights (terminal EE weight = row N-1's ee entry). Enables
via-points, terminal-only goals, and horizon masking at runtime.
``None`` clears the table (back to the scalar weights)."""
if knot_weights is None:
self.solver.clear_cost_weights_per_knot()
return
w = np.ascontiguousarray(np.asarray(knot_weights, dtype=np.float32)).reshape(self.N, 3)
self.solver.set_cost_weights_per_knot(w)
# ---- lazy pinocchio model (FK helpers / sim / estimators / tests) ----
@property
def model(self):
"""pinocchio Model for model_path (built on first use; needs the [test] extra)."""
if self._model is None:
pin = _require_pin()
if self.floating_base:
# floating-base modules: the stored state carries the
# free-flyer q [p; quat xyzw] — mirror it in the pin model so
# nq/nv/nx and every oracle computation match the module.
self._model = pin.buildModelFromUrdf(self.model_path, pin.JointModelFreeFlyer())
else:
self._model = pin.buildModelFromUrdf(self.model_path)
return self._model
@property
def data(self):
if self._data is None:
self._data = self.model.createData()
return self._data
@property
def ee_frame_id(self):
"""Frame id of the codegen'd EE target frame, or None if the URDF lacks it."""
if self._ee_frame_id is None:
m = self.model
self._ee_frame_id = m.getFrameId(self.ee_frame) if m.existFrame(self.ee_frame) else -1
return None if self._ee_frame_id == -1 else self._ee_frame_id
[docs]
def ee_pos(self, q, frame="ee"):
"""Frame position via pinocchio FK.
frame="ee": the URDF ee_frame (fixed-joint child, e.g. tcp). Since
GRiD e31f7bd the device FK (tracking cost AND EE row-groups) uses the
named-target ``*_EE`` codegen, which INCLUDES the terminal fixed-joint
origin — device == this frame to f32 precision (~1e-7).
frame="solver": historical alias for the device frame; now identical
to "ee" (the old dropped-origin convention is gone upstream).
Any other string is a URDF frame name (e.g. a baked contact frame from
``contact_frames``): its origin, the point the fc wrench acts about.
"""
pin = _require_pin()
pin.forwardKinematics(self.model, self.data, q)
if frame in ("ee", "solver"):
fid = self.ee_frame_id
if fid is None:
return np.array(self.data.oMi[self.model.njoints - 1].translation)
else:
if not self.model.existFrame(frame):
raise ValueError(f"frame {frame!r} is not in the URDF {self.model_path!r}")
fid = self.model.getFrameId(frame)
pin.updateFramePlacement(self.model, self.data, fid)
return np.array(self.data.oMf[fid].translation)
[docs]
def reset(self):
"""Clear all solver state carried across solves: duals (AL/ADMM), the
adapted trust-region rho, and the external-force buffer. (Trajectory
warm starts are the caller's — see solve().)"""
self.reset_dual()
self.reset_rho() # adapted rho is solver state -> a full reset must clear it too
self.set_f_ext_B(np.zeros((self.batch_size, 6)))
[docs]
def sim_forward(self, xk, uk, sim_dt):
"""One plant step on the DEVICE integrator for every batch entry:
(B, nx) x (B, nu) -> (B, nx) (knot-0 external wrench applies)."""
xk = np.asarray(xk, dtype=np.float32)
uk = np.asarray(uk, dtype=np.float32)
return self.solver.sim_forward(xk, uk, sim_dt)
[docs]
def set_f_ext_B(self, f_ext_B):
"""Upload the external-wrench band the dynamics see.
The GPU buffer is body-major: 6*NUM_BODIES per (solve, knot); each 6-slot is
that body's spatial force in its JOINT-LOCAL frame about the joint origin,
Featherstone-ordered [angular(3); linear(3)] (verified vs pin.aba 2026-07-07).
World wrenches go through gato.common.world_wrench_to_joint_local and are
reordered — see hypotheses.world_wrench_to_gato_slot.
Accepted shapes (N = knot count; wrench k applies to interval [k, k+1]):
- ``(B, 6)``: per-solve EE wrench, broadcast over knots
- ``(B, 6*NUM_BODIES)``: per-solve body-major, broadcast over knots
- ``(B, N, 6)``: per-knot EE wrench (scattered into the EE body slot)
- ``(B, N, 6*NUM_BODIES)``: per-knot body-major
A correctly-sized contiguous buffer is always uploaded (a short buffer
would make the device copy over-read -> garbage wrench -> NaN dynamics)."""
f_ext_B = np.asarray(f_ext_B, dtype=np.float32)
nb = self.n_bodies or self.nv
per_knot = (f_ext_B.ndim == 3)
if per_knot:
if f_ext_B.shape[:2] != (self.batch_size, self.N):
raise ValueError(
f"per-knot f_ext_B must be (batch, N, ...) = ({self.batch_size}, "
f"{self.N}, ...); got {f_ext_B.shape}"
)
width = f_ext_B.shape[2]
body_major = np.zeros((self.batch_size, self.N, 6 * nb), dtype=np.float32)
else:
f_ext_B = f_ext_B.reshape(self.batch_size, -1)
width = f_ext_B.shape[1]
body_major = np.zeros((self.batch_size, 6 * nb), dtype=np.float32)
self.f_ext_B = f_ext_B
if width == 6:
ee = 6 * (nb - 1) # end-effector body slot
body_major[..., ee:ee + 6] = f_ext_B
elif width == 6 * nb:
body_major[...] = f_ext_B
else:
raise ValueError(
f"f_ext_B must have width 6 (EE wrench) or {6 * nb} (body-major "
f"6*NUM_BODIES); got {width}"
)
if per_knot:
self.solver.set_f_ext_knot_batch(np.ascontiguousarray(body_major))
else:
self.solver.set_f_ext_batch(np.ascontiguousarray(body_major))
[docs]
def reset_rho(self):
self.solver.reset_rho()
[docs]
def reset_dual(self):
self.solver.reset_dual()