Solvers

MJINX provides multiple solver implementations for inverse kinematics problems, each with different characteristics suitable for various applications.

Base Solver

The abstract base class defining the interface for all solvers.

class SolverData

Bases: object

Base class for solver-specific data.

This class serves as a placeholder for any data that a specific solver needs to maintain between iterations or function calls. It enables solver implementations to preserve state information and warm-start subsequent optimization steps.

class SolverSolution(v_opt)

Bases: object

Base class for solver solutions.

This class provides the structure for returning optimization results. It contains the optimal velocity solution and can be extended by specific solvers to include additional solution information.

Parameters:: v_opt (Array) – Optimal velocity solution.

v_opt: Array

class Solver(model)

Bases: Generic[SolverDataType, SolverSolutionType], ABC

Abstract base class for solvers.

This class defines the interface for solvers used in inverse kinematics problems. Solvers transform task and barrier constraints into optimization problems, which can be solved to find joint configurations or velocities that satisfy the constraints.

In general, the optimization problem can be formulated as:

\[\min_{q} \sum_{i} \|e_i(q)\|^2_{W_i} \quad \text{subject to} \quad h_j(q) \geq 0\]

where:

\(e_i(q)\) are the task errors
\(W_i\) are the task weight matrices
\(h_j(q)\) are the barrier constraints

Different solver implementations use different approaches to solve this problem, such as local linearization (QP) or global nonlinear optimization.

Parameters:: model (Model) – The MuJoCo model used by the solver.

model: Model

abstract solve_from_data(solver_data, problem_data, model_data)

Solve the inverse kinematics problem using pre-computed model data.

Parameters:

solver_data (TypeVar(SolverDataType, bound= SolverData)) – Solver-specific data.
problem_data (JaxProblemData) – Problem-specific data.
model_data (Data) – MuJoCo model data.

Return type:

tuple[TypeVar(SolverSolutionType, bound= SolverSolution), TypeVar(SolverDataType, bound= SolverData)]

Returns:

A tuple containing the solver solution and updated solver data.

solve(q, solver_data, problem_data)

Solve the inverse kinematics problem for a given configuration.

This method creates mjx.Data instance and updates it under the hood. To avoid doing an extra update, consider solve_from_data method.

Parameters:

q (Array) – The current joint configuration.
solver_data (TypeVar(SolverDataType, bound= SolverData)) – Solver-specific data.
problem_data (JaxProblemData) – Problem-specific data.

Return type:

tuple[TypeVar(SolverSolutionType, bound= SolverSolution), TypeVar(SolverDataType, bound= SolverData)]

Returns:

A tuple containing the solver solution and updated solver data.

Raises:

ValueError – If the input configuration has incorrect dimensions.

abstract init(q)

Initialize solver-specific data.

Parameters:: q (ndarray | Array) – The initial joint configuration.
Return type:: TypeVar(SolverDataType, bound= SolverData)
Returns:: Initialized solver-specific data.

Local IK Solver

A Quadratic Programming (QP) based solver that linearizes the problem at each step. This solver is efficient for real-time control and tracking applications.

Build and solve the inverse kinematics problem.

class OSQPParameters

Bases: TypedDict

Parameters for configuring the OSQP solver.

Parameters:

check_primal_dual_infeasability – If True, populates the status field of state with one of BoxOSQP.PRIMAL_INFEASIBLE, BoxOSQP.DUAL_INFEASIBLE (default: True). If False, improves speed but does not check feasibility. When jit=False and the problem is primal or dual infeasible, a ValueError is raised.
sigma – Ridge regularization parameter for stabilizing the linear system solution (default: 1e-6).
momentum – Relaxation parameter (default: 1.6). Must be in the open interval (0, 2). A value of 1 means no relaxation, less than 1 means under-relaxation, and greater than 1 means over-relaxation. For best results, choose momentum in the range [1.5, 1.8].
eq_qp_solve – Method for solving equality-constrained QP subproblems. Options are ‘cg’ (conjugate gradient), ‘cg+jacobi’ (with Jacobi preconditioning), and ‘lu’ (direct solving via LU factorization). Default is ‘cg’.
rho_start – Initial step size for the primal-dual algorithm (default: 1e-1).
rho_min – Minimum allowed step size to prevent excessively small steps (default: 1e-6).
rho_max – Maximum allowed step size to prevent overly large steps (default: 1e6).
stepsize_updates_frequency – Number of iterations between step size adjustments (default: 10).
primal_infeasible_tol – Tolerance for detecting primal infeasibility (default: 1e-4).
dual_infeasible_tol – Tolerance for detecting dual infeasibility (default: 1e-4).
maxiter – Maximum number of iterations allowed before termination (default: 4000).
tol – Absolute tolerance for the stopping criterion (default: 1e-3).
termination_check_frequency – Number of iterations between convergence checks (default: 5).
implicit_diff_solve – Solver for linear systems in implicit differentiation. Must be a callable that solves Ax = b for x.

