curobo.wrap.reacher.mpc module

This module contains MpcSolver that provides a high-level interface to for model predictive control (MPC) for reaching Cartesian poses and also joint configurations while avoiding obstacles. The solver uses Model Predictive Path Integral (MPPI) optimization as the solver. MPC only optimizes locally so the robot can get stuck near joint limits or behind obstacles. To generate global trajectories, use MotionGen.

A python example is available at Model Predictive Control (MPC).

Note

Gradient-based MPC is also implemented with L-BFGS but is highly experimental and not recommended for real robots.

class MpcSolverConfig(
solver: WrapMpc,
world_coll_checker: WorldCollision | None = None,
tensor_args: TensorDeviceType = TensorDeviceType(device=device(type='cuda', index=0), dtype=torch.float32, collision_geometry_dtype=torch.float32, collision_gradient_dtype=torch.float32, collision_distance_dtype=torch.float32),
use_cuda_graph_full_step: bool = False,
)

Bases: object

Configuration dataclass for MPC.

solver: WrapMpc

MPC Solver.

world_coll_checker: WorldCollision | None = None

World Collision Checker.

tensor_args: TensorDeviceType = TensorDeviceType(device=device(type='cuda', index=0), dtype=torch.float32, collision_geometry_dtype=torch.float32, collision_gradient_dtype=torch.float32, collision_distance_dtype=torch.float32)

Numeric precision and device to run computations.

use_cuda_graph_full_step: bool = False

Capture full step in MPC as a single CUDA graph. This is not supported currently.

static load_from_robot_config(
robot_cfg: str | dict | RobotConfig,
world_model: str | dict | WorldConfig,
base_cfg: dict | None = None,
tensor_args: TensorDeviceType = TensorDeviceType(device=device(type='cuda', index=0), dtype=torch.float32, collision_geometry_dtype=torch.float32, collision_gradient_dtype=torch.float32, collision_distance_dtype=torch.float32),
compute_metrics: bool = True,
use_cuda_graph: bool | None = None,
particle_opt_iters: int | None = None,
self_collision_check: bool = True,
collision_checker_type: CollisionCheckerType | None = CollisionCheckerType.MESH,
use_es: bool | None = None,
es_learning_rate: float | None = 0.01,
use_cuda_graph_metrics: bool = False,
store_rollouts: bool = True,
use_cuda_graph_full_step: bool = False,
sync_cuda_time: bool = True,
collision_cache: Dict[str, int] | None = None,
n_collision_envs: int | None = None,
collision_activation_distance: float | None = None,
world_coll_checker=None,
step_dt: float | None = None,
use_lbfgs: bool = False,
use_mppi: bool = True,
)

Create an MPC solver configuration from robot and world configuration.

Parameters:
  • robot_cfg – Robot configuration. Can be a path to a YAML file or a dictionary or an instance of RobotConfig.

  • world_model – World configuration. Can be a path to a YAML file or a dictionary or an instance of WorldConfig.

  • base_cfg – Base configuration for the solver. This file is used to check constraints and convergence. If None, the default configuration from base_cfg.yml is used.

  • tensor_args – Numeric precision and device to run computations.

  • compute_metrics – Compute metrics on MPC step.

  • use_cuda_graph – Use CUDA graph for the optimization step.

  • particle_opt_iters – Number of iterations for the particle optimization.

  • self_collision_check – Enable self-collision check during MPC optimization.

  • collision_checker_type – Type of collision checker to use. See Collision World Representation.

  • use_es – Use Evolution Strategies (ES) solver for MPC. Highly experimental.

  • es_learning_rate – Learning rate for ES solver.

  • use_cuda_graph_metrics – Use CUDA graph for computing metrics.

  • store_rollouts – Store rollouts information for debugging. This will also store the trajectory taken by the end-effector across the horizon.

  • use_cuda_graph_full_step – Capture full step in MPC as a single CUDA graph. This is experimental and might not work reliably.

  • sync_cuda_time – Synchronize CUDA device with host using torch.cuda.synchronize before calculating compute time.

  • collision_cache – Cache of obstacles to create to load obstacles between planning calls. An example: {"obb": 10, "mesh": 10}, to create a cache of 10 cuboids and 10 meshes.

  • n_collision_envs – Number of collision environments to create for batched planning across different environments. Only used for MpcSolver.setup_solve_batch_env and MpcSolver.setup_solve_batch_env_goalset.

  • collision_activation_distance – Distance in meters to activate collision cost. A good value to start with is 0.01 meters. Increase the distance if the robot needs to stay further away from obstacles.

  • world_coll_checker – Instance of world collision checker to use for MPC. Leaving this to None will create a new instance of world collision checker using the provided world_model.

  • step_dt – Time step to use between each step in the trajectory. If None, the default time step from the configuration~(particle_mpc.yml or gradient_mpc.yml) is used. This dt should match the control frequency at which you are sending commands to the robot. This dt should also be greater than than the compute time for a single step.

  • use_lbfgs – Use L-BFGS solver for MPC. Highly experimental.

  • use_mppi – Use MPPI solver for MPC.

