hofee/gen_data_curobo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add re-timing, minimum dt robustness

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-04-25 12:24:17 -07:00
parent d6e600c88c
commit 7362ccd4c2
54 changed files with 4773 additions and 2189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/curobo/__init__.py
View File

@@ -10,17 +10,18 @@
#
"""
cuRobo provides accelerated modules for robotics which can be used to build high-performance
cuRobo provides accelerated modules for robotics which can be used to build high-performance
robotics applications. The library has several modules for numerical optimization, robot kinematics,
geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for
performing tasks like collision-free inverse kinematics, model predictive control, and motion
geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for
performing tasks like collision-free inverse kinematics, model predictive control, and motion
planning.
High-level APIs:
- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`.
- Inverse Kinematics: :mod:`curobo.wrap.reacher.ik_solver`.
- Model Predictive Control: :mod:`curobo.wrap.reacher.mpc`.
- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`.
- Trajectory Optimization: :mod:`curobo.wrap.reacher.trajopt`.
cuRobo package is split into several modules:
@@ -33,10 +34,11 @@ cuRobo package is split into several modules:
- :mod:`curobo.rollout` contains methods that map actions to costs. This class wraps instances of
:mod:`curobo.cuda_robot_model` and :mod:`curobo.geom` to compute costs given trajectory of actions.
- :mod:`curobo.util` contains utility methods.
- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of
- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of
collision-free reacher and batched robot world collision checking.
- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
:meth:`types.state.JointState`, :meth:`types.camera.CameraObservation`, :meth:`types.math.Pose`.
- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
:py:meth:`~types.state.JointState`, :py:meth:`~types.camera.CameraObservation`,
:py:meth:`~types.math.Pose`.
"""

12
src/curobo/content/configs/robot/franka.yml
View File

@@ -28,7 +28,7 @@ robot_cfg:
"panda_finger_joint1": "Y",
"panda_finger_joint2": "Y",
}
usd_flip_joint_limits: ["panda_finger_joint2"]
urdf_path: "robot/franka_description/franka_panda.urdf"
asset_root_path: "robot/franka_description"
@@ -67,13 +67,13 @@ robot_cfg:
"panda_hand": ["panda_leftfinger", "panda_rightfinger","attached_object"],
"panda_leftfinger": ["panda_rightfinger", "attached_object"],
"panda_rightfinger": ["attached_object"],
}
self_collision_buffer:
{
"panda_link0": 0.1,
"panda_link1": 0.05,
"panda_link0": 0.1,
"panda_link1": 0.05,
"panda_link2": 0.0,
"panda_link3": 0.0,
"panda_link4": 0.0,
@@ -101,7 +101,7 @@ robot_cfg:
"panda_rightfinger",
]
lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}
extra_links: {"attached_object":{"parent_link_name": "panda_hand" ,
extra_links: {"attached_object":{"parent_link_name": "panda_hand" ,
"link_name": "attached_object", "fixed_transform": [0,0,0,1,0,0,0], "joint_type":"FIXED",
"joint_name": "attach_joint" }}
cspace:

8
src/curobo/content/configs/robot/spheres/franka_mesh.yml
View File

@@ -134,11 +134,11 @@ collision_spheres:
"radius": 0.022
panda_leftfinger:
- "center": [0.0, 0.01, 0.043]
"radius": 0.011 #0.025 # 0.011
"radius": 0.011
- "center": [0.0, 0.02, 0.015]
"radius": 0.011 #0.025 # 0.011
"radius": 0.011
panda_rightfinger:
- "center": [0.0, -0.01, 0.043]
"radius": 0.011 #0.025 #0.011
"radius": 0.011
- "center": [0.0, -0.02, 0.015]
"radius": 0.011 #0.025 #0.011
"radius": 0.011

2
src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml
View File

@@ -23,7 +23,7 @@ robot_cfg:
base_link: "base_link"
ee_link: "grasp_frame"
link_names: null
lock_joints: {'finger_joint': 0.7}
lock_joints: {'finger_joint': 0.3}
extra_links: null #{"attached_object":{"parent_link_name": "grasp_frame" ,
#"link_name": "attached_object", "fixed_transform": [0,0,0,1,0,0,0], "joint_type":"FIXED",

4
src/curobo/content/configs/task/base_cfg.yml
View File

@@ -11,8 +11,8 @@
world_collision_checker_cfg:
cache: null #{"cube": 41, "capsule": 0, "sphere": 0}
checker_type: "PRIMITIVE" # ["PRIMITIVE", "BLOX", "MESH"]
cache: null
checker_type: "PRIMITIVE"
max_distance: 0.1

34
src/curobo/content/configs/task/finetune_trajopt.yml
View File

@@ -19,12 +19,12 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
state_finite_difference_mode: "CENTRAL"
teleport_mode: False
return_full_act_buffer: True
@@ -35,7 +35,7 @@ cost:
weight: [2000,500000.0,30,60]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 1.0
run_weight: 1.0
use_metric: True
link_pose_cfg:
@@ -47,19 +47,19 @@ cost:
run_weight: 1.0
use_metric: True
cspace_cfg:
weight: 10000.0
terminal: True
run_weight: 0.00 #1
bound_cfg:
weight: [10000.0, 500000.0, 500.0, 500.0]
smooth_weight: [0.0,10000.0, 5.0, 0.0]
weight: [10000.0, 500000.0, 500.0, 500.0]
smooth_weight: [0.0,10000.0, 5.0, 0.0]
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk
activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -75,21 +75,21 @@ cost:
self_collision_cfg:
weight: 5000.0
classify: False
lbfgs:
n_iters: 300
inner_iters: 25
cold_start_n_iters: null
n_iters: 400
inner_iters: 25
cold_start_n_iters: null
min_iters: 25
line_search_scale: [0.1, 0.3, 0.7,1.0] # #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 1.0
cost_relative_threshold: 0.999 #0.999
cost_convergence: 0.001
cost_delta_threshold: 0.0001
cost_relative_threshold: 0.99
epsilon: 0.01
history: 27 #15
history: 27
use_cuda_graph: True
n_problems: 1
store_debug: False

34
src/curobo/content/configs/task/finetune_trajopt_slow.yml
View File

@@ -19,9 +19,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
@@ -31,22 +31,21 @@ model:
cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position
weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
weight: [2000,500000.0,30,60]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 0.0 #0.05
run_weight: 1.0
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position
weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
weight: [2000,500000.0,30,60] #[150.0, 2000.0, 30, 40]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 0.001
run_weight: 1.0
use_metric: True
cspace_cfg:
weight: 10000.0
@@ -54,12 +53,12 @@ cost:
run_weight: 0.00 #1
bound_cfg:
weight: [10000.0, 100000.0, 500.0, 500.0]
smooth_weight: [0.0,10000.0, 50.0, 0.0]
weight: [10000.0, 500000.0, 500.0, 500.0]
smooth_weight: [0.0,10000.0, 5.0, 0.0]
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
activation_distance: [0.1,0.1,0.1,10.0] # for position, velocity, acceleration and jerk
activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -75,21 +74,20 @@ cost:
self_collision_cfg:
weight: 5000.0
classify: False
lbfgs:
n_iters: 400 # 400
n_iters: 400
inner_iters: 25
cold_start_n_iters: 200
cold_start_n_iters: null
min_iters: 25
line_search_scale: [0.01, 0.3, 0.7,1.0] # #
line_search_scale: [0.1, 0.3, 0.7,1.0] # #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 1.0
cost_relative_threshold: 0.999 #0.999
epsilon: 0.01
history: 15 #15
history: 27 #15
use_cuda_graph: True
n_problems: 1
store_debug: False

18
src/curobo/content/configs/task/gradient_ik.yml
View File

@@ -29,32 +29,30 @@ cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [2000,10000,20,40]
vec_convergence: [0.0, 0.00]
vec_convergence: [0.0, 0.00]
terminal: False
use_metric: True
run_weight: 1.0
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [2000,10000,20,40]
vec_convergence: [0.00, 0.000]
vec_convergence: [0.00, 0.000]
terminal: False
use_metric: True
run_weight: 1.0
cspace_cfg:
weight: 0.000
bound_cfg:
weight: 500.0
activation_distance: [0.1]
null_space_weight: [1.0]
weight: 5000.0
activation_distance: [0.001]
null_space_weight: [0.1]
primitive_collision_cfg:
weight: 5000.0
use_sweep: False
classify: False
activation_distance: 0.01
self_collision_cfg:
weight: 500.0
weight: 5000.0
classify: False
@@ -81,10 +79,10 @@ lbfgs:
use_cuda_line_search_kernel: True
use_cuda_update_best_kernel: True
sync_cuda_time: True
step_scale: 1.0
step_scale: 1.0
use_coo_sparse: True
last_best: 10
debug_info:
visual_traj : null #'ee_pos_seq'

22
src/curobo/content/configs/task/gradient_trajopt.yml
View File

@@ -19,12 +19,12 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
state_finite_difference_mode: "CENTRAL"
teleport_mode: False
return_full_act_buffer: True
@@ -48,17 +48,17 @@ cost:
use_metric: True
cspace_cfg:
weight: 1000.0
weight: 10000.0
terminal: True
run_weight: 0.0
bound_cfg:
weight: [50000.0, 50000.0, 50000.0,50000.0]
smooth_weight: [0.0,10000.0,5.0, 0.0]
weight: [50000.0, 50000.0, 500.0,500.0]
smooth_weight: [0.0,10000.0,5.0, 0.0]
run_weight_velocity: 0.00
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk
activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -75,13 +75,13 @@ cost:
self_collision_cfg:
weight: 5000.0
classify: False
lbfgs:
n_iters: 100 #175
inner_iters: 25
cold_start_n_iters: null
inner_iters: 25
cold_start_n_iters: null
min_iters: 25
line_search_scale: [0.1,0.3,0.7,1.0]
fixed_iters: True

24
src/curobo/content/configs/task/particle_trajopt.yml
View File

@@ -18,9 +18,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
teleport_mode: False
@@ -30,22 +30,22 @@ model:
cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
weight: [250.0, 5000.0, 20, 20]
vec_convergence: [0.0,0.0,1000.0,1000.0]
vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True
run_weight: 0.0
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
weight: [0.0, 5000.0, 40, 40]
vec_convergence: [0.0,0.0,1000.0,1000.0]
vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True
run_weight: 0.00
run_weight: 0.00
use_metric: True
cspace_cfg:
weight: 10000.0
terminal: True
@@ -74,8 +74,8 @@ cost:
self_collision_cfg:
weight: 500.0
classify: False
mppi:
@@ -89,7 +89,7 @@ mppi:
alpha : 1
num_particles : 25 # 100
update_cov : True
cov_type : "DIAG_A" #
cov_type : "DIAG_A" #
kappa : 0.001
null_act_frac : 0.0
sample_mode : 'MEAN'

71
src/curobo/cuda_robot_model/cuda_robot_model.py
View File

@@ -24,11 +24,7 @@ from curobo.cuda_robot_model.cuda_robot_generator import (
CudaRobotGeneratorConfig,
)
from curobo.cuda_robot_model.kinematics_parser import KinematicsParser
from curobo.cuda_robot_model.types import (
CudaRobotModelState,
KinematicsTensorConfig,
SelfCollisionKinematicsConfig,
)
from curobo.cuda_robot_model.types import KinematicsTensorConfig, SelfCollisionKinematicsConfig
from curobo.curobolib.kinematics import get_cuda_kinematics
from curobo.geom.types import Mesh, Sphere
from curobo.types.base import TensorDeviceType
@@ -148,14 +144,67 @@ class CudaRobotModelConfig:
return self.kinematics_config.n_dof
@dataclass
class CudaRobotModelState:
"""Dataclass that stores kinematics information."""
#: End-effector position stored as x,y,z in meters [b, 3]. End-effector is defined by
#: :attr:`CudaRobotModel.ee_link`.
ee_position: torch.Tensor
#: End-effector orientaiton stored as quaternion qw, qx, qy, qz [b,4]. End-effector is defined
#: by :attr:`CudaRobotModel.ee_link`.
ee_quaternion: torch.Tensor
#: Linear Jacobian. Currently not supported.
lin_jacobian: Optional[torch.Tensor] = None
#: Angular Jacobian. Currently not supported.
ang_jacobian: Optional[torch.Tensor] = None
#: Position of links specified by link_names (:attr:`CudaRobotModel.link_names`).
links_position: Optional[torch.Tensor] = None
#: Quaternions of links specified by link names (:attr:`CudaRobotModel.link_names`).
links_quaternion: Optional[torch.Tensor] = None
#: Position of spheres specified by collision spheres (:attr:`CudaRobotModel.robot_spheres`)
#: in x, y, z, r format [b,n,4].
link_spheres_tensor: Optional[torch.Tensor] = None
#: Names of links that each index in :attr:`links_position` and :attr:`links_quaternion`
#: corresponds to.
link_names: Optional[str] = None
@property
def ee_pose(self) -> Pose:
"""Get end-effector pose as a Pose object."""
return Pose(self.ee_position, self.ee_quaternion)
def get_link_spheres(self) -> torch.Tensor:
"""Get spheres representing robot geometry as a tensor with [batch,4], [x,y,z,radius]."""
return self.link_spheres_tensor
@property
def link_pose(self) -> Dict[str, Pose]:
"""Get link poses as a dictionary of link name to Pose object."""
link_poses = None
if self.link_names is not None:
link_poses = {}
link_pos = self.links_position.contiguous()
link_quat = self.links_quaternion.contiguous()
for i, v in enumerate(self.link_names):
link_poses[v] = Pose(link_pos[..., i, :], link_quat[..., i, :])
return link_poses
class CudaRobotModel(CudaRobotModelConfig):
"""
CUDA Accelerated Robot Model
NOTE: Currently dof is created only for links that we need to compute kinematics.
E.g., for robots with many serial chains, add all links of the robot to get the correct dof.
This is not an issue if you are loading collision spheres as that will cover the full geometry
of the robot.
Currently dof is created only for links that we need to compute kinematics. E.g., for robots
with many serial chains, add all links of the robot to get the correct dof. This is not an
issue if you are loading collision spheres as that will cover the full geometry of the robot.
"""
def __init__(self, config: CudaRobotModelConfig):
@@ -421,6 +470,10 @@ class CudaRobotModel(CudaRobotModelConfig):
def base_link(self):
return self.kinematics_config.base_link
@property
def robot_spheres(self):
return self.kinematics_config.link_spheres
def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig):
self.kinematics_config.copy_(new_kin_config)

49
src/curobo/cuda_robot_model/types.py
View File

@@ -496,52 +496,3 @@ class SelfCollisionKinematicsConfig:
collision_matrix: Optional[torch.Tensor] = None
experimental_kernel: bool = True
checks_per_thread: int = 32
@dataclass
class CudaRobotModelState:
"""Dataclass that stores kinematics information."""
#: End-effector position stored as x,y,z in meters [b, 3]. End-effector is defined by
#: :py:attr:`curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGeneratorConfig.ee_link`.
ee_position: torch.Tensor
#: End-effector orientaiton stored as quaternion qw, qx, qy, qz [b,4]. End-effector is defined
# by :py:attr:`CudaRobotModelConfig.ee_link`.
ee_quaternion: torch.Tensor
#: Linear Jacobian. Currently not supported.
lin_jacobian: Optional[torch.Tensor] = None
#: Angular Jacobian. Currently not supported.
ang_jacobian: Optional[torch.Tensor] = None
#: Position of links specified by link_names (:py:attr:`CudaRobotModelConfig.link_names`).
links_position: Optional[torch.Tensor] = None
#: Quaternions of links specified by link names (:py:attr:`CudaRobotModelConfig.link_names`).
links_quaternion: Optional[torch.Tensor] = None
#: Position of spheres specified by collision spheres (:py:attr:`CudaRobotModelConfig.collision_spheres`)
#: in x, y, z, r format [b,n,4].
link_spheres_tensor: Optional[torch.Tensor] = None
link_names: Optional[str] = None
@property
def ee_pose(self):
return Pose(self.ee_position, self.ee_quaternion)
def get_link_spheres(self):
return self.link_spheres_tensor
@property
def link_pose(self):
link_poses = None
if self.link_names is not None:
link_poses = {}
link_pos = self.links_position.contiguous()
link_quat = self.links_quaternion.contiguous()
for i, v in enumerate(self.link_names):
link_poses[v] = Pose(link_pos[..., i, :], link_quat[..., i, :])
return link_poses

17
src/curobo/curobolib/cpp/sphere_obb_kernel.cu
View File

@@ -194,13 +194,13 @@ namespace Curobo
const scalar_t x = transform_mat[4];
const scalar_t y = transform_mat[5];
const scalar_t z = transform_mat[6];
if ((x != 0) || (y != 0) || (z != 0))
if(x != 0 || y!= 0 || z!=0)
{
C.x = p_arr.x + w * w * sphere_pos.x + 2 * y * w * sphere_pos.z -
2 * z * w * sphere_pos.y + x * x * sphere_pos.x +
2 * y * x * sphere_pos.y + 2 * z * x * sphere_pos.z -
z * z * sphere_pos.x - y * y * sphere_pos.x;
C.y = p_arr.y + 2 * x * y * sphere_pos.x + y * y * sphere_pos.y +
2 * z * y * sphere_pos.z + 2 * w * z * sphere_pos.x -
z * z * sphere_pos.y + w * w * sphere_pos.y - 2 * x * w * sphere_pos.z -
@@ -714,7 +714,7 @@ float4 &sum_pt)
// compute interpolation distance between voxel origin and sphere location:
get_array_value(grid_features, voxel_idx, interpolated_distance);// + 0.5 * loc_grid_params.w;
get_array_value(grid_features, voxel_idx, interpolated_distance);
if(!INTERPOLATION)
{
interpolated_distance += 0.5 * loc_grid_params.w;//max(0.0, (0.3 * loc_grid_params.w) - loc_sphere.w);
@@ -821,7 +821,7 @@ float4 &sum_pt)
d_grad = (d_plus - d_minus) * (1/(2*voxel_size));
}
else
if (!CENTRAL_DIFFERENCE)
{
// x difference:
int x_minus,y_minus, z_minus;
@@ -1252,7 +1252,9 @@ float4 &sum_pt)
const int32_t env_idx,
const int bn_sph_idx,
const int sph_idx,
dist_scalar_t *out_distance, grad_scalar_t *closest_pt, uint8_t *sparsity_idx,
dist_scalar_t *out_distance,
grad_scalar_t *closest_pt,
uint8_t *sparsity_idx,
const float *weight,
const float *activation_distance,
const float *max_distance,
@@ -1863,7 +1865,9 @@ float4 &sum_pt)
const int32_t env_idx,
const int bn_sph_idx,
const int sph_idx,
dist_scalar_t *out_distance, grad_scalar_t *closest_pt, uint8_t *sparsity_idx,
dist_scalar_t *out_distance,
grad_scalar_t *closest_pt,
uint8_t *sparsity_idx,
const float *weight,
const float *activation_distance,
const float *max_distance,
@@ -1879,7 +1883,6 @@ float4 &sum_pt)
const float eta = max_distance_local;
float max_dist = -1 * max_distance_local;
max_distance_local = -1 * max_distance_local;
;
float3 max_grad = make_float3(0.0, 0.0, 0.0);

6
src/curobo/geom/cv.py
View File

@@ -47,7 +47,9 @@ def project_depth_to_pointcloud(depth_image: torch.Tensor, intrinsics_matrix: to
@get_torch_jit_decorator()
def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor):
def get_projection_rays(
height: int, width: int, intrinsics_matrix: torch.Tensor, depth_to_meter: float = 0.001
):
"""Projects numpy depth image to point cloud.
Args:
@@ -82,7 +84,7 @@ def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor
rays = torch.stack([output_x, output_y, input_z], -1).reshape(
intrinsics_matrix.shape[0], width * height, 3
)
rays = rays * (1.0 / 1000.0)
rays = rays * depth_to_meter
return rays

4
src/curobo/geom/sdf/world.py
View File

@@ -264,7 +264,7 @@ class WorldCollisionConfig:
n_envs: int = 1
checker_type: CollisionCheckerType = CollisionCheckerType.PRIMITIVE
max_distance: Union[torch.Tensor, float] = 0.1
max_esdf_distance: Union[torch.Tensor, float] = 1000.0
max_esdf_distance: Union[torch.Tensor, float] = 100.0
def __post_init__(self):
if self.world_model is not None and isinstance(self.world_model, list):
@@ -398,6 +398,7 @@ class WorldCollision(WorldCollisionConfig):
self._cache_voxelization is None
or self._cache_voxelization.voxel_size != new_grid.voxel_size
or self._cache_voxelization.dims != new_grid.dims
or self._cache_voxelization.xyzr_tensor is None
):
self._cache_voxelization = new_grid
self._cache_voxelization.xyzr_tensor = self._cache_voxelization.create_xyzr_tensor(
@@ -458,7 +459,6 @@ class WorldCollision(WorldCollisionConfig):
self.update_cache_voxelization(new_grid)
xyzr = self._cache_voxelization.xyzr_tensor
voxel_shape = xyzr.shape
xyzr = xyzr.view(-1, 1, 1, 4)
weight = self.tensor_args.to_device([1.0])

12
src/curobo/geom/sdf/world_blox.py
View File

@@ -127,6 +127,7 @@ class WorldBloxCollision(WorldVoxelCollision):
collision_query_buffer: CollisionQueryBuffer,
weight,
activation_distance,
return_loss: bool = False,
compute_esdf: bool = False,
):
d = self._blox_mapper.query_sphere_sdf_cost(
@@ -136,8 +137,11 @@ class WorldBloxCollision(WorldVoxelCollision):
collision_query_buffer.blox_collision_buffer.sparsity_index_buffer,
weight,
activation_distance,
self.max_esdf_distance,
self._blox_tensor_list[0],
self._blox_tensor_list[1],
return_loss,
compute_esdf,
)
return d
@@ -191,12 +195,12 @@ class WorldBloxCollision(WorldVoxelCollision):
sum_collisions=sum_collisions,
compute_esdf=compute_esdf,
)
d = self._get_blox_sdf(
query_sphere,
collision_query_buffer,
weight=weight,
activation_distance=activation_distance,
return_loss=return_loss,
compute_esdf=compute_esdf,
)
@@ -246,6 +250,7 @@ class WorldBloxCollision(WorldVoxelCollision):
collision_query_buffer,
weight=weight,
activation_distance=activation_distance,
return_loss=return_loss,
)
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
"mesh" not in self.collision_types or not self.collision_types["mesh"]
@@ -404,7 +409,7 @@ class WorldBloxCollision(WorldVoxelCollision):
cuboid: Cuboid,
layer_name: Optional[str] = None,
):
raise NotImplementedError
log_error("Not implemented")
def get_bounding_spheres(
self,
@@ -418,7 +423,8 @@ class WorldBloxCollision(WorldVoxelCollision):
clear_region: bool = False,
) -> List[Sphere]:
mesh = self.get_mesh_in_bounding_box(bounding_box, obstacle_name, clear_region=clear_region)
if clear_region:
self.clear_bounding_box(bounding_box, obstacle_name)
spheres = mesh.get_bounding_spheres(
n_spheres=n_spheres,
surface_sphere_radius=surface_sphere_radius,

6
src/curobo/geom/sdf/world_voxel.py
View File

@@ -80,6 +80,9 @@ class WorldVoxelCollision(WorldMeshCollision):
if feature_dtype in [torch.float32, torch.float16, torch.bfloat16]:
voxel_features[:] = -1.0 * self.max_esdf_distance
else:
if self.max_esdf_distance > 100.0:
log_warn("Using fp8 for WorldVoxelCollision will reduce max_esdf_distance to 100")
self.max_esdf_distance = 100.0
voxel_features = (voxel_features.to(dtype=torch.float16) - self.max_esdf_distance).to(
dtype=feature_dtype
)
@@ -503,9 +506,10 @@ class WorldVoxelCollision(WorldMeshCollision):
False,
use_batch_env,
False,
True,
False,
False,
)
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
"mesh" not in self.collision_types or not self.collision_types["mesh"]
):

1
src/curobo/geom/sphere_fit.py
View File

@@ -92,6 +92,7 @@ def voxel_fit_surface_mesh(
def get_voxelgrid_from_mesh(mesh: trimesh.Trimesh, n_spheres: int, voxelize_method: str = "ray"):
"""Get voxel grid from mesh using :py:func:`trimesh.voxel.creation.voxelize`."""
pitch = get_voxel_pitch(mesh, n_spheres)
radius = pitch / 2.0
try:

4
src/curobo/geom/types.py
View File

@@ -58,9 +58,9 @@ class Obstacle:
texture: Optional[str] = None
#: material properties to apply in visualization.
material: Material = Material()
material: Material = field(default_factory=Material)
tensor_args: TensorDeviceType = TensorDeviceType()
tensor_args: TensorDeviceType = field(default_factory=TensorDeviceType)
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True) -> trimesh.Trimesh:
"""Create a trimesh instance from the obstacle representation.

6
src/curobo/opt/newton/newton_base.py
View File

@@ -125,7 +125,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
def _shift(self, shift_steps=1):
# TODO: shift best q?:
self.best_cost[:] = 500000.0
self.best_cost[:] = 5000000.0
self.best_iteration[:] = 0
self.current_iteration[:] = 0
return True
@@ -159,8 +159,9 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
with profiler.record_function("newton/reset"):
self.i = -1
self._opt_finished = False
self.best_cost[:] = 500000.0
self.best_cost[:] = 5000000.0
self.best_iteration[:] = 0
self.current_iteration[:] = 0
super().reset()
@@ -448,6 +449,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self.cost_delta_threshold,
self.cost_relative_threshold,
)
# print(self.best_cost[0], self.best_q[0])
else:
cost = cost.detach()
q = q.detach()

2
src/curobo/opt/opt_base.py
View File

@@ -169,7 +169,7 @@ class Optimizer(OptimizerConfig):
st_time = time.time()
out = self._optimize(opt_tensor, shift_steps, n_iters)
if self.sync_cuda_time:
torch.cuda.synchronize()
torch.cuda.synchronize(device=self.tensor_args.device)
self.opt_dt = time.time() - st_time
return out

3
src/curobo/rollout/arm_base.py
View File

@@ -384,6 +384,7 @@ class ArmBase(RolloutBase, ArmBaseConfig):
# setup constraint terms:
constraint = self.bound_constraint.forward(state.state_seq)
constraint_list = [constraint]
if (
self.constraint_cfg.primitive_collision_cfg is not None
@@ -407,7 +408,9 @@ class ArmBase(RolloutBase, ArmBaseConfig):
self_constraint = self.robot_self_collision_constraint.forward(state.robot_spheres)
constraint_list.append(self_constraint)
constraint = cat_sum(constraint_list)
feasible = constraint == 0.0
if out_metrics is None:
out_metrics = RolloutMetrics()
out_metrics.feasible = feasible

3
src/curobo/rollout/arm_reacher.py
View File

@@ -263,7 +263,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
goal_cost = self.goal_cost.forward(
ee_pos_batch, ee_quat_batch, self._goal_buffer
)
# print(self._compute_g_dist, goal_cost.view(-1))
cost_list.append(goal_cost)
with profiler.record_function("cost/link_poses"):
if self._goal_buffer.links_goal_pose is not None and self.cost_cfg.pose_cfg is not None:
@@ -286,6 +286,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
and self.cost_cfg.cspace_cfg is not None
and self.dist_cost.enabled
):
joint_cost = self.dist_cost.forward_target_idx(
self._goal_buffer.goal_state.position,
state_batch.position,

119
src/curobo/rollout/cost/bound_cost.py
View File

@@ -21,6 +21,7 @@ import warp as wp
from curobo.cuda_robot_model.types import JointLimits
from curobo.types.robot import JointState
from curobo.types.tensor import T_DOF
from curobo.util.logger import log_error
from curobo.util.torch_utils import get_torch_jit_decorator
from curobo.util.warp import init_warp
@@ -53,6 +54,16 @@ class BoundCostConfig(CostConfig):
self.joint_limits = bounds.clone()
if teleport_mode:
self.cost_type = BoundCostType.POSITION
if self.cost_type != BoundCostType.POSITION:
if torch.max(self.joint_limits.velocity[1] - self.joint_limits.velocity[0]) == 0.0:
log_error("Joint velocity limits is zero")
if (
torch.max(self.joint_limits.acceleration[1] - self.joint_limits.acceleration[0])
== 0.0
):
log_error("Joint acceleration limits is zero")
if torch.max(self.joint_limits.jerk[1] - self.joint_limits.jerk[0]) == 0.0:
log_error("Joint jerk limits is zero")
def __post_init__(self):
if isinstance(self.activation_distance, List):
@@ -675,8 +686,14 @@ def forward_bound_pos_warp(
target_p = retract_config[target_id]
p_l = p_b[d_id]
p_u = p_b[dof + d_id]
p_range = p_u - p_l
eta_p = eta_p * (p_range)
p_l += eta_p
p_u -= eta_p
if p_range < 1.0:
w = (1.0 / p_range) * w
# compute cost:
b_addrs = b_id * horizon * dof + h_id * dof + d_id
@@ -690,33 +707,15 @@ def forward_bound_pos_warp(
g_p = 2.0 * n_w * error
# bound cost:
# if c_p < p_l:
# delta = p_l - c_p
# if (delta) > eta_p or eta_p == 0.0:
# c_total += w * (delta - 0.5 * eta_p)
# g_p += -w
# else:
# c_total += w * (0.5 / eta_p) * delta * delta
# g_p += -w * (1.0 / eta_p) * delta
# elif c_p > p_u:
# delta = c_p - p_u
# if (delta) > eta_p or eta_p == 0.0:
# c_total += w * (delta - 0.5 * eta_p)
# g_p += w
# else:
# c_total += w * (0.5 / eta_p) * delta * delta
# g_p += w * (1.0 / eta_p) * delta
# bound cost:
delta = 0.0
if c_p < p_l:
delta = c_p - p_l
c_total += w * delta * delta
g_p += 2.0 * w * delta
elif c_p > p_u:
delta = c_p - p_u
c_total += w * delta * delta
g_p += 2.0 * w * delta
c_total += w * delta * delta
g_p += 2.0 * w * delta
out_cost[b_addrs] = c_total
# compute gradient
@@ -807,16 +806,33 @@ def forward_bound_warp(
c_j = jerk[b_addrs]
p_l = p_b[d_id] + eta_p
p_u = p_b[dof + d_id] - eta_p
p_l = p_b[d_id]
p_u = p_b[dof + d_id]
p_range = p_u - p_l
eta_p = eta_p * (p_range)
p_l += eta_p
p_u -= eta_p
v_l = v_b[d_id] + eta_v
v_u = v_b[dof + d_id] - eta_v
a_l = a_b[d_id] + eta_a
a_u = a_b[dof + d_id] - eta_a
v_l = v_b[d_id]
v_u = v_b[dof + d_id]
v_range = v_u - v_l
eta_v = eta_v * (v_range)
v_l += eta_v
v_u -= eta_v
j_l = j_b[d_id] + eta_j
j_u = j_b[dof + d_id] - eta_j
a_l = a_b[d_id]
a_u = a_b[dof + d_id]
a_range = a_u - a_l
eta_a = eta_a * (a_range)
a_l += eta_a
a_u -= eta_a
j_l = j_b[d_id]
j_u = j_b[dof + d_id]
j_range = j_u - j_l
eta_j = eta_j * (j_range)
j_l += eta_j
j_u -= eta_j
delta = float(0.0)
if n_w > 0.0:
@@ -826,6 +842,7 @@ def forward_bound_warp(
# bound cost:
delta = 0.0
if c_p < p_l:
delta = c_p - p_l
elif c_p > p_u:
@@ -995,19 +1012,47 @@ def forward_bound_smooth_warp(
r_wj = run_weight_jerk[h_id]
c_j = jerk[b_addrs]
p_l = p_b[d_id] + eta_p
p_u = p_b[dof + d_id] - eta_p
p_l = p_b[d_id]
p_u = p_b[dof + d_id]
p_range = p_u - p_l
eta_p = eta_p * (p_range)
p_l += eta_p
p_u -= eta_p
v_l = v_b[d_id] + eta_v
v_u = v_b[dof + d_id] - eta_v
a_l = a_b[d_id] + eta_a
a_u = a_b[dof + d_id] - eta_a
v_l = v_b[d_id]
v_u = v_b[dof + d_id]
v_range = v_u - v_l
eta_v = eta_v * (v_range)
v_l += eta_v
v_u -= eta_v
j_l = j_b[d_id] + eta_j
j_u = j_b[dof + d_id] - eta_j
a_l = a_b[d_id]
a_u = a_b[dof + d_id]
a_range = a_u - a_l
eta_a = eta_a * (a_range)
a_l += eta_a
a_u -= eta_a
j_l = j_b[d_id]
j_u = j_b[dof + d_id]
j_range = j_u - j_l
eta_j = eta_j * (j_range)
j_l += eta_j
j_u -= eta_j
delta = float(0.0)
if p_range < 1.0:
w = w * (1.0 / (p_range * p_range))
if v_range < 1.0:
b_wv = b_wv * (1.0 / (v_range * v_range))
if a_range < 1.0:
b_wa = b_wa * (1.0 / (a_range * a_range))
w_a = w_a * (1.0 / (a_range * a_range))
if j_range < 1.0:
b_wj = b_wj * (1.0 / (j_range * j_range))
w_j = w_j * (1.0 / (j_range * j_range))
# position:
if n_w > 0.0:
error = c_p - target_p

1
src/curobo/rollout/cost/self_collision_cost.py
View File

@@ -68,7 +68,6 @@ class SelfCollisionCost(CostBase, SelfCollisionCostConfig):
self.self_collision_kin_config.thread_location,
self.self_collision_kin_config.thread_max,
self.self_collision_kin_config.checks_per_thread,
# False,
self.self_collision_kin_config.experimental_kernel,
self.return_loss,
)

2
src/curobo/types/math.py
View File

@@ -271,7 +271,7 @@ class Pose(Sequence):
if tensor_args is None and device is None:
log_error("Pose.to() requires tensor_args or device")
if tensor_args is not None:
t_type = vars(tensor_args.as_torch_dict())
t_type = tensor_args.as_torch_dict()
else:
t_type = {"device": device}
if self.position is not None:

2
src/curobo/types/robot.py
View File

@@ -43,6 +43,8 @@ class RobotConfig:
@staticmethod
def from_dict(data_dict_in, tensor_args=TensorDeviceType()):
if "robot_cfg" in data_dict_in:
data_dict_in = data_dict_in["robot_cfg"]
data_dict = deepcopy(data_dict_in)
if isinstance(data_dict["kinematics"], dict):
data_dict["kinematics"] = CudaRobotModelConfig.from_config(

78
src/curobo/util/logger.py
View File

@@ -8,11 +8,26 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
"""
This module provides logging API, wrapping :py:class:`logging.Logger`. These functions are used
to log messages in the cuRobo package. The functions can also be used in other packages by
creating a new logger (:py:meth:`setup_logger`) with the desired name.
"""
# Standard Library
import logging
import sys
def setup_curobo_logger(level="info"):
def setup_logger(level="info", logger_name: str = "curobo"):
"""Set up logger level.
Args:
level: Log level. Default is "info". Other options are "debug", "warning", "error".
logger_name: Name of the logger. Default is "curobo".
Raises:
ValueError: If log level is not one of [info, debug, warning, error].
"""
FORMAT = "[%(levelname)s] [%(name)s] %(message)s"
if level == "info":
level = logging.INFO
@@ -25,21 +40,64 @@ def setup_curobo_logger(level="info"):
else:
raise ValueError("Log level should be one of [info,debug, warn, error]")
logging.basicConfig(format=FORMAT, level=level)
logger = logging.getLogger("curobo")
logger = logging.getLogger(logger_name)
logger.setLevel(level=level)
def log_warn(txt: str, *args, **kwargs):
logger = logging.getLogger("curobo")
def setup_curobo_logger(level="info"):
"""Set up logger level for curobo package. Deprecated. Use :py:meth:`setup_logger` instead."""
return setup_logger(level, "curobo")
def log_warn(txt: str, logger_name: str = "curobo", *args, **kwargs):
"""Log warning message. Also see :py:meth:`logging.Logger.warning`.
Args:
txt: Warning message.
logger_name: Name of the logger. Default is "curobo".
"""
logger = logging.getLogger(logger_name)
logger.warning(txt, *args, **kwargs)
def log_info(txt: str, *args, **kwargs):
logger = logging.getLogger("curobo")
def log_info(txt: str, logger_name: str = "curobo", *args, **kwargs):
"""Log info message. Also see :py:meth:`logging.Logger.info`.
Args:
txt: Info message.
logger_name: Name of the logger. Default is "curobo".
"""
logger = logging.getLogger(logger_name)
logger.info(txt, *args, **kwargs)
def log_error(txt: str, exc_info=True, stack_info=True, *args, **kwargs):
logger = logging.getLogger("curobo")
logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs)
raise
def log_error(
txt: str,
logger_name: str = "curobo",
exc_info=True,
stack_info=False,
stacklevel: int = 2,
*args,
**kwargs
):
"""Log error and raise ValueError.
Args:
txt: Helpful message that conveys the error.
logger_name: Name of the logger. Default is "curobo".
exc_info: Add exception info to message. See :py:meth:`logging.Logger.error`.
stack_info: Add stacktracke to message. See :py:meth:`logging.Logger.error`.
stacklevel: See :py:meth:`logging.Logger.error`. Default value of 2 removes this function
from the stack trace.
Raises:
ValueError: Error message with exception.
"""
logger = logging.getLogger(logger_name)
if sys.version_info.major == 3 and sys.version_info.minor <= 7:
logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs)
else:
logger.error(
txt, exc_info=exc_info, stack_info=stack_info, stacklevel=stacklevel, *args, **kwargs
)
raise ValueError(txt)

3
src/curobo/util/metrics.py
View File

@@ -38,6 +38,7 @@ class CuroboMetrics(TrajectoryMetrics):
perception_success: bool = False
perception_interpolated_success: bool = False
jerk: float = np.inf
perception_time: float = 0.0
@dataclass
@@ -47,6 +48,7 @@ class CuroboGroupMetrics(TrajectoryGroupMetrics):
perception_success: float = 0.0
perception_interpolated_success: float = 0.0
jerk: float = np.inf
perception_time: float = np.inf
@classmethod
def from_list(cls, group: List[CuroboMetrics]):
@@ -60,5 +62,6 @@ class CuroboGroupMetrics(TrajectoryGroupMetrics):
[m.perception_interpolated_success for m in group]
)
data.jerk = Statistic.from_list([m.jerk for m in successes])
data.perception_time = Statistic.from_list([m.perception_time for m in successes])
return data

103
src/curobo/util/trajectory.py
View File

@@ -11,7 +11,7 @@
# Standard Library
import math
from enum import Enum
from typing import List, Optional
from typing import List, Optional, Tuple
# Third Party
import numpy as np
@@ -121,44 +121,58 @@ def get_spline_interpolated_trajectory(raw_traj, des_horizon, degree=5):
def get_batch_interpolated_trajectory(
raw_traj: JointState,
raw_dt: torch.Tensor,
interpolation_dt: float,
max_vel: torch.Tensor,
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
raw_dt: float,
max_vel: Optional[torch.Tensor] = None,
max_acc: Optional[torch.Tensor] = None,
max_jerk: Optional[torch.Tensor] = None,
kind: InterpolateType = InterpolateType.LINEAR_CUDA,
out_traj_state: Optional[JointState] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
max_deviation: float = 0.1,
min_dt: float = 0.02,
max_dt: float = 0.15,
optimize_dt: bool = True,
):
# compute dt across trajectory:
if len(raw_traj.shape) == 2:
raw_traj = raw_traj.unsqueeze(0)
b, horizon, dof = raw_traj.position.shape # horizon
# given the dt required to run trajectory at maximum velocity,
# we find the number of timesteps required:
traj_steps, steps_max, opt_dt = calculate_tsteps(
raw_traj.velocity,
raw_traj.acceleration,
raw_traj.jerk,
interpolation_dt,
max_vel,
max_acc,
max_jerk,
raw_dt,
min_dt,
horizon,
optimize_dt,
)
if optimize_dt:
traj_steps, steps_max, opt_dt = calculate_tsteps(
raw_traj.velocity,
raw_traj.acceleration,
raw_traj.jerk,
interpolation_dt,
max_vel,
max_acc,
max_jerk,
raw_dt,
min_dt,
max_dt,
horizon,
optimize_dt,
)
else:
traj_steps, steps_max = calculate_traj_steps(raw_dt, interpolation_dt, horizon)
opt_dt = raw_dt
# traj_steps contains the tsteps for each trajectory
assert steps_max > 0
# To do linear interpolation, we
if steps_max <= 0:
log_error("Steps max is less than 0")
if out_traj_state is not None and out_traj_state.position.shape[1] < steps_max:
log_warn(
"Interpolation buffer shape is smaller than steps_max,"
+ " creating new buffer of shape "
+ str(steps_max)
)
out_traj_state = None
if out_traj_state is None:
out_traj_state = JointState.zeros([b, steps_max, dof], tensor_args)
if out_traj_state.position.shape[1] < steps_max:
log_error("Interpolation buffer shape is smaller than steps_max")
if kind in [InterpolateType.LINEAR, InterpolateType.CUBIC]:
# plot and save:
out_traj_state = get_cpu_linear_interpolation(
@@ -187,7 +201,7 @@ def get_batch_interpolated_trajectory(
)
else:
raise ValueError("Unknown interpolation type")
# opt_dt = (raw_dt) / opt_dt
return out_traj_state, traj_steps, opt_dt
@@ -448,7 +462,9 @@ def calculate_dt_fixed(
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
raw_dt: torch.Tensor,
interpolation_dt: float,
min_dt: float,
max_dt: float,
epsilon: float = 1e-6,
):
# compute scaled dt:
max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof
@@ -465,9 +481,9 @@ def calculate_dt_fixed(
dt_score_jerk = raw_dt * torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3)
dt_score = torch.maximum(dt_score_vel, dt_score_acc)
dt_score = torch.maximum(dt_score, dt_score_jerk)
dt_score = torch.clamp(dt_score, interpolation_dt, raw_dt)
# NOTE: this dt score is not dt, rather a scaling to convert velocity, acc, jerk that was
# computed with raw_dt to a new dt
dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, max_dt)
return dt_score
@@ -480,7 +496,8 @@ def calculate_dt(
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
raw_dt: float,
interpolation_dt: float,
min_dt: float,
epsilon: float = 1e-6,
):
# compute scaled dt:
max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof
@@ -497,7 +514,8 @@ def calculate_dt(
dt_score_jerk = raw_dt * torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3)
dt_score = torch.maximum(dt_score_vel, dt_score_acc)
dt_score = torch.maximum(dt_score, dt_score_jerk)
dt_score = torch.clamp(dt_score, interpolation_dt, raw_dt)
dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, raw_dt)
# NOTE: this dt score is not dt, rather a scaling to convert velocity, acc, jerk that was
# computed with raw_dt to a new dt
return dt_score
@@ -511,6 +529,7 @@ def calculate_dt_no_clamp(
max_vel: torch.Tensor,
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
epsilon: float = 1e-6,
):
# compute scaled dt:
max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof
@@ -527,9 +546,19 @@ def calculate_dt_no_clamp(
dt_score_jerk = torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3)
dt_score = torch.maximum(dt_score_vel, dt_score_acc)
dt_score = torch.maximum(dt_score, dt_score_jerk)
dt_score = dt_score * (1.0 + epsilon)
return dt_score
@get_torch_jit_decorator()
def calculate_traj_steps(
opt_dt: torch.Tensor, interpolation_dt: float, horizon: int
) -> Tuple[torch.Tensor, torch.Tensor]:
traj_steps = (torch.ceil((horizon - 1) * ((opt_dt) / interpolation_dt))).to(dtype=torch.int32)
steps_max = torch.max(traj_steps)
return traj_steps, steps_max
@get_torch_jit_decorator()
def calculate_tsteps(
vel: torch.Tensor,
@@ -541,15 +570,23 @@ def calculate_tsteps(
max_jerk: torch.Tensor,
raw_dt: torch.Tensor,
min_dt: float,
max_dt: float,
horizon: int,
optimize_dt: bool = True,
):
# compute scaled dt:
opt_dt = calculate_dt_fixed(
vel, acc, jerk, max_vel, max_acc, max_jerk, raw_dt, interpolation_dt
vel,
acc,
jerk,
max_vel,
max_acc,
max_jerk,
raw_dt,
min_dt,
max_dt,
)
if not optimize_dt:
opt_dt[:] = raw_dt
traj_steps = (torch.ceil((horizon - 1) * ((opt_dt) / interpolation_dt))).to(dtype=torch.int32)
steps_max = torch.max(traj_steps)
traj_steps, steps_max = calculate_traj_steps(opt_dt, interpolation_dt, horizon)
return traj_steps, steps_max, opt_dt

7
src/curobo/wrap/model/robot_segmenter.py
View File

@@ -40,6 +40,7 @@ class RobotSegmenter:
distance_threshold: float = 0.05,
use_cuda_graph: bool = True,
ops_dtype: torch.dtype = torch.float32,
depth_to_meter: float = 0.001,
):
self._robot_world = robot_world
self._projection_rays = None
@@ -53,6 +54,7 @@ class RobotSegmenter:
self.tensor_args = robot_world.tensor_args
self.distance_threshold = distance_threshold
self._ops_dtype = ops_dtype
self._depth_to_meter = depth_to_meter
@staticmethod
def from_robot_file(
@@ -95,7 +97,10 @@ class RobotSegmenter:
if len(intrinsics.shape) == 2:
intrinsics = intrinsics.unsqueeze(0)
project_rays = get_projection_rays(
camera_obs.depth_image.shape[-2], camera_obs.depth_image.shape[-1], intrinsics
camera_obs.depth_image.shape[-2],
camera_obs.depth_image.shape[-1],
intrinsics,
self._depth_to_meter,
).to(dtype=self._ops_dtype)
if self._projection_rays is None:

3
src/curobo/wrap/model/robot_world.py
View File

@@ -18,8 +18,7 @@ from typing import Dict, List, Optional, Tuple, Union
import torch
# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.cuda_robot_model.types import CudaRobotModelState
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelState
from curobo.geom.sdf.utils import create_collision_checker
from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig
from curobo.geom.types import WorldConfig

9
src/curobo/wrap/reacher/evaluator.py
View File

@@ -78,7 +78,7 @@ class TrajEvaluatorConfig:
max_jerk: float = 500.0,
cost_weight: float = 0.01,
min_dt: float = 0.001,
max_dt: float = 0.1,
max_dt: float = 0.15,
tensor_args: TensorDeviceType = TensorDeviceType(),
) -> TrajEvaluatorConfig:
"""Creates TrajEvaluatorConfig object from basic parameters.
@@ -169,6 +169,7 @@ def compute_smoothness(
traj_dt: torch.Tensor,
min_dt: float,
max_dt: float,
epsilon: float = 1e-6,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""JIT compatible function to compute smoothness.
@@ -185,6 +186,8 @@ def compute_smoothness(
traj_dt: dt of trajectory tensor of shape (batch, horizon) or (1, horizon)
min_dt: minimum delta time allowed between steps/waypoints in a trajectory.
max_dt: maximum delta time allowed between steps/waypoints in a trajectory.
epsilon: relaxes evaluation of velocity, acceleration, and jerk limits to allow for
numerical errors.
Returns:
Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and
@@ -211,7 +214,7 @@ def compute_smoothness(
# clamp dt score:
dt_score = torch.clamp(dt_score, min_dt, max_dt)
dt_score = torch.clamp(dt_score * (1 + epsilon), min_dt, max_dt)
scale_dt = (1 / dt_score).view(-1, 1, 1)
abs_acc = torch.abs(acc) * (scale_dt**2)
# mean_acc_val = torch.max(torch.mean(abs_acc, dim=-1), dim=-1)[0]
@@ -376,7 +379,9 @@ class TrajEvaluator(TrajEvaluatorConfig):
label, cost = compute_smoothness_opt_dt(
js.velocity, js.acceleration, js.jerk, max_vel, self.max_acc, self.max_jerk, opt_dt
)
label = torch.logical_and(label, opt_dt <= self.max_dt)
pl_cost = compute_path_length_cost(js.velocity, cspace_distance_weight)
return label, pl_cost + self.cost_weight * cost

45
src/curobo/wrap/reacher/ik_solver.py
View File

@@ -9,9 +9,10 @@
# its affiliates is strictly prohibited.
#
"""
This module contains :meth:`IKSolver` to solve inverse kinematics, along with
:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result.
"""
This module contains :meth:`IKSolver` to solve inverse kinematics, along with
:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result. A minimal
example to solve IK is available at :ref:`python_ik_example`.
.. raw:: html
@@ -20,6 +21,7 @@ This module contains :meth:`IKSolver` to solve inverse kinematics, along with
</p>
This demo is available in :ref:`isaac_ik_reachability`.
"""
from __future__ import annotations
@@ -190,9 +192,9 @@ class IKSolverConfig:
grad_iters: Number of iterations for gradient optimization.
use_particle_opt: Flag to enable particle optimization.
collision_checker_type: Type of collision checker to use for collision checking.
sync_cuda_time: Flag to enable synchronization of cuda device with host before
measuring compute time. If you set this to False, then measured time will not
include completion of any launched CUDA kernels.
sync_cuda_time: Flag to enable synchronization of cuda device with host using
:py:func:`torch.cuda.synchronize` before measuring compute time. If you set this to
False, then measured time will not include completion of any launched CUDA kernels.
use_gradient_descent: Flag to enable gradient descent optimization instead of LBFGS.
collision_cache: Number of objects to cache for collision checking.
n_collision_envs: Number of collision environments to use for IK. This is useful when
@@ -1065,7 +1067,6 @@ class IKSolver(IKSolverConfig):
IKResult object with solutions to the IK problems.
"""
success = self._get_success(result.metrics, num_seeds=num_seeds)
if result.metrics.null_space_error is not None:
result.metrics.pose_error += result.metrics.null_space_error
if result.metrics.cspace_error is not None:
@@ -1524,6 +1525,36 @@ class IKSolver(IKSolverConfig):
if isinstance(rollout, ArmReacher)
]
def get_full_js(self, active_js: JointState) -> JointState:
"""Get full joint state from controlled joint state, appending locked joints.
Args:
active_js: Controlled joint state
Returns:
JointState: Full joint state.
"""
return self.rollout_fn.get_full_dof_from_solution(active_js)
def get_active_js(
self,
in_js: JointState,
):
"""Get controlled joint state from input joint state.
This is used to get the joint state for only joints that are optimization variables. This
also re-orders the joints to match the order of optimization variables.
Args:
in_js: Input joint state.
Returns:
JointState: Active joint state.
"""
opt_jnames = self.rollout_fn.joint_names
opt_js = in_js.get_ordered_joint_state(opt_jnames)
return opt_js
@property
def joint_names(self) -> List[str]:
"""Get ordered names of all joints used in optimization with IKSolver."""

3670
src/curobo/wrap/reacher/motion_gen.py
View File

File diff suppressed because it is too large Load Diff

701
src/curobo/wrap/reacher/mpc.py
View File

@@ -8,7 +8,31 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
"""
This module contains :meth:`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
:py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`.
A python example is available at :ref:`python_mpc_example`.
.. note::
Gradient-based MPC is also implemented with L-BFGS but is highly experimental and not
recommended for real robots.
.. raw:: html
<p>
<video autoplay="True" loop="True" muted="True" preload="auto" width="100%"><source src="../videos/mpc_clip.mp4" type="video/mp4"></video>
</p>
"""
# Standard Library
import time
@@ -19,6 +43,7 @@ from typing import Dict, Optional, Union
import torch
# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.geom.sdf.utils import create_collision_checker
from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig
from curobo.geom.types import WorldConfig
@@ -44,22 +69,31 @@ from curobo.wrap.wrap_mpc import WrapConfig, WrapMpc
@dataclass
class MpcSolverConfig:
"""Configuration dataclass for MPC."""
#: MPC Solver.
solver: WrapMpc
#: World Collision Checker.
world_coll_checker: Optional[WorldCollision] = None
#: Numeric precision and device to run computations.
tensor_args: TensorDeviceType = TensorDeviceType()
#: Capture full step in MPC as a single CUDA graph. This is not supported currently.
use_cuda_graph_full_step: bool = False
@staticmethod
def load_from_robot_config(
robot_cfg: Union[Union[str, dict], RobotConfig],
world_cfg: Union[Union[str, dict], WorldConfig],
world_model: Union[Union[str, dict], WorldConfig],
base_cfg: Optional[dict] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
compute_metrics: bool = True,
use_cuda_graph: Optional[bool] = None,
particle_opt_iters: Optional[int] = None,
self_collision_check: bool = True,
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE,
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH,
use_es: Optional[bool] = None,
es_learning_rate: Optional[float] = 0.01,
use_cuda_graph_metrics: bool = False,
@@ -74,9 +108,56 @@ class MpcSolverConfig:
use_lbfgs: bool = False,
use_mppi: bool = True,
):
"""Create an MPC solver configuration from robot and world configuration.
Args:
robot_cfg: Robot configuration. Can be a path to a YAML file or a dictionary or
an instance of :class:`~curobo.types.robot.RobotConfig`.
world_model: World configuration. Can be a path to a YAML file or a dictionary or
an instance of :class:`~curobo.geom.types.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 :ref:`world_collision`.
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
:py:func:`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 :py:meth:`MpcSolver.setup_solve_batch_env`
and :py:meth:`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
:attr:`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:
MpcSolverConfig: Configuration for the MPC solver.
"""
if use_cuda_graph_full_step:
log_error("use_cuda_graph_full_step currently is not supported")
raise ValueError("use_cuda_graph_full_step currently is not supported")
task_file = "particle_mpc.yml"
config_data = load_yaml(join_path(get_task_configs_path(), task_file))
@@ -108,14 +189,14 @@ class MpcSolverConfig:
if isinstance(robot_cfg, dict):
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
if isinstance(world_cfg, str):
world_cfg = load_yaml(join_path(get_world_configs_path(), world_cfg))
if isinstance(world_model, str):
world_model = load_yaml(join_path(get_world_configs_path(), world_model))
if world_coll_checker is None and world_cfg is not None:
world_cfg = WorldCollisionConfig.load_from_dict(
base_cfg["world_collision_checker_cfg"], world_cfg, tensor_args
if world_coll_checker is None and world_model is not None:
world_model = WorldCollisionConfig.load_from_dict(
base_cfg["world_collision_checker_cfg"], world_model, tensor_args
)
world_coll_checker = create_collision_checker(world_cfg)
world_coll_checker = create_collision_checker(world_model)
grad_config_data = None
if use_lbfgs:
grad_config_data = load_yaml(join_path(get_task_configs_path(), "gradient_mpc.yml"))
@@ -134,7 +215,7 @@ class MpcSolverConfig:
base_cfg["constraint"],
base_cfg["convergence"],
base_cfg["world_collision_checker_cfg"],
world_cfg,
world_model,
world_coll_checker=world_coll_checker,
tensor_args=tensor_args,
)
@@ -172,7 +253,7 @@ class MpcSolverConfig:
base_cfg["constraint"],
base_cfg["convergence"],
base_cfg["world_collision_checker_cfg"],
world_cfg,
world_model,
world_coll_checker=world_coll_checker,
tensor_args=tensor_args,
)
@@ -201,13 +282,42 @@ class MpcSolverConfig:
class MpcSolver(MpcSolverConfig):
"""Model Predictive Control Solver for Arm Reacher task.
"""High-level interface for Model Predictive Control (MPC).
Args:
MpcSolverConfig: _description_
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 :py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`.
See :ref:`python_mpc_example` for an example. This MPC solver implementation can be used in the
following steps:
1. Create a :py:class:`~curobo.rollout.rollout_base.Goal` object with the target pose or joint
configuration.
2. Create a goal buffer for the problem type using :meth:`setup_solve_single`,
:meth:`setup_solve_goalset`, :meth:`setup_solve_batch`, :meth:`setup_solve_batch_goalset`,
:meth:`setup_solve_batch_env`, or :meth:`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 :meth:`step` with the current joint state to get the next action.
4. To change the goal, create a :py:class:`~curobo.types.math.Pose` object with new pose or
:py:class:`~curobo.types.state.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 :meth:`update_goal` with the augmented goal buffer to update the goal for MPC.
6. Call :meth:`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 :meth:`enable_cspace_cost` and
:meth:`enable_pose_cost` to enable or disable reaching joint configuration cost and pose cost.
"""
def __init__(self, config: MpcSolverConfig) -> None:
"""Initializes the MPC solver.
Args:
config: Configuration parameters for MPC.
"""
super().__init__(**vars(config))
self.tensor_args = self.solver.rollout_fn.tensor_args
self._goal_buffer = Goal()
@@ -222,15 +332,326 @@ class MpcSolver(MpcSolverConfig):
self._cu_step_graph = None
self._cu_result = None
def _update_batch_size(self, batch_size):
if self.batch_size != batch_size:
self.batch_size = batch_size
def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a robot to reach target pose or joint configuration.
def update_goal_buffer(
Args:
goal: goal object containing target pose or joint configuration.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a robot to reach a pose in a set of poses.
Args:
goal: goal object containing target goalset or joint configuration.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.GOALSET,
num_mpc_seeds=num_seeds,
batch_size=1,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots to reach targets.
Args:
goal: goal object containing target poses or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots to reach a set of poses.
Args:
goal: goal object containing target goalset or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots in different collision worlds.
Args:
goal: goal object containing target poses or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots in different collision worlds.
Args:
goal: goal object containing target goalset or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def step(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
max_attempts: int = 1,
):
"""Solve for the next action given the current state.
Args:
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:
WrapResult: Result of the optimization.
"""
converged = True
for _ in range(max_attempts):
result = self._step_once(current_state.clone(), shift_steps, seed_traj)
if (
torch.count_nonzero(torch.isnan(result.action.position)) == 0
and torch.count_nonzero(~result.metrics.feasible) == 0
):
converged = True
break
self.reset()
if not converged:
result.action.copy_(current_state)
log_warn("MPC didn't converge")
return result
def update_goal(self, goal: Goal):
"""Update the goal for MPC.
Args:
goal: goal object containing target pose or joint configuration. This goal instance
should be created using one of the setup_solve functions.
"""
self.solver.update_params(goal)
def reset(self):
"""Reset the solver."""
# reset warm start
self.solver.reset()
def reset_cuda_graph(self):
"""Reset captured CUDA graph. This does not work."""
self.solver.reset_cuda_graph()
def enable_cspace_cost(self, enable=True):
"""Enable or disable reaching joint configuration cost in the solver.
Args:
enable: Enable or disable reaching joint configuration cost. When False, cspace cost
is disabled.
"""
self.solver.safety_rollout.enable_cspace_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_cspace_cost(enable)
def enable_pose_cost(self, enable=True):
"""Enable or disable reaching pose cost in the solver.
Args:
enable: Enable or disable reaching pose cost. When False, pose cost is disabled.
"""
self.solver.safety_rollout.enable_pose_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_pose_cost(enable)
def get_active_js(
self,
in_js: JointState,
):
"""Get controlled joints indexed in MPC order from the input joint state.
Args:
in_js: Input joint state.
Returns:
JointState: Joint state with controlled joints.
"""
opt_jnames = self.rollout_fn.joint_names
opt_js = in_js.get_ordered_joint_state(opt_jnames)
return opt_js
def update_world(self, 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 :attr:`MpcSolver.collision_cache` as
created during initialization of :class:`MpcSolverConfig`.
Args:
world: New collision world configuration. See :ref:`world_collision` for more details.
"""
self.world_coll_checker.load_collision_model(world)
def get_visual_rollouts(self):
"""Get rollouts for debugging."""
return self.solver.optimizers[0].get_rollouts()
@property
def joint_names(self):
"""Get the ordered joint names of the robot."""
return self.rollout_fn.joint_names
@property
def collision_cache(self) -> Dict[str, int]:
"""Returns the collision cache created by the world collision checker."""
return self.world_coll_checker.cache
@property
def kinematics(self) -> CudaRobotModel:
"""Get kinematics instance of the robot."""
return self.solver.safety_rollout.dynamics_model.robot_model
@property
def world_collision(self) -> WorldCollision:
"""Get the world collision checker."""
return self.world_coll_checker
@property
def rollout_fn(self) -> ArmReacher:
"""Get the rollout function."""
return self.solver.safety_rollout
def _step_once(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
"""Solve for the next action given the current state.
Args:
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:
WrapResult: Result of the optimization.
"""
# Create cuda graph for whole solve step including computation of metrics
# Including updation of goal buffers
if self._solve_state is None:
log_error("Need to first setup solve state before calling solve()")
if self.use_cuda_graph_full_step:
st_time = time.time()
if not self._cu_step_init:
self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
self._cu_state_in.copy_(current_state)
if seed_traj is not None:
self._cu_seed.copy_(seed_traj)
self._cu_step_graph.replay()
result = self._cu_result.clone()
torch.cuda.synchronize(device=self.tensor_args.device)
result.solve_time = time.time() - st_time
else:
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
self._step_goal_buffer,
shift_steps,
seed_traj,
)
return result
def _update_solve_state_and_goal_buffer(
self,
solve_state: ReacherSolveState,
goal: Goal,
) -> Goal:
"""Update solve state and goal for MPC.
Args:
solve_state: New solve state.
goal: New goal buffer.
Returns:
Goal: Updated goal buffer.
"""
self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal(
goal,
self._solve_state,
@@ -250,71 +671,64 @@ class MpcSolver(MpcSolverConfig):
)
return self._goal_buffer
def step(
def _update_batch_size(self, batch_size: int):
"""Update the batch size of the solver.
Args:
batch_size: Number of problems to solve in parallel.
"""
if self.batch_size != batch_size:
self.batch_size = batch_size
def _solve_from_solve_state(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
max_attempts: int = 1,
):
converged = True
for _ in range(max_attempts):
result = self.step_once(current_state.clone(), shift_steps, seed_traj)
if (
torch.count_nonzero(torch.isnan(result.action.position)) == 0
and torch.max(torch.abs(result.action.position)) < 10.0
and torch.count_nonzero(~result.metrics.feasible) == 0
):
converged = True
break
self.reset()
if not converged:
result.action.copy_(current_state)
log_warn("NOT CONVERGED")
return result
def step_once(
self,
current_state: JointState,
solve_state: ReacherSolveState,
goal: Goal,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
# Create cuda graph for whole solve step including computation of metrics
# Including updation of goal buffers
"""Solve for the next action given the current state.
if self._solve_state is None:
log_error("Need to first setup solve state before calling solve()")
Args:
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.
if self.use_cuda_graph_full_step:
st_time = time.time()
if not self._cu_step_init:
self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
self._cu_state_in.copy_(current_state)
if seed_traj is not None:
self._cu_seed.copy_(seed_traj)
self._cu_step_graph.replay()
result = self._cu_result.clone()
torch.cuda.synchronize()
result.solve_time = time.time() - st_time
else:
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
self._step_goal_buffer,
shift_steps,
seed_traj,
)
Returns:
WrapResult: Result of the optimization.
"""
if solve_state.batch_env:
if solve_state.batch_size > self.world_coll_checker.n_envs:
log_error("Batch Env is less that goal batch")
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
if seed_traj is not None:
self.solver.update_init_seed(seed_traj)
result = self.solver.solve(goal_buffer, seed_traj, shift_steps)
result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action)
return result
def _step(
def _mpc_step(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
):
"""One step function that is used to create a CUDA graph.
Args:
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:
WrapResult: Result of the optimization.
"""
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
@@ -331,6 +745,14 @@ class MpcSolver(MpcSolverConfig):
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
):
"""Create a CUDA graph for the full step of MPC.
Args:
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.
"""
log_info("MpcSolver: Creating Cuda Graph")
self._cu_state_in = current_state.clone()
if seed_traj is not None:
@@ -339,7 +761,7 @@ class MpcSolver(MpcSolverConfig):
s.wait_stream(torch.cuda.current_stream(device=self.tensor_args.device))
with torch.cuda.stream(s):
for _ in range(3):
self._cu_result = self._step(
self._cu_result = self._mpc_step(
self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed
)
torch.cuda.current_stream(device=self.tensor_args.device).wait_stream(s)
@@ -350,142 +772,3 @@ class MpcSolver(MpcSolverConfig):
self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed
)
self._cu_step_init = True
def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.GOALSET,
num_mpc_seeds=num_seeds,
batch_size=1,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def _solve_from_solve_state(
self,
solve_state: ReacherSolveState,
goal: Goal,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
if solve_state.batch_env:
if solve_state.batch_size > self.world_coll_checker.n_envs:
raise ValueError("Batch Env is less that goal batch")
goal_buffer = self.update_goal_buffer(solve_state, goal)
# NOTE: implement initialization from seed set here:
if seed_traj is not None:
self.solver.update_init_seed(seed_traj)
result = self.solver.solve(goal_buffer, seed_traj, shift_steps)
result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action)
return result
def fn(self):
# this will run one step of optimization and get new command
pass
def update_goal(self, goal: Goal):
self.solver.update_params(goal)
return True
def reset(self):
# reset warm start
self.solver.reset()
pass
def reset_cuda_graph(self):
self.solver.reset_cuda_graph()
@property
def rollout_fn(self):
return self.solver.safety_rollout
def enable_cspace_cost(self, enable=True):
self.solver.safety_rollout.enable_cspace_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_cspace_cost(enable)
def enable_pose_cost(self, enable=True):
self.solver.safety_rollout.enable_pose_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_pose_cost(enable)
def get_active_js(
self,
in_js: JointState,
):
opt_jnames = self.rollout_fn.joint_names
opt_js = in_js.get_ordered_joint_state(opt_jnames)
return opt_js
@property
def joint_names(self):
return self.rollout_fn.joint_names
def update_world(self, world: WorldConfig):
self.world_coll_checker.load_collision_model(world)
return True
def get_visual_rollouts(self):
return self.solver.optimizers[0].get_rollouts()
@property
def kinematics(self):
return self.solver.safety_rollout.dynamics_model.robot_model
@property
def world_collision(self):
return self.world_coll_checker

856
src/curobo/wrap/reacher/trajopt.py
View File

File diff suppressed because it is too large Load Diff

4
src/curobo/wrap/wrap_base.py
View File

@@ -127,6 +127,10 @@ class WrapBase(WrapConfig):
def rollout_fn(self):
return self.safety_rollout
@property
def tensor_args(self):
return self.safety_rollout.tensor_args
def solve(self, goal: Goal, seed: Optional[torch.Tensor] = None):
metrics = None

4
src/curobo/wrap/wrap_mpc.py
View File

@@ -53,13 +53,13 @@ class WrapMpc(WrapBase):
self.update_params(goal)
if self.sync_cuda_time:
torch.cuda.synchronize()
torch.cuda.synchronize(device=self.tensor_args.device)
# print("In: ", seed[0,:,0])
start_time = time.time()
with profiler.record_function("mpc/opt"):
act_seq = self.optimize(seed, shift_steps=shift_steps)
if self.sync_cuda_time:
torch.cuda.synchronize()
torch.cuda.synchronize(device=self.tensor_args.device)
self.opt_dt = time.time() - start_time
with profiler.record_function("mpc/filter"):
act = self.safety_rollout.get_robot_command(