curobo.wrap.reacher.types module

Module contains custom types and dataclasses used across reacher solvers.

class ReacherSolveType(value)

Bases: Enum

Enum for different types of problems solved with reacher solvers.

SINGLE = 0
GOALSET = 1
BATCH = 2
BATCH_GOALSET = 3
BATCH_ENV = 4
BATCH_ENV_GOALSET = 5
class ReacherSolveState(
solve_type: ReacherSolveType,
batch_size: int,
n_envs: int,
n_goalset: int = 1,
batch_env: bool = False,
batch_retract: bool = False,
batch_mode: bool = False,
num_seeds: int | None = None,
num_ik_seeds: int | None = None,
num_graph_seeds: int | None = None,
num_trajopt_seeds: int | None = None,
num_mpc_seeds: int | None = None,
)

Bases: object

Dataclass for storing the current problem type of a reacher solver.

solve_type: ReacherSolveType

Type of problem solved by the reacher solver.

batch_size: int

Number of problems in the batch.

n_envs: int

Number of environments in the batch.

n_goalset: int = 1

Number of goals per problem. Only valid for goalset problems.

batch_env: bool = False

Flag to indicate if the problems use different world environments in the batch.

batch_retract: bool = False

Flag to indicate if the problems use different retract configurations in the batch.

batch_mode: bool = False

Flag to indicate if there is more than 1 problem to be solved.

num_seeds: int | None = None

Number of seeds for each problem.

num_ik_seeds: int | None = None

Number of seeds for inverse kinematics problems.

num_graph_seeds: int | None = None

Number of seeds for graph search problems.

num_trajopt_seeds: int | None = None

Number of seeds for trajectory optimization problems.

num_mpc_seeds: int | None = None

Number of seeds for model predictive control problems.

clone() → ReacherSolveState

Method to create a deep copy of the current reacher solve state.

get_batch_size() → int

Method to get total number of optimization problems in the batch including seeds.

get_ik_batch_size() → int

Method to get total number of IK problems in the batch including seeds.

create_goal_buffer(
goal_pose: Pose,
goal_state: JointState | None = None,
retract_config: Tensor | None = None,
link_poses: Dict[str, Pose] | 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),
) → Goal

Method to create a goal buffer from goal pose and other problem targets.

Parameters:

  • goal_pose – Pose to reach with the end effector.

  • goal_state – Joint configuration to reach. If None, the goal is to reach the pose.

  • retract_config – Joint configuration to use for L2 regularization. If None, retract_config from robot configuration file is used. An alternative value is to use the start state as the retract configuration.

  • link_poses – Dictionary of link poses to reach. This is only required for multi-link pose reaching, where the goal is to reach multiple poses with different links.

  • tensor_args – Device and floating precision.

Returns:

Goal buffer with the goal pose, goal state, retract state, and link poses.

update_goal_buffer(
goal_pose: Pose,
goal_state: JointState | None = None,
retract_config: Tensor | None = None,
link_poses: List[Pose] | None = None,
current_solve_state: ReacherSolveState | None = None,
current_goal_buffer: Goal | 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),
) → Tuple[ReacherSolveState, Goal, bool]

Method to update the goal buffer with new goal pose and other problem targets.

Parameters:

  • goal_pose – Pose to reach with the end effector.

  • goal_state – Joint configuration to reach. If None, the goal is to reach the pose.

  • retract_config – Joint configuration to use for L2 regularization. If None, retract_config from robot configuration file is used. An alternative value is to use the start state as the retract configuration.

  • link_poses – Dictionary of link poses to reach. This is only required for multi-link pose reaching, where the goal is to reach multiple poses with different links. To use this, curobo.cuda_robot_model.cuda_robot_model.CudaRobotModelConfig.link_names should have the link names to reach.

  • current_solve_state – Current reacher solve state.

  • current_goal_buffer – Current goal buffer.

  • tensor_args – Device and floating precision.

Returns:

Tuple of updated reacher solve state, goal buffer, and a flag indicating if the goal buffer reference has changed which is useful to break existing CUDA Graphs.

update_goal(
goal: Goal,
current_solve_state: ReacherSolveState | None = None,
current_goal_buffer: Goal | 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),
) → Tuple[ReacherSolveState, Goal, bool]

Method to update the goal buffer with values from new Rollout goal.

Parameters:

  • goal – Rollout goal to update the goal buffer.

  • current_solve_state – Current reacher solve state.

  • current_goal_buffer – Current goal buffer.

  • tensor_args – Device and floating precision.

Returns:

Tuple of updated reacher solve state, goal buffer, and a flag indicating if the goal buffer reference has changed which is useful to break existing CUDA Graphs.

class MotionGenSolverState(
solve_type: ReacherSolveType,
ik_solve_state: ReacherSolveState,
trajopt_solve_state: ReacherSolveState,
)

Bases: object

Dataclass for storing the current state of a motion generation solver.

solve_type: ReacherSolveType
ik_solve_state: ReacherSolveState
trajopt_solve_state: ReacherSolveState
get_padded_goalset(
solve_state: ReacherSolveState,
current_solve_state: ReacherSolveState,
current_goal_buffer: Goal,
new_goal_pose: Pose,
) → Pose | None

Method to pad number of goals in goalset to match the cached goal buffer.

This allows for creating a goalset problem with large number of goals during the first call, and subsequent calls can have fewer goals. This function will pad the new goalset with the first goal to match the cached goal buffer’s shape.

Parameters:

  • solve_state – New problem’s solve state.

  • current_solve_state – Current solve state.

  • current_goal_buffer – Current goal buffer.

  • new_goal_pose – Padded goal pose to match the cached goal buffer’s shape.

Returns:

Padded goal pose to match the cached goal buffer’s shape. If the new goal can’t be padded, returns None.