check_primal_dual_infeasability: str | bool

sigma: float

momentum: float

eq_qp_solve: str

rho_start: float

rho_min: float

rho_max: float

stepsize_updates_frequency: int

primal_infeasible_tol: float

dual_infeasible_tol: float

maxiter: int

tol: float

termination_check_frequency: int

implicit_diff_solve: Callable

class LocalIKData(v_prev)

Bases: SolverData

Data class for the Local Inverse Kinematics solver.

Parameters:: v_prev (Array) – The previous velocity solution.

v_prev: Array

class LocalIKSolution(v_opt, dual_var_eq, dual_var_ineq, iterations, error, status)

Bases: SolverSolution

Solution class for the Local Inverse Kinematics solver.

This class extends the base SolverSolution to include additional information specific to the Local IK solver’s quadratic programming solution.

Parameters:

v_opt (Array) – The optimal velocity solution.
dual_var_eq (Array) – Dual variables for equality constraints.
dual_var_ineq (Array) – Dual variables for inequality constraints.
iterations (int) – Number of iterations performed by the solver.
error (float) – Final error value at convergence.
status (int) – Solver status code indicating outcome (success, infeasible, etc.).

dual_var_eq: Array

dual_var_ineq: Array

iterations: int

error: float

status: int

class LocalIKSolver(model, **kwargs)

Bases: Solver[LocalIKData, LocalIKSolution]

Local Inverse Kinematics solver using Quadratic Programming (QP).

This solver uses a local linearization approach to solve the inverse kinematics problem. At each step, it formulates a Quadratic Program (QP) that approximates the nonlinear optimization problem and solves for joint velocities that minimize task errors while respecting constraints.

The QP is formulated as:

\[\min_{v} \frac{1}{2} v^T P v + c^T v \quad \text{subject to} \quad G v \leq h\]

where:

\(v\) is the vector of joint velocities
\(P\) is a positive-definite matrix constructed from task Jacobians
\(c\) is a vector constructed from task errors
\(G\) and \(h\) encode barrier constraints and velocity limits

For tasks, the contribution to the objective is:

\[\begin{split}P_{task} &= J^T W J \\ c_{task} &= -J^T W e\end{split}\]

where:

\(J\) is the task Jacobian matrix (∂f/∂q)
\(W\) is the task weight matrix (cost)
\(e\) is the task error (f(q) - f_desired)

For barriers, the constraints are linearized as:

\[\begin{split}G_{barrier} &= -J_b \\ h_{barrier} &= \alpha h(q)\end{split}\]

where:

\(J_b\) is the barrier Jacobian (∂h/∂q)
\(h(q)\) is the barrier function value
\(\alpha\) is a gain parameter that controls constraint relaxation

Additionally, velocity limits are incorporated as:

\[\begin{split}G_{limits} &= \begin{bmatrix} I \\ -I \end{bmatrix} \\ h_{limits} &= \begin{bmatrix} v_{max} \\ -v_{min} \end{bmatrix}\end{split}\]

The solver also includes a safe displacement term to push away from constraint boundaries:

\[\begin{split}P_{safe} &= \beta I \\ c_{safe} &= -\beta v_{safe}\end{split}\]

where \(v_{safe}\) is a velocity that pushes away from constraint boundaries.

The complete QP matrices are assembled by combining these components:

\[\begin{split}P &= \sum_i P_{task,i} + P_{safe} \\ c &= \sum_i c_{task,i} + c_{safe} \\ G &= \begin{bmatrix} G_{barrier} \\ G_{limits} \end{bmatrix} \\ h &= \begin{bmatrix} h_{barrier} \\ h_{limits} \end{bmatrix}\end{split}\]

The QP is solved using the OSQP solver, which implements an efficient primal-dual interior point method specifically designed for convex quadratic programs.

Parameters:

model (Model) – The MuJoCo model.
dt – The time step for integration.
osqp_params – Parameters for the OSQP solver.

solve_from_data(solver_data, problem_data, model_data)

Solve the Local IK problem using pre-computed data.

Parameters:

solver_data (LocalIKData) – The solver-specific data.
problem_data (JaxProblemData) – The problem-specific data.
model_data (Data) – The MuJoCo model data.

Return type:

tuple[LocalIKSolution, LocalIKData]

Returns:

A tuple containing the solver solution and updated solver data.

init(v_init=None)

Initialize the Local IK solver data.

Parameters:: v_init (ndarray | Array | None) – The initial velocity. If None, zero velocity is used.
Return type:: LocalIKData
Returns:: Initialized solver-specific data.
Raises:: ValueError – If the input velocity has incorrect dimensions.

Global IK Solver

A nonlinear optimization solver that directly optimizes joint positions. This solver can find solutions that avoid local minima and is suitable for complex positioning tasks.

class GlobalIKData(optax_state)

Bases: SolverData

Data class for the Global Inverse Kinematics solver.

Parameters:: optax_state (Array | ndarray | bool | number | Iterable[ArrayTree] | Mapping[Any, ArrayTree]) – The state of the Optax optimizer.

optax_state: Array | ndarray | bool | number | Iterable[ArrayTree] | Mapping[Any, ArrayTree]

class GlobalIKSolution(v_opt, q_opt)

Bases: SolverSolution

Solution class for the Global Inverse Kinematics solver.

Parameters:

q_opt (Array) – The optimal joint configuration.
v_opt (Array) – The optimal joint velocities.

q_opt: Array

class GlobalIKSolver(model, optimizer, dt=0.01)

Bases: Solver[GlobalIKData, GlobalIKSolution]

Global Inverse Kinematics solver using gradient-based optimization.

This solver uses Optax for gradient-based optimization to solve the inverse kinematics problem globally. Unlike the local solver, it optimizes over joint positions directly rather than velocities, and can potentially find solutions that avoid local minima.

The optimization problem is formulated as:

\[\min_{q} \mathcal{L}(q) = \sum_{i} \|e_i(q)\|^2_{W_i} - \sum_{j} \lambda_j \log(h_j(q))\]

where:

\(e_i(q)\) are the task errors (f_i(q) - f_i_desired)
\(W_i\) are the task weight matrices
\(h_j(q)\) are the barrier constraints
\(\lambda_j\) are barrier gains
\(\log(h_j(q))\) is a logarithmic barrier that approaches infinity as \(h_j(q)\) approaches zero

The gradient of this loss function is:

\[\nabla_q \mathcal{L}(q) = \sum_{i} J_i^T W_i e_i(q) - \sum_{j} \lambda_j \frac{\nabla_q h_j(q)}{h_j(q)}\]

where:

\(J_i\) is the Jacobian of task i (∂e_i/∂q)
\(\nabla_q h_j(q)\) is the gradient of barrier j

The optimization is performed using Optax’s gradient-based optimizers, which update the joint positions iteratively:

\[q_{t+1} = q_t - \alpha_t \nabla_q \mathcal{L}(q_t)\]

where \(\alpha_t\) is the step size determined by the optimizer.

The logarithmic barrier function ensures that constraints are satisfied by creating an infinite penalty as the system approaches constraint boundaries. This allows the optimization to explore the full feasible space without explicit inequality constraints.

Compared to the local IK solver: 1. Global IK optimizes joint positions directly rather than velocities 2. It can potentially escape local minima through gradient-based exploration 3. It uses soft constraints via logarithmic barriers rather than hard inequality constraints 4. It typically requires more iterations but may find better solutions for complex problems

Parameters:

model (Model) – The MuJoCo model.
optimizer (GradientTransformation) – The Optax optimizer to use.
dt (float) – The time step for velocity computation.

loss_fn(q, problem_data)

Compute the loss function for the given joint configuration.

This function evaluates the total loss for the optimization problem:

\[\mathcal{L}(q) = \sum_{i} \|e_i(q)\|^2_{W_i} - \sum_{j} \lambda_j \log(h_j(q))\]

The first term represents the weighted sum of squared task errors, while the second term represents the logarithmic barrier penalties for constraint satisfaction. The logarithmic barrier approach transforms inequality constraints into penalty terms, creating a barrier that prevents the optimizer from violating the constraints.

Parameters:

q (Array) – The joint configuration.
problem_data (JaxProblemData) – The problem-specific data.

Return type:

float

Returns:

The computed loss value.

solve_from_data(solver_data, problem_data, model_data)

Solve the Global IK problem using pre-computed data.

Parameters:

solver_data (GlobalIKData) – The solver-specific data.
problem_data (JaxProblemData) – The problem-specific data.
model_data (Data) – The MuJoCo model data.

Return type:

tuple[GlobalIKSolution, GlobalIKData]

Returns:

A tuple containing the solver solution and updated solver data.

solve(q, solver_data, problem_data)

Solve the Global IK problem for a given configuration.

This method performs gradient-based optimization to find a configuration that minimizes the task errors while satisfying barrier constraints. Unlike the local solver, it operates directly on joint positions rather than velocities.

Parameters:

q (Array) – The current joint configuration.
solver_data (GlobalIKData) – The solver-specific data containing optimization state.
problem_data (JaxProblemData) – The problem-specific data containing tasks and barriers.

Return type:

tuple[GlobalIKSolution, GlobalIKData]

Returns:

A tuple containing the solver solution and updated solver data.

Raises:

ValueError – If the input configuration has incorrect dimensions.

init(q)

Initialize the Global IK solver data.

Parameters:: q (ndarray | Array) – The initial joint configuration.
Return type:: GlobalIKData
Returns:: Initialized solver-specific data.
Raises:: ValueError – If the input configuration has incorrect dimensions.