Returns:

Configuration for the MPC solver.

Return type:

MpcSolverConfig

class MpcSolver(
config: MpcSolverConfig,
)

Bases: MpcSolverConfig

High-level interface for Model Predictive Control (MPC).

MPC can reach Cartesian poses and joint configurations while avoiding obstacles. The solver uses Model Predictive Path Integral (MPPI) optimization as the solver. MPC only optimizes locally so the robot can get stuck near joint limits or behind obstacles. To generate global trajectories, use MotionGen.

See Model Predictive Control (MPC) for an example. This MPC solver implementation can be used in the following steps:

  1. Create a Goal object with the target pose or joint configuration.

  2. Create a goal buffer for the problem type using setup_solve_single, setup_solve_goalset, setup_solve_batch, setup_solve_batch_goalset, setup_solve_batch_env, or setup_solve_batch_env_goalset. Pass the goal object from the previous step to this function. This function will update the internal solve state of MPC and also the goal for MPC. An augmented goal buffer is returned.

  3. Call step with the current joint state to get the next action.

  4. To change the goal, create a Pose object with new pose or JointState object with new joint configuration. Then copy the target into the augmented goal buffer using goal_buffer.goal_pose.copy_(new_pose) or goal_buffer.goal_state.copy_(new_state).

  5. Call update_goal with the augmented goal buffer to update the goal for MPC.

  6. Call step with the current joint state to get the next action.

To dynamically change the type of goal reached between pose and joint configuration targets, create the goal object in step 1 with both targets and then use enable_cspace_cost and enable_pose_cost to enable or disable reaching joint configuration cost and pose cost.

Initializes the MPC solver.

Parameters:

config – Configuration parameters for MPC.

tensor_args: TensorDeviceType = TensorDeviceType(device=device(type='cuda', index=0), dtype=torch.float32, collision_geometry_dtype=torch.float32, collision_gradient_dtype=torch.float32, collision_distance_dtype=torch.float32)

Numeric precision and device to run computations.

setup_solve_single(
goal: Goal,
num_seeds: int | None = None,
) Goal

Creates a goal buffer to solve for a robot to reach target pose or joint configuration.

Parameters:
  • goal – goal object containing target pose or joint configuration.

  • num_seeds – Number of seeds to use in the solver.

Returns:

Instance of augmented goal buffer.

Return type:

Goal

setup_solve_goalset(
goal: Goal,
num_seeds: int | None = None,
) Goal

Creates a goal buffer to solve for a robot to reach a pose in a set of poses.

Parameters:
  • goal – goal object containing target goalset or joint configuration.

  • num_seeds – Number of seeds to use in the solver.

Returns:

Instance of augmented goal buffer.

Return type:

Goal

setup_solve_batch(
goal: Goal,
num_seeds: int | None = None,
) Goal

Creates a goal buffer to solve for a batch of robots to reach targets.

Parameters:
  • goal – goal object containing target poses or joint configurations.

  • num_seeds – Number of seeds to use in the solver.

Returns:

Instance of augmented goal buffer.

Return type:

Goal

setup_solve_batch_goalset(
goal: Goal,
num_seeds: int | None = None,
) Goal

Creates a goal buffer to solve for a batch of robots to reach a set of poses.

Parameters:
  • goal – goal object containing target goalset or joint configurations.

  • num_seeds – Number of seeds to use in the solver.

Returns:

Instance of augmented goal buffer.

Return type:

Goal

setup_solve_batch_env(
goal: Goal,
num_seeds: int | None = None,
) Goal

Creates a goal buffer to solve for a batch of robots in different collision worlds.

Parameters:
  • goal – goal object containing target poses or joint configurations.

  • num_seeds – Number of seeds to use in the solver.

Returns:

Instance of augmented goal buffer.

Return type:

Goal

setup_solve_batch_env_goalset(
goal: Goal,
num_seeds: int | None = None,
) Goal

Creates a goal buffer to solve for a batch of robots in different collision worlds.

Parameters:
  • goal – goal object containing target goalset or joint configurations.

  • num_seeds – Number of seeds to use in the solver.

Returns:

Instance of augmented goal buffer.

Return type:

Goal

step(
current_state: JointState,
shift_steps: int = 1,
seed_traj: JointState | None = None,
max_attempts: int = 1,
)

Solve for the next action given the current state.

Parameters:
  • current_state – Current joint state of the robot.

  • shift_steps – Number of steps to shift the trajectory.

  • seed_traj – Initial trajectory to seed the optimization. If None, the solver uses the solution from the previous step.

  • max_attempts – Maximum number of attempts to solve the problem.

Returns:

Result of the optimization.

Return type:

WrapResult

update_goal(
goal: Goal,
)

Update the goal for MPC.

Parameters:

goal – goal object containing target pose or joint configuration. This goal instance should be created using one of the setup_solve functions.

reset()

Reset the solver.

reset_cuda_graph()

Reset captured CUDA graph. This does not work.

enable_cspace_cost(
enable=True,
)

Enable or disable reaching joint configuration cost in the solver.

Parameters:

enable – Enable or disable reaching joint configuration cost. When False, cspace cost is disabled.

enable_pose_cost(enable=True)

Enable or disable reaching pose cost in the solver.

Parameters:

enable – Enable or disable reaching pose cost. When False, pose cost is disabled.

get_active_js(
in_js: JointState,
)

Get controlled joints indexed in MPC order from the input joint state.

Parameters:

in_js – Input joint state.

Returns:

Joint state with controlled joints.

Return type:

JointState

update_world(
world: WorldConfig,
)

Update the collision world for the solver.

This allows for updating the world representation as long as the new world representation does not have a larger number of obstacles than the MpcSolver.collision_cache as created during initialization of MpcSolverConfig.

Parameters:

world – New collision world configuration. See Collision World Representation for more details.

get_visual_rollouts()

Get rollouts for debugging.

property joint_names

Get the ordered joint names of the robot.

property collision_cache: Dict[str, int]

Returns the collision cache created by the world collision checker.

property kinematics: CudaRobotModel

Get kinematics instance of the robot.

property world_collision: WorldCollision

Get the world collision checker.

property rollout_fn: ArmReacher

Get the rollout function.

_step_once(
current_state: JointState,
shift_steps: int = 1,
seed_traj: JointState | None = None,
) WrapResult

Solve for the next action given the current state.

Parameters:
  • current_state – Current joint state of the robot.

  • shift_steps – Number of steps to shift the trajectory.

  • seed_traj – Initial trajectory to seed the optimization. If None, the solver uses the solution from the previous step.

Returns:

Result of the optimization.

Return type:

WrapResult

_update_solve_state_and_goal_buffer(
solve_state: ReacherSolveState,
goal: Goal,
) Goal

Update solve state and goal for MPC.

Parameters:
  • solve_state – New solve state.

  • goal – New goal buffer.

Returns:

Updated goal buffer.

Return type:

Goal

_update_batch_size(
batch_size: int,
)

Update the batch size of the solver.

Parameters:

batch_size – Number of problems to solve in parallel.

_solve_from_solve_state(
solve_state: ReacherSolveState,
goal: Goal,
shift_steps: int = 1,
seed_traj: JointState | None = None,
) WrapResult

Solve for the next action given the current state.

Parameters:
  • solve_state – solve state object containing information about the current MPC problem.

  • goal – goal object containing target pose or joint configuration.

  • shift_steps – Number of steps to shift the trajectory before optimization.

  • seed_traj – Initial trajectory to seed the optimization. If None, the solver uses the solution from the previous step.

Returns:

Result of the optimization.

Return type:

WrapResult

_mpc_step(
current_state: JointState,
shift_steps: int = 1,
seed_traj: JointState | None = None,
)

One step function that is used to create a CUDA graph.

Parameters:
  • current_state – Current joint state of the robot.

  • shift_steps – Number of steps to shift the trajectory.

  • seed_traj – Initial trajectory to seed the optimization. If None, the solver uses the solution from the previous step.

Returns:

Result of the optimization.

Return type:

WrapResult

_initialize_cuda_graph_step(
current_state: JointState,
shift_steps: int = 1,
seed_traj: JointState | None = None,
)

Create a CUDA graph for the full step of MPC.

Parameters:
  • current_state – Current joint state of the robot.

  • shift_steps – Number of steps to shift the trajectory.

  • seed_traj – Initial trajectory to seed the optimization. If None, the solver uses the solution from the previous step.

static load_from_robot_config(
robot_cfg: str | dict | RobotConfig,
world_model: str | dict | WorldConfig,
base_cfg: dict | None = None,
tensor_args: TensorDeviceType = TensorDeviceType(device=device(type='cuda', index=0), dtype=torch.float32, collision_geometry_dtype=torch.float32, collision_gradient_dtype=torch.float32, collision_distance_dtype=torch.float32),
compute_metrics: bool = True,
use_cuda_graph: bool | None = None,
particle_opt_iters: int | None = None,
self_collision_check: bool = True,
collision_checker_type: CollisionCheckerType | None = CollisionCheckerType.MESH,
use_es: bool | None = None,
es_learning_rate: float | None = 0.01,
use_cuda_graph_metrics: bool = False,
store_rollouts: bool = True,
use_cuda_graph_full_step: bool = False,
sync_cuda_time: bool = True,
collision_cache: Dict[str, int] | None = None,
n_collision_envs: int | None = None,
collision_activation_distance: float | None = None,
world_coll_checker=None,
step_dt: float | None = None,
use_lbfgs: bool = False,
use_mppi: bool = True,
)

Create an MPC solver configuration from robot and world configuration.

Parameters:
  • robot_cfg – Robot configuration. Can be a path to a YAML file or a dictionary or an instance of RobotConfig.

  • world_model – World configuration. Can be a path to a YAML file or a dictionary or an instance of WorldConfig.

  • base_cfg – Base configuration for the solver. This file is used to check constraints and convergence. If None, the default configuration from base_cfg.yml is used.

  • tensor_args – Numeric precision and device to run computations.

  • compute_metrics – Compute metrics on MPC step.

  • use_cuda_graph – Use CUDA graph for the optimization step.

  • particle_opt_iters – Number of iterations for the particle optimization.

  • self_collision_check – Enable self-collision check during MPC optimization.

  • collision_checker_type – Type of collision checker to use. See Collision World Representation.

  • use_es – Use Evolution Strategies (ES) solver for MPC. Highly experimental.

  • es_learning_rate – Learning rate for ES solver.

  • use_cuda_graph_metrics – Use CUDA graph for computing metrics.

  • store_rollouts – Store rollouts information for debugging. This will also store the trajectory taken by the end-effector across the horizon.

  • use_cuda_graph_full_step – Capture full step in MPC as a single CUDA graph. This is experimental and might not work reliably.

  • sync_cuda_time – Synchronize CUDA device with host using torch.cuda.synchronize before calculating compute time.

  • collision_cache – Cache of obstacles to create to load obstacles between planning calls. An example: {"obb": 10, "mesh": 10}, to create a cache of 10 cuboids and 10 meshes.

  • n_collision_envs – Number of collision environments to create for batched planning across different environments. Only used for MpcSolver.setup_solve_batch_env and MpcSolver.setup_solve_batch_env_goalset.

  • collision_activation_distance – Distance in meters to activate collision cost. A good value to start with is 0.01 meters. Increase the distance if the robot needs to stay further away from obstacles.

  • world_coll_checker – Instance of world collision checker to use for MPC. Leaving this to None will create a new instance of world collision checker using the provided world_model.

  • step_dt – Time step to use between each step in the trajectory. If None, the default time step from the configuration~(particle_mpc.yml or gradient_mpc.yml) is used. This dt should match the control frequency at which you are sending commands to the robot. This dt should also be greater than than the compute time for a single step.

  • use_lbfgs – Use L-BFGS solver for MPC. Highly experimental.

  • use_mppi – Use MPPI solver for MPC.

Returns:

Configuration for the MPC solver.

Return type:

MpcSolverConfig

use_cuda_graph_full_step: bool = False

Capture full step in MPC as a single CUDA graph. This is not supported currently.

world_coll_checker: WorldCollision | None = None

World Collision Checker.

solver: WrapMpc

MPC Solver.