Problem
At the heart of MJINX lies the Problem module - a structured framework that elegantly handles inverse kinematics challenges. This module serves as the central hub where various components come together to form a cohesive mathematical formulation.
When working with MJINX, a Problem instance orchestrates several key elements:
Tasks: Objective functions that define desired behaviors, such as reaching specific poses or following trajectories
Barriers: Smooth constraint functions that naturally keep the system within valid states
Velocity limits: Physical bounds on joint velocities to ensure feasible motion
The module’s modular architecture shines in its flexibility - users can begin with simple scenarios like positioning a single end-effector, then naturally build up to complex whole-body motions with multiple objectives and constraints.
Under the hood, the Problem class transforms these high-level specifications into optimized computations through its JaxProblemData representation. By leveraging JAX’s JIT compilation, it ensures that even sophisticated inverse kinematics problems run with maximum efficiency.
- class JaxProblemData(model, v_min, v_max, components)
A dataclass representing the compiled problem data for JAX optimization.
This container holds all the necessary elements required for solving an inverse kinematics problem with JAX-based optimization techniques.
- Parameters:
model (
Model) – The MuJoCo model.v_min (
Array) – Minimum velocity limits.v_max (
Array) – Maximum velocity limits.components (
dict[str,JaxComponent]) – Dictionary of JAX components (tasks and barriers).
- class Problem(model, v_min=-1000.0, v_max=1000.0)
A class representing an optimization problem for robotics applications.
This class manages the components (tasks and barriers), velocity limits, and model for the optimization problem. It serves as the central container for defining the constraints and objectives that characterize the inverse kinematics problem.
- Parameters:
- update_v_min(v_min)
Update the minimum velocity limits.
- update_v_max(v_max)
Update the maximum velocity limits.
- add_component(component)
Add a component to the problem.
- Parameters:
component (
Component) – The component to add.- Raises:
ValueError – If a component with the same name already exists.
- remove_component(name)
Remove a component from the problem.
- Parameters:
name (
str) – The name of the component to remove.
- compile()
Compile the problem into a JaxProblemData object.
- Return type:
- Returns:
A JaxProblemData object containing the compiled problem data.
- component(name)
Get a component by name.
- Parameters:
name (
str) – The name of the component.- Return type:
- Returns:
The requested component.
- Raises:
ValueError – If the component is not found.
- task(name)
Get a task component by name.
- Parameters:
name (
str) – The name of the task component.- Return type:
- Returns:
The requested task component.
- Raises:
ValueError – If the component is not found or is not a task.
- barrier(name)
Get a barrier component by name.
- Parameters:
name (
str) – The name of the barrier component.- Return type:
- Returns:
The requested barrier component.
- Raises:
ValueError – If the component is not found or is not a barrier.
- set_vmap_dimension()
Set the vmap dimension for the problem.
- Return type:
AbstractContextManager[JaxProblemData]- Returns:
A context manager for the JaxProblemData with vmap dimension set.