Components
MJINX employs a component-based architecture to formulate inverse kinematics problems through functional decomposition. Each component encapsulates a mathematical mapping from the configuration space to a task or constraint space, functioning either as an objective function (task) or an inequality constraint (barrier).
This modular structure facilitates the systematic construction of complex kinematic problems through composition of elementary components. The approach aligns with established practices in robotics control theory, where complex behaviors emerge from the coordination of simpler control objectives.
Components are characterized by several key attributes:
A unique identifier for reference within the problem formulation
A gain parameter that determines its relative weight in the optimization
An optional dimensional mask for selective application
A differentiable function mapping robot state to output values
This formulation follows the task-priority framework established in robotics literature, where multiple objectives are managed through appropriate weighting and prioritization. The separation of concerns between tasks and constraints provides a natural expression of both the desired behavior and the feasible region of operation.
When integrated into a Problem instance, components form a well-posed optimization problem. Tasks define the objective function to be minimized, while barriers establish the constraint manifold. The solver then computes solutions that optimize the weighted task objectives while maintaining feasibility with respect to all constraints.
Base Component
The Component class serves as the abstract base class from which all specific component implementations derive. This inheritance hierarchy ensures a consistent interface while enabling specialized behavior for different component types.
- class JaxComponent(dim, model, vector_gain, gain_fn, mask_idxs)
Base class for all JAX-based components in the inverse kinematics framework.
This class provides the fundamental structure for components that can be evaluated and differentiated within the JAX ecosystem. Components represent either tasks (objectives) or barriers (constraints) in the optimization problem.
- Parameters:
gain – The gain factor applied to the component.
gain_fn (
Callable[[float],float]) – A function to compute the gain dynamically.mask – A sequence of integers to mask certain dimensions.
- abstract __call__(data)
Compute the component’s value \(f(q, t)\).
This method should be implemented by subclasses to provide specific component calculations.
- Parameters:
data (
Data) – The MuJoCo simulation data.- Return type:
- Returns:
The computed component value.
- copy_and_set(**kwargs)
Create a copy of the component with updated attributes.
- Parameters:
kwargs – Keyword arguments specifying the attributes to update.
- Return type:
- Returns:
A new instance of the component with updated attributes.
- compute_jacobian(data)
Compute the Jacobian of the component with respect to the joint positions. The jacobian is a matrix \(J \in R^{dim \times nv}\) that is defined as:
\[J(q, t) = \frac{\partial f}{\partial q}\]The main property that is desirable for us is that given this matrix, we obtain component derivative as a linear matrix w.r.t. velocities:
\[\dot{f}(q, t) = \frac{\partial f}{\partial q} v = J(q, t)\]The jacobian is calculated via automatic differentiation, if it is not overwritten. If \(nq \\neq nv\), a special mapping is computed to transform \(J_{dq}\) into \(J_v\). For details, see
mjinx.configuration.jac_dq2v().- Parameters:
data (
Data) – The MuJoCo simulation data.- Return type:
- Returns:
The computed Jacobian matrix.
- class Component(name, gain, gain_fn=None, mask=None)
High-level interface for components in the inverse kinematics framework.
This class provides a Python-friendly interface for creating and manipulating components, which are then compiled into JAX-compatible representations for efficient computation.
- Parameters:
name (
str) – The name of the component.gain (
mjinx.typing.ArrayOrFloat) – The gain factor applied to the component.gain_fn (
Callable[[float],float] |None) – A function to compute the gain dynamically.mask (
Sequence[int] |None) – A sequence of integers to mask certain dimensions.
- property modified: bool
Check if the component has been modified.
- Returns:
True if the component has been modified, False otherwise.
- property model: Model
Get the MuJoCo model associated with the component.
- Returns:
The MuJoCo model.
- update_model(model)
Update the MuJoCo model for the component.
- Parameters:
model (
Model) – The new MuJoCo model.
- update_gain(gain)
Update the gain for the component.
- Parameters:
gain (
mjinx.typing.ArrayOrFloat) – The new gain value.- Raises:
ValueError – If the gain has an invalid dimension.
- property vector_gain: Array
Get the vector gain for the component.
- The vector gain is generated as follows:
gain is scalar -> vector_gain = jnp.ones(self.dim) * gain
gain is vector -> vector_gain = gain
- Returns:
The vector gain array.
- Raises:
ValueError – If the dimension is not set or the gain size is invalid.
- property gain_fn: Callable[[float], float]
Get the gain function for the component.
- Returns:
The gain function.
- property name: str
Get the name of the component.
- Returns:
The component name.
- property dim: int
Get the dimension of the component.
- Returns:
The component dimension.
- Raises:
ValueError – If the dimension is not set.
- property mask: Array
Get the mask for the component.
- Returns:
The mask array.
- Raises:
ValueError – If the mask is not set and the dimension is not defined.
- property mask_idxs: tuple[int, ...]
Get the mask indices for the component.
- Returns:
A tuple of mask indices.
- Raises:
ValueError – If the mask is not set and the dimension is not defined.
- property jax_component: AtomicComponentType
Get the JAX implementation of the component.
- Returns:
The JAX component instance.
- Raises:
ValueError – If the model or dimension is not set.
Tasks
Tasks define objective functions that map from configuration space to task space, with the solver minimizing weighted errors between current and desired values [Kanoun et al., 2011]. Each task \(f: \mathcal{Q} \rightarrow \mathbb{R}^m\) produces an error \(e(q) = f(q) - f_{desired}\) that is minimized according to \(\|e(q)\|^2_W\).
MJINX provides task implementations for common robotics objectives:
Barriers
Barriers implement inequality constraints through scalar functions \(h(q) \geq 0\) that create boundaries the solver must respect. Based on control barrier functions (CBFs) [Ames et al., 2019], these constraints are enforced through differential inequality: \(\nabla h(q)^T v \geq -\alpha h(q)\), with \(\alpha\) controls constraint enforcement and \(v\) is the velocity vector.
MJINX provides several barrier implementations:
Constraints
Constraints represent a new type of component with a strictly enforced equality condition \(f(q) = 0\). Those constraints might be either treated strictly as differentiated exponentially stable equality: \(\nabla h(q)^T v = -\alpha h(q)\), with \(\alpha\) controls constraint enforcement and \(v\) is the velocity vector, or as a soft constraint – task with high gain.
Yet, only the following constraints are implemented: