Constraints represent a new type of component with a strictly enforced equality, either hardly or softly.

The constratin is formulated as \(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.

Base Constraint

The foundational class that all constraints extend. It defines the core interface and mathematical properties for constraint objectives.

class JaxConstraint(dim, model, vector_gain, gain_fn, mask_idxs, active, hard_constraint, soft_constraint_cost)

A JAX-based representation of an equality constraint.

This class defines an equality constraint function f(x) = 0 for optimization tasks.

Parameters:

  • active (bool) – Indicates whether the constraint is active.

  • hard_constraint (bool) – A flag that specifies if the constraint is hard (True) or soft (False).

  • soft_constraint_cost (Array) – The cost matrix used for soft constraint relaxation.

compute_constraint(data)

Compute the equality constraint function value.

Evaluates the constraint function f(x) = 0 based on the current simulation data. For soft constraints, the computed value is later penalized by the associated cost; for hard constraints, the evaluation is directly used in the Ax = b formulation.

Parameters:

data (Data) – The MuJoCo simulation data.

Return type:

Array

Returns:

The computed constraint value.

class Constraint(name, gain, mask=None, hard_constraint=False, soft_constraint_cost=None)

A high-level component for formulating equality constraints.

This class wraps an atomic JAX constraint (JaxConstraint) and provides a framework to manage constraints within the optimization problem. Equality constraints are specified as:

\[f(x) = 0\]
They can be treated in one of two ways:

  • As a soft constraint. In this mode, the constraint violation is penalized by a cost (typically with a high gain), transforming the constraint into a task.

  • As a hard constraint. Here, the constraint is enforced as a strict equality in the following form:

\[\nabla h(q)^T v = -\alpha h(q),\]

where \(\alpha\) controls constraint enforcement and \(v\) is the velocity vector.

Parameters:

  • matrix_cost – The cost matrix associated with the task.

  • lm_damping – The Levenberg-Marquardt damping factor.

  • active – Determines if the constraint is active.

  • hard_constraint (bool) – Indicates whether the constraint is hard (True) or soft (False).

property soft_constraint_cost: Array

Get the cost matrix associated with a soft constraint.

For soft constraints, any violation is penalized by a cost if self._soft_constraint_cost is not None else 1e2 * jnp.eye(self.dim)”

Returns:

The cost matrix for the soft constraint.

Model Equality Constraint

The constraints, described in MuJoCo model as equality constrainst.

class JaxModelEqualityConstraint(dim, model, vector_gain, gain_fn, mask_idxs, active, hard_constraint, soft_constraint_cost)

A JAX-based equality constraint derived from the simulation model.

Equality constraints respresent all constrainst in the mujoco, which could be defined in <equality> tag in the xml file. More details could be found here: https://mujoco.readthedocs.io/en/stable/XMLreference.html#equality.

This class utilized the fact, that during kinematics computations, the equality value are computed as well, and the jacobian of the equalities are stored in data.efc_J, and the residuals are stored in data.efc_pos.

Parameters:

data – A MuJoCo simulation data structure containing model-specific constraint information.

Returns:

A jax.numpy.ndarray with the positions corresponding to the equality constraints.

compute_jacobian(data)

Compute the Jacobian matrix of the equality constraint.

Retrieves the relevant rows of the Jacobian matrix from the simulation data, filtering by equality constraint type and applying the mask indices.

Parameters:

data (Data) – The MuJoCo simulation data.

Return type:

Array

Returns:

A jax.numpy.ndarray representing the Jacobian matrix of the equality constraints.

class ModelEqualityConstraint(name='equality_constraint', gain=100, hard_constraint=False, soft_constraint_cost=None)

High-level component that aggregates all equality constraints from the simulation model.

The main purpose of the wrapper is to recalculate mask based on the dimensions of the equality constrain and compute proper dimensionality.

Parameters:

  • name (str) – The unique identifier for the constraint.

  • gain (ndarray | Array | float) – The gain for the constraint function, affecting its impact.

  • hard_constraint (bool) – If True, the constraint is handled as a hard constraint. Defaults to False (i.e., soft constraint).

  • soft_constraint_cost (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | bool | int | float | complex | str | bytes | _NestedSequence[bool | int | float | complex | str | bytes] | Array | ndarray | bool | number | None) – The cost used to relax a soft constraint. If not provided, a default high gain cost matrix (scaled identity matrix) based on the component dimension will be used.

JaxComponentType

alias of JaxModelEqualityConstraint

update_model(model)

Update the equality constraint using the provided MuJoCo model.

This method computes the total dimension of the equality constraints by iterating over the equality types in the model. Based on the type of each equality constraint, it sets the mask indices and component dimension accordingly.

Parameters:

model (MjModel) – The MuJoCo model instance.

Return type:

None

Returns:

None