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

update to 0.6.2

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2023-12-15 02:01:33 -08:00
parent d85ae41fba
commit 58958bbcce
105 changed files with 2514 additions and 934 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/curobo/content/assets/robot/franka_description/franka_panda.urdf
View File

@@ -243,10 +243,11 @@
<parent link="panda_hand"/>
<child link="panda_rightfinger"/>
<origin rpy="0 0 0" xyz="0 0 0.0584"/>
<axis xyz="0 1 0"/>
<axis xyz="0 -1 0"/>
<dynamics damping="10.0"/>
<limit effort="20" lower="-0.04" upper="0.0" velocity="0.2"/>
<limit effort="20" lower="0.0" upper="0.04" velocity="0.2"/>
</joint>
<link name="ee_link"/>
<joint name="ee_fixed_joint" type="fixed">

5
src/curobo/content/assets/robot/franka_description/franka_panda_mobile_xy_tz.urdf
View File

@@ -321,9 +321,10 @@
<parent link="panda_hand"/>
<child link="panda_rightfinger"/>
<origin rpy="0 0 0" xyz="0 0 0.0584"/>
<axis xyz="0 1 0"/>
<axis xyz="0 -1 0"/>
<dynamics damping="10.0"/>
<limit effort="20" lower="-0.04" upper="0.0" velocity="0.2"/>
<limit effort="20" lower="0.0" upper="0.04" velocity="0.2"/>
<mimic joint="panda_finger_joint1"/>
</joint>
</robot>

7
src/curobo/content/assets/robot/ur_description/ur10e.urdf
View File

@@ -356,4 +356,11 @@
<parent link="flange"/>
<child link="tool0"/>
</joint>
<joint name="camera_mount_joint" type="fixed">
<origin rpy="0 0 -3.14" xyz="0 0 0"/>
<parent link="flange"/>
<child link="camera_mount"/>
</joint>
<link name="camera_mount"/>
</robot>

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

@@ -47,6 +47,8 @@ robot_cfg:
}
lock_joints: null
mesh_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', 'wrist_2_link' ,'wrist_3_link' ,
'shoulder_link_1','upper_arm_link_1', 'forearm_link_1', 'wrist_1_link_1', 'wrist_2_link_1' ,'wrist_3_link_1']
cspace:
joint_names: [

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

@@ -28,6 +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"
@@ -84,7 +85,7 @@ robot_cfg:
"panda_rightfinger": 0.01,
"attached_object": 0.0,
}
#link_names: ["panda_link4"]
mesh_link_names:
[
"panda_link0",
@@ -99,14 +100,14 @@ robot_cfg:
"panda_leftfinger",
"panda_rightfinger",
]
lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": -0.04}
lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}
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:
joint_names: ["panda_joint1","panda_joint2","panda_joint3","panda_joint4", "panda_joint5",
"panda_joint6","panda_joint7","panda_finger_joint1", "panda_finger_joint2"]
retract_config: [0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0., 0.04, -0.04]
retract_config: [0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0., 0.04, 0.04]
null_space_weight: [1,1,1,1,1,1,1,1,1]
cspace_distance_weight: [1,1,1,1,1,1,1,1,1]
max_acceleration: 15.0

4
src/curobo/content/configs/robot/franka_mobile.yml
View File

@@ -100,7 +100,7 @@ robot_cfg:
"panda_leftfinger",
"panda_rightfinger",
]
lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": -0.04}
lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}
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" }}
@@ -109,7 +109,7 @@ robot_cfg:
"base_x", "base_y", "base_z",
"panda_joint1","panda_joint2","panda_joint3","panda_joint4", "panda_joint5",
"panda_joint6","panda_joint7","panda_finger_joint1", "panda_finger_joint2"]
retract_config: [0,0,0,0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0., 0.04, -0.04]
retract_config: [0,0,0,0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0., 0.04, 0.04]
null_space_weight: [1,1,1,1,1,1,1,1,1,1,1,1]
cspace_distance_weight: [1,1,1,1,1,1,1,1,1,1,1,1]
max_acceleration: 15.0 #15.0

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

@@ -32,7 +32,7 @@ robot_cfg:
- thumb_link_3
collision_sphere_buffer: 0.005
collision_spheres: spheres/iiwa_allegro.yml
ee_link: index_link_3
ee_link: palm_link
link_names:
- palm_link
- index_link_3

4
src/curobo/content/configs/robot/quad_ur10e.yml
View File

@@ -62,6 +62,10 @@ robot_cfg:
'shoulder_link_3': 0.05,
}
lock_joints: null
mesh_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', 'wrist_2_link' ,'wrist_3_link',
'shoulder_link_1','upper_arm_link_1', 'forearm_link_1', 'wrist_1_link_1', 'wrist_2_link_1' ,'wrist_3_link_1',
'shoulder_link_2','upper_arm_link_2', 'forearm_link_2', 'wrist_1_link_2', 'wrist_2_link_2' ,'wrist_3_link_2',
'shoulder_link_3','upper_arm_link_3', 'forearm_link_3', 'wrist_1_link_3', 'wrist_2_link_3' ,'wrist_3_link_3']
cspace:
joint_names: [

17
src/curobo/content/configs/robot/spheres/ur10e.yml
View File

@@ -11,18 +11,20 @@
robot: 'UR10'
collision_spheres:
shoulder_link:
- center: [0, 0, 0.0]
radius: 0.02
- center: [0, 0, 0]
radius: 0.05
#- center: [0, 0, -0.18]
# radius: 0.09
upper_arm_link:
- center: [-0, -0, 0.18]
radius: 0.09
- center: [-0.102167, 0, 0.18]
radius: 0.05
- center: [-0.204333, -8.88178e-19, 0.18]
- center: [-0.204333, 0, 0.18]
radius: 0.05
- center: [-0.3065, -1.33227e-18, 0.18]
- center: [-0.3065, 0, 0.18]
radius: 0.05
- center: [-0.408667, -1.77636e-18, 0.18]
- center: [-0.408667, 0, 0.18]
radius: 0.05
- center: [-0.510833, 0, 0.18]
radius: 0.05
@@ -53,10 +55,9 @@ collision_spheres:
- center: [0, 0, 0]
radius: 0.05
- center: [0, 0, 0.06]
radius: 0.08
radius: 0.07
tool0:
- center: [0, 0, 0.03]
- center: [0, 0, 0.12]
radius: 0.05
camera_mount:
- center: [0, 0.11, -0.01]

5
src/curobo/content/configs/robot/template.yml
View File

@@ -31,11 +31,12 @@ robot_cfg:
collision_spheres: null #
collision_sphere_buffer: 0.005
extra_collision_spheres: {}
self_collision_ignore: null # Dict[str, List[str]]
self_collision_buffer: null # Dict[str, float]
self_collision_ignore: {} # Dict[str, List[str]]
self_collision_buffer: {} # Dict[str, float]
use_global_cumul: True
mesh_link_names: null # List[str]
external_asset_path: null # Use this to add path for externally located assets/robot folder.
cspace:
joint_names: [] # List[str]

3
src/curobo/content/configs/robot/tri_ur10e.yml
View File

@@ -56,6 +56,9 @@ robot_cfg:
'shoulder_link_2': 0.05,
}
lock_joints: null
mesh_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', 'wrist_2_link' ,'wrist_3_link',
'shoulder_link_1','upper_arm_link_1', 'forearm_link_1', 'wrist_1_link_1', 'wrist_2_link_1' ,'wrist_3_link_1',
'shoulder_link_2','upper_arm_link_2', 'forearm_link_2', 'wrist_1_link_2', 'wrist_2_link_2' ,'wrist_3_link_2']
cspace:
joint_names: [

7
src/curobo/content/configs/robot/ur10e.yml
View File

@@ -22,7 +22,7 @@ robot_cfg:
link_names: null
collision_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', 'wrist_2_link' ,'wrist_3_link', 'tool0']
collision_spheres: 'spheres/ur10e.yml'
collision_sphere_buffer: 0.02
collision_sphere_buffer: 0.01
self_collision_ignore: {
"upper_arm_link": ["forearm_link", "shoulder_link"],
"forarm_link": ["wrist_1_link"],
@@ -49,4 +49,7 @@ robot_cfg:
null_space_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
cspace_distance_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] #[2.5, 1.75, 1.5, 1.25, 0.7, 0.4]
max_jerk: 500.0
max_acceleration: 15.0
max_acceleration: 12.0
position_limit_clip: 0.1
# add velocity scaling
# add joint position limit clipping

18
src/curobo/content/configs/robot/ur5e.yml
View File

@@ -27,7 +27,8 @@ robot_cfg:
extra_links: null
collision_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', 'wrist_2_link' ,'wrist_3_link'] # List[str]
collision_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link',
'wrist_2_link' ,'wrist_3_link', 'tool0'] # List[str]
collision_spheres:
shoulder_link:
- "center": [0.0, 0.0, 0.0]
@@ -89,14 +90,19 @@ robot_cfg:
wrist_3_link:
- "center": [0.001, 0.001, -0.029]
"radius": 0.043
tool0:
- "center": [0.001, 0.001, 0.05]
"radius": 0.05
collision_sphere_buffer: 0.005
extra_collision_spheres: {}
self_collision_ignore: {
"upper_arm_link": ["forearm_link", "shoulder_link"],
"forearm_link": ["wrist_1_link", "wrist_2_link", "wrist_3_link"],
"wrist_1_link": ["wrist_2_link","wrist_3_link"],
"wrist_2_link": ["wrist_3_link"],
"wrist_1_link": ["wrist_2_link","wrist_3_link","tool0"],
"wrist_2_link": ["wrist_3_link", "tool0"],
"wrist_3_link": ["tool0"],
}
self_collision_buffer: {'upper_arm_link': 0,
'forearm_link': 0,
@@ -106,7 +112,8 @@ robot_cfg:
}
use_global_cumul: True
mesh_link_names: null # List[str]
mesh_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link',
'wrist_2_link' ,'wrist_3_link'] # List[str]
cspace:
joint_names: ['shoulder_pan_joint', 'shoulder_lift_joint', 'elbow_joint', 'wrist_1_joint', 'wrist_2_joint', 'wrist_3_joint']
@@ -114,4 +121,5 @@ robot_cfg:
null_space_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
cspace_distance_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
max_jerk: 500.0
max_acceleration: 15.0
max_acceleration: 12.0
position_limit_clip: 0.1

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

@@ -22,6 +22,11 @@ cost:
weight: [0.0, 0.0]
vec_convergence: [0.0, 0.00] # orientation, position
terminal: False
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [0.0, 0.0]
vec_convergence: [0.0, 0.00] # orientation, position
terminal: False
bound_cfg:
@@ -43,9 +48,14 @@ constraint:
convergence:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [0.1,10.0] #[0.1, 100.0]
vec_convergence: [0.0, 0.0] # orientation, position
terminal: False
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [1.0, 1.0]
vec_convergence: [0.0, 0.00] # orientation, position
vec_convergence: [0.0, 0.0] # orientation, position
terminal: False
cspace_cfg:

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

@@ -31,8 +31,17 @@ 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.0,1.0,1.0,1.0,1.0] # running weight orientation, position
weight: [5000,30000.0,50,70] #[150.0, 2000.0, 30, 40]
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]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 0.0 #0.05
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]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 0.0
@@ -44,16 +53,16 @@ cost:
run_weight: 0.00 #1
bound_cfg:
weight: [5000.0, 5000.0, 5000.0, 5000.0] # needs to be 3 values
smooth_weight: [0.0,1000.0, 500.0, 0.0] # [vel, acc, jerk,]
weight: [5000.0, 50000.0, 50000.0, 50000.0] # needs to be 3 values
smooth_weight: [0.0,2000.0, 50.0, 0.0] # [vel, acc, jerk,]
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
null_space_weight: [1.0]
null_space_weight: [0.0]
primitive_collision_cfg:
weight: 50000.0
weight: 1000000.0
use_sweep: True
sweep_steps: 6
classify: False
@@ -69,17 +78,17 @@ cost:
lbfgs:
n_iters: 300 #125 #@200 #250 #250 # 150 #25
inner_iters: 25 #$25 # 25
cold_start_n_iters: 300 #125 #200 #250 #$150 #25
n_iters: 400
inner_iters: 25
cold_start_n_iters: 200
min_iters: 25
line_search_scale: [0.01,0.3,0.7,1.0] # #
fixed_iters: False
line_search_scale: [0.01, 0.3,0.7,1.0] # #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 10.0 #10.0 # 5.0 #2.0
cost_delta_threshold: 1.0
cost_relative_threshold: 0.999 #0.999
epsilon: 0.01
history: 4 #15
history: 15 #15
use_cuda_graph: True
n_envs: 1
store_debug: False
@@ -91,7 +100,7 @@ lbfgs:
use_temporal_smooth: False
sync_cuda_time: True
step_scale: 1.0
last_best: 5
last_best: 10
use_coo_sparse: True
debug_info:
visual_traj : null #'ee_pos_seq'

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

@@ -28,8 +28,13 @@ model:
cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
#weight: [100, 500, 10, 20]
weight: [500, 5000, 30, 30]
weight: [1000, 20000, 30, 50]
vec_convergence: [0.0, 0.00]
terminal: False
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [1000, 20000, 30, 50]
vec_convergence: [0.00, 0.000]
terminal: False
use_metric: True
@@ -37,31 +42,31 @@ cost:
cspace_cfg:
weight: 0.000
bound_cfg:
weight: 50.0
weight: 100.0
activation_distance: [0.1]
null_space_weight: [1.0]
primitive_collision_cfg:
weight: 1000.0
weight: 50000.0
use_sweep: False
classify: False
activation_distance: 0.025
activation_distance: 0.01
self_collision_cfg:
weight: 1000.0
weight: 500.0
classify: False
lbfgs:
n_iters: 80 #60
inner_iters: 20
cold_start_n_iters: 80
n_iters: 100 #60
inner_iters: 25
cold_start_n_iters: 100
min_iters: 20
line_search_scale: [0.01, 0.3, 0.7, 1.0] #[0.01,0.4, 0.9, 1.0] # #
line_search_scale: [0.01, 0.3, 0.7, 1.0]
fixed_iters: True
cost_convergence: 1e-7
cost_delta_threshold: 1e-6 #0.0001
cost_relative_threshold: 1.0
epsilon: 0.01 # used only in stable_mode
history: 4
epsilon: 0.01 #0.01 # used only in stable_mode
history: 6
horizon: 1
use_cuda_graph: True
n_envs: 1

14
src/curobo/content/configs/task/gradient_mpc.yml
View File

@@ -10,7 +10,7 @@
##
model:
horizon: 40
horizon: 30
state_filter_cfg:
filter_coeff:
position: 0.0
@@ -18,9 +18,9 @@ model:
acceleration: 0.0
enable: True
dt_traj_params:
base_dt: 0.02
base_dt: 0.05
base_ratio: 0.5
max_dt: 0.02
max_dt: 0.05
vel_scale: 1.0
control_space: 'POSITION'
teleport_mode: False
@@ -72,16 +72,16 @@ cost:
max_nlimit: 0.5 #0.2
lbfgs:
n_iters: 150 #125 #@200 #250 #250 # 150 #25
n_iters: 500 #125 #@200 #250 #250 # 150 #25
inner_iters: 25
cold_start_n_iters: 500 #125 #200 #250 #$150 #25
min_iters: 50
line_search_scale: [0.01,0.25, 0.75,1.0] #[0.01,0.25,0.7, 1.0] # [0.01, 0.8, 1.0] #
line_search_scale: [0.01,0.3, 0.7,1.0] #[0.01,0.25,0.7, 1.0] # [0.01, 0.8, 1.0] #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 0.0001
epsilon: 0.01
history: 6 #15
history: 15 #15
use_cuda_graph: True
n_envs: 1
store_debug: False
@@ -92,7 +92,7 @@ lbfgs:
use_cuda_update_best_kernel: True
sync_cuda_time: True
use_temporal_smooth: False
last_best: 26
last_best: 10
step_scale: 1.0
use_coo_sparse: True
debug_info:

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

@@ -32,10 +32,19 @@ 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,20000.0,30,30] #[150.0, 2000.0, 30, 40]
weight: [2000,50000.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.0 #0.05
run_weight: 0.00
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,50000.0,30,50] #[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.0
use_metric: True
cspace_cfg:
@@ -44,18 +53,18 @@ cost:
run_weight: 0.00 #1
bound_cfg:
weight: [5000.0, 5000.0, 5000.0,5000.0] # needs to be 3 values
smooth_weight: [0.0,3000.0,10.0, 0.0] # [vel, acc, jerk,]
weight: [5000.0, 50000.0, 50000.0,50000.0] # needs to be 3 values
smooth_weight: [0.0,10000.0,10.0, 0.0] # [vel, acc, jerk, alpha_v-not-used]
run_weight_velocity: 0.00
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
null_space_weight: [1.0]
null_space_weight: [0.0]
primitive_collision_cfg:
weight: 10000.0
weight: 100000.0
use_sweep: True
sweep_steps: 4
sweep_steps: 6
classify: False
use_sweep_kernel: True
use_speed_metric: True
@@ -70,17 +79,17 @@ cost:
lbfgs:
n_iters: 100 #125 #@200 #250 #250 # 150 #25
inner_iters: 25 #25 # 25
cold_start_n_iters: 100 #125 #200 #250 #$150 #25
n_iters: 175 #175
inner_iters: 25
cold_start_n_iters: 150
min_iters: 25
line_search_scale: [0.01,0.3,0.7,1.0] #[0.01,0.2, 0.3,0.5,0.7,0.9, 1.0] #
fixed_iters: False
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 1.0
cost_delta_threshold: 2000.0
cost_relative_threshold: 0.9999
epsilon: 0.01
history: 4 #15 #$14
history: 15
use_cuda_graph: True
n_envs: 1
store_debug: False
@@ -92,7 +101,7 @@ lbfgs:
use_temporal_smooth: False
sync_cuda_time: True
step_scale: 1.0 #1.0
last_best: 5
last_best: 10
use_coo_sparse: True
debug_info:
visual_traj : null #'ee_pos_seq'

2
src/curobo/content/configs/task/graph.yml
View File

@@ -33,7 +33,7 @@ graph:
sample_pts: 1500
node_similarity_distance: 0.1
rejection_ratio: 20 #$20
rejection_ratio: 20
k_nn: 15
max_buffer: 10000
max_cg_buffer: 1000

6
src/curobo/content/configs/task/particle_ik.yml
View File

@@ -31,6 +31,12 @@ cost:
vec_convergence: [0.00, 0.000] # orientation, position
terminal: False
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [30, 50, 10, 10] #[20.0, 100.0]
vec_convergence: [0.00, 0.000] # orientation, position
terminal: False
use_metric: True
cspace_cfg:
weight: 0.00
bound_cfg:

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

@@ -29,6 +29,15 @@ model:
cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
run_vec_weight: [1.0,1.0,1.0,1.0,1.0,1.0] # running weight
weight: [250.0, 5000.0, 20, 20]
vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True
run_weight: 0.00
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
run_vec_weight: [1.0,1.0,1.0,1.0,1.0,1.0] # running weight
weight: [250.0, 5000.0, 40, 40]
@@ -38,14 +47,14 @@ cost:
use_metric: True
cspace_cfg:
weight: 500.0
weight: 10000.0
terminal: True
run_weight: 0.001
run_weight: 0.0
bound_cfg:
weight: [0.1, 0.1,0.0,0.0]
activation_distance: [0.0,0.0,0.0,0.0] #-#0.01
smooth_weight: [0.0,100.0,1.0,0.0]
smooth_weight: [0.0,20.0,0.0,0.0]
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
@@ -68,7 +77,7 @@ cost:
mppi:
init_cov : 0.1 #0.5
init_cov : 0.5
gamma : 1.0
n_iters : 2
cold_start_n_iters: 2

2
src/curobo/content/configs/world/collision_nvblox_online.yml
View File

@@ -12,6 +12,6 @@ blox:
world:
pose: [0,0,0,1,0,0,0]
integrator_type: "tsdf"
voxel_size: 0.03
voxel_size: 0.01

38
src/curobo/cuda_robot_model/cuda_robot_generator.py
View File

@@ -136,6 +136,9 @@ class CudaRobotGeneratorConfig:
use_external_assets: bool = False
external_asset_path: Optional[str] = None
external_robot_configs_path: Optional[str] = None
#: Create n collision spheres for links with name
extra_collision_spheres: Optional[Dict[str, int]] = None
@@ -144,18 +147,20 @@ class CudaRobotGeneratorConfig:
def __post_init__(self):
# add root path:
if self.urdf_path is not None:
self.urdf_path = join_path(get_assets_path(), self.urdf_path)
if self.usd_path is not None:
self.usd_path = join_path(get_assets_path(), self.usd_path)
if self.asset_root_path != "":
if self.use_external_assets:
with importlib_resources.files("content") as path:
content_dir_posix = path
self.asset_root_path = join_path(content_dir_posix, self.asset_root_path)
# Check if an external asset path is provided:
asset_path = get_assets_path()
robot_path = get_robot_configs_path()
if self.external_asset_path is not None:
asset_path = self.external_asset_path
if self.external_robot_configs_path is not None:
robot_path = self.external_robot_configs_path
else:
self.asset_root_path = join_path(get_assets_path(), self.asset_root_path)
if self.urdf_path is not None:
self.urdf_path = join_path(asset_path, self.urdf_path)
if self.usd_path is not None:
self.usd_path = join_path(asset_path, self.usd_path)
if self.asset_root_path != "":
self.asset_root_path = join_path(asset_path, self.asset_root_path)
elif self.urdf_path is not None:
self.asset_root_path = os.path.dirname(self.urdf_path)
@@ -178,7 +183,7 @@ class CudaRobotGeneratorConfig:
self.link_names.append(self.ee_link)
if self.collision_spheres is not None:
if isinstance(self.collision_spheres, str):
coll_yml = join_path(get_robot_configs_path(), self.collision_spheres)
coll_yml = join_path(robot_path, self.collision_spheres)
coll_params = load_yaml(coll_yml)
self.collision_spheres = coll_params["collision_spheres"]
@@ -399,7 +404,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
joint_map = [
-1 if i not in self._controlled_joints else i for i in range(len(self._bodies))
]
] #
joint_map_type = [
-1 if i not in self._controlled_joints else i for i in range(len(self._bodies))
]
@@ -885,4 +890,11 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
self.cspace.max_acceleration.unsqueeze(0),
]
)
# clip joint position:
# TODO: change this to be per joint
joint_limits["position"][0] += self.cspace.position_limit_clip
joint_limits["position"][1] -= self.cspace.position_limit_clip
joint_limits["velocity"][0] *= self.cspace.velocity_scale
joint_limits["velocity"][1] *= self.cspace.velocity_scale
self._joint_limits = JointLimits(joint_names=self.joint_names, **joint_limits)

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

@@ -116,7 +116,8 @@ class CudaRobotModelConfig:
@staticmethod
def from_data_dict(
data_dict: Dict[str, Any], tensor_args: TensorDeviceType = TensorDeviceType()
data_dict: Dict[str, Any],
tensor_args: TensorDeviceType = TensorDeviceType(),
):
return CudaRobotModelConfig.from_config(
CudaRobotGeneratorConfig(**data_dict, tensor_args=tensor_args)

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

@@ -34,6 +34,12 @@ class JointType(Enum):
X_ROT = 3
Y_ROT = 4
Z_ROT = 5
X_PRISM_NEG = 6
Y_PRISM_NEG = 7
Z_PRISM_NEG = 8
X_ROT_NEG = 9
Y_ROT_NEG = 10
Z_ROT_NEG = 11
@dataclass
@@ -82,6 +88,7 @@ class CSpaceConfig:
velocity_scale: Union[float, List[float]] = 1.0
acceleration_scale: Union[float, List[float]] = 1.0
jerk_scale: Union[float, List[float]] = 1.0
position_limit_clip: Union[float, List[float]] = 0.01
def __post_init__(self):
if self.retract_config is not None:

19
src/curobo/cuda_robot_model/urdf_kinematics_parser.py
View File

@@ -95,7 +95,7 @@ class UrdfKinematicsParser(KinematicsParser):
"velocity": joint.limit.velocity,
}
else:
# log_warn("Converting continuous joint to revolute")
log_warn("Converting continuous joint to revolute")
joint_type = "revolute"
joint_limits = {
"effort": joint.limit.effort,
@@ -105,9 +105,7 @@ class UrdfKinematicsParser(KinematicsParser):
}
joint_axis = joint.axis
if (joint_axis < 0).any():
log_warn("Joint axis has negative value (-1). This is not supported.")
joint_axis = np.abs(joint_axis)
body_params["joint_limits"] = [joint_limits["lower"], joint_limits["upper"]]
body_params["joint_velocity_limits"] = [
-1.0 * joint_limits["velocity"],
@@ -128,6 +126,13 @@ class UrdfKinematicsParser(KinematicsParser):
joint_type = JointType.Y_PRISM
if joint_axis[2] == 1:
joint_type = JointType.Z_PRISM
if joint_axis[0] == -1:
joint_type = JointType.X_PRISM_NEG
if joint_axis[1] == -1:
joint_type = JointType.Y_PRISM_NEG
if joint_axis[2] == -1:
joint_type = JointType.Z_PRISM_NEG
elif joint_type == "revolute":
if joint_axis[0] == 1:
joint_type = JointType.X_ROT
@@ -135,6 +140,12 @@ class UrdfKinematicsParser(KinematicsParser):
joint_type = JointType.Y_ROT
if joint_axis[2] == 1:
joint_type = JointType.Z_ROT
if joint_axis[0] == -1:
joint_type = JointType.X_ROT_NEG
if joint_axis[1] == -1:
joint_type = JointType.Y_ROT_NEG
if joint_axis[2] == -1:
joint_type = JointType.Z_ROT_NEG
else:
log_error("Joint type not supported")

306
src/curobo/curobolib/cpp/kinematics_fused_kernel.cu
View File

@@ -28,8 +28,15 @@
#define Y_ROT 4
#define Z_ROT 5
#define X_PRISM_NEG 6
#define Y_PRISM_NEG 7
#define Z_PRISM_NEG 8
#define X_ROT_NEG 9
#define Y_ROT_NEG 10
#define Z_ROT_NEG 11
#define MAX_BATCH_PER_BLOCK 32 // tunable parameter for improving occupancy
#define MAX_BW_BATCH_PER_BLOCK 8 // tunable parameter for improving occupancy
#define MAX_BW_BATCH_PER_BLOCK 16 // tunable parameter for improving occupancy
#define MAX_TOTAL_LINKS \
750 // limited by shared memory size. We need to fit 16 * float32 per link
@@ -238,67 +245,46 @@ __device__ __forceinline__ void prism_fn(const scalar_t *fixedTransform,
}
}
template <typename scalar_t>
__device__ __forceinline__ void xprism_fn(const scalar_t *fixedTransform,
const float angle, const int col_idx,
float *JM) {
__device__ __forceinline__ void update_joint_type_direction(int &j_type, int8_t &axis_sign)
{
// Assume that input j_type >= 0 . Check fixed joint outside of this function.
switch (col_idx) {
case 0:
case 1:
case 2:
fixed_joint_fn(&fixedTransform[col_idx], &JM[0]);
break;
case 3:
JM[0] = fixedTransform[0] * angle + fixedTransform[3]; // FT_0[1];
JM[1] = fixedTransform[M] * angle + fixedTransform[M + 3]; // FT_1[1];
JM[2] =
fixedTransform[M + M] * angle + fixedTransform[M + M + 3]; // FT_2[1];
JM[3] = 1;
break;
// Don't do anything if j_type < 6. j_type range is [0, 11]
// divergence here.
axis_sign = 1;
if (j_type >= 6)
{
j_type -= 6;
axis_sign = -1;
}
}
template <typename scalar_t>
__device__ __forceinline__ void yprism_fn(const scalar_t *fixedTransform,
const float angle, const int col_idx,
float *JM) {
switch (col_idx) {
case 0:
case 1:
case 2:
fixed_joint_fn(&fixedTransform[col_idx], &JM[0]);
break;
case 3:
JM[0] = fixedTransform[1] * angle + fixedTransform[3]; // FT_0[1];
JM[1] = fixedTransform[M + 1] * angle + fixedTransform[M + 3]; // FT_1[1];
JM[2] = fixedTransform[M + M + 1] * angle +
fixedTransform[M + M + 3]; // FT_2[1];
JM[3] = 1;
break;
__device__ __forceinline__ void update_joint_type_direction(int &j_type)
{
// Assume that input j_type >= 0 . Check fixed joint outside of this function.
// Don't do anything if j_type < 6. j_type range is [0, 11]
// divergence here.
if (j_type >= 6)
{
j_type -= 6;
}
}
template <typename scalar_t>
__device__ __forceinline__ void zprism_fn(const scalar_t *fixedTransform,
const float angle, const int col_idx,
float *JM) {
switch (col_idx) {
case 0:
case 1:
case 2:
fixed_joint_fn(&fixedTransform[col_idx], &JM[0]);
break;
case 3:
JM[0] = fixedTransform[2] * angle + fixedTransform[3]; // FT_0[1];
JM[1] = fixedTransform[M + 2] * angle + fixedTransform[M + 3]; // FT_1[1];
JM[2] = fixedTransform[M + M + 2] * angle +
fixedTransform[M + M + 3]; // FT_2[1];
JM[3] = 1;
break;
}
__device__ __forceinline__ void update_axis_direction(
float &angle,
int &j_type)
{
// Assume that input j_type >= 0 . Check fixed joint outside of this function.
// sign should be +ve <= 5 and -ve >5
// j_type range is [0, 11].
// cuda code treats -1.0 * 0.0 as negative. Hence we subtract 6. If in future, -1.0 * 0.0 = +ve,
// then this code should be j_type - 5.
angle = -1 * copysignf(1.0, j_type - 6) * angle;
update_joint_type_direction(j_type);
}
// In the following versions of rot_fn, some non-nan values may become nan as we
@@ -423,7 +409,7 @@ __device__ __forceinline__ void zrot_fn(const scalar_t *fixedTransform,
template <typename psum_t>
__device__ __forceinline__ void
rot_backward_translation(const float3 &vec, float *cumul_mat, float *l_pos,
const float3 &loc_grad, psum_t &grad_q) {
const float3 &loc_grad, psum_t &grad_q, const int8_t axis_sign=1) {
float3 e_pos, j_pos;
@@ -433,136 +419,141 @@ rot_backward_translation(const float3 &vec, float *cumul_mat, float *l_pos,
// compute position gradient:
j_pos = *(float3 *)&l_pos[0] - e_pos; // - e_pos;
scale_cross_sum(vec, j_pos, loc_grad, grad_q); // cross product
float3 scale_grad = axis_sign * loc_grad;
scale_cross_sum(vec, j_pos, scale_grad, grad_q); // cross product
}
template <typename psum_t>
__device__ __forceinline__ void
rot_backward_rotation(const float3 vec, const float3 grad_vec, psum_t &grad_q, const int8_t axis_sign=1) {
grad_q += axis_sign * dot(vec, grad_vec);
}
template <typename psum_t>
__device__ __forceinline__ void
prism_backward_translation(const float3 vec, const float3 grad_vec,
psum_t &grad_q) {
grad_q += dot(vec, grad_vec);
psum_t &grad_q, const int8_t axis_sign=1) {
grad_q += axis_sign * dot(vec, grad_vec);
}
template <typename psum_t>
__device__ __forceinline__ void
rot_backward_rotation(const float3 vec, const float3 grad_vec, psum_t &grad_q) {
grad_q += dot(vec, grad_vec);
}
template <typename psum_t>
__device__ __forceinline__ void
z_rot_backward(float *link_cumul_mat, float *l_pos, float3 &loc_grad_position,
float3 &loc_grad_orientation, psum_t &grad_q) {
float3 &loc_grad_orientation, psum_t &grad_q, const int8_t axis_sign=1) {
float3 vec =
make_float3(link_cumul_mat[2], link_cumul_mat[6], link_cumul_mat[10]);
// get rotation vector:
rot_backward_translation(vec, &link_cumul_mat[0], &l_pos[0],
loc_grad_position, grad_q);
loc_grad_position, grad_q, axis_sign);
rot_backward_rotation(vec, loc_grad_orientation, grad_q);
rot_backward_rotation(vec, loc_grad_orientation, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void
x_rot_backward(float *link_cumul_mat, float *l_pos, float3 &loc_grad_position,
float3 &loc_grad_orientation, psum_t &grad_q) {
float3 &loc_grad_orientation, psum_t &grad_q, const int8_t axis_sign=1) {
float3 vec =
make_float3(link_cumul_mat[0], link_cumul_mat[4], link_cumul_mat[8]);
// get rotation vector:
rot_backward_translation(vec, &link_cumul_mat[0], &l_pos[0],
loc_grad_position, grad_q);
loc_grad_position, grad_q, axis_sign);
rot_backward_rotation(vec, loc_grad_orientation, grad_q);
rot_backward_rotation(vec, loc_grad_orientation, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void
y_rot_backward(float *link_cumul_mat, float *l_pos, float3 &loc_grad_position,
float3 &loc_grad_orientation, psum_t &grad_q) {
float3 &loc_grad_orientation, psum_t &grad_q, const int8_t axis_sign=1) {
float3 vec =
make_float3(link_cumul_mat[1], link_cumul_mat[5], link_cumul_mat[9]);
// get rotation vector:
rot_backward_translation(vec, &link_cumul_mat[0], &l_pos[0],
loc_grad_position, grad_q);
loc_grad_position, grad_q, axis_sign);
rot_backward_rotation(vec, loc_grad_orientation, grad_q);
rot_backward_rotation(vec, loc_grad_orientation, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void
xyz_prism_backward_translation(float *cumul_mat, float3 &loc_grad,
psum_t &grad_q, int xyz) {
psum_t &grad_q, int xyz, const int8_t axis_sign=1) {
prism_backward_translation(
make_float3(cumul_mat[0 + xyz], cumul_mat[4 + xyz], cumul_mat[8 + xyz]),
loc_grad, grad_q);
loc_grad, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void x_prism_backward_translation(float *cumul_mat,
float3 &loc_grad,
psum_t &grad_q) {
psum_t &grad_q,
const int8_t axis_sign=1) {
// get rotation vector:
prism_backward_translation(
make_float3(cumul_mat[0], cumul_mat[4], cumul_mat[8]), loc_grad, grad_q);
make_float3(cumul_mat[0], cumul_mat[4], cumul_mat[8]), loc_grad, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void y_prism_backward_translation(float *cumul_mat,
float3 &loc_grad,
psum_t &grad_q) {
psum_t &grad_q,
const int8_t axis_sign=1) {
// get rotation vector:
prism_backward_translation(
make_float3(cumul_mat[1], cumul_mat[5], cumul_mat[9]), loc_grad, grad_q);
make_float3(cumul_mat[1], cumul_mat[5], cumul_mat[9]), loc_grad, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void z_prism_backward_translation(float *cumul_mat,
float3 &loc_grad,
psum_t &grad_q) {
psum_t &grad_q,
const int8_t axis_sign=1) {
// get rotation vector:
prism_backward_translation(
make_float3(cumul_mat[2], cumul_mat[6], cumul_mat[10]), loc_grad, grad_q);
make_float3(cumul_mat[2], cumul_mat[6], cumul_mat[10]), loc_grad, grad_q, axis_sign);
}
__device__ __forceinline__ void
xyz_rot_backward_translation(float *cumul_mat, float *l_pos, float3 &loc_grad,
float &grad_q, int xyz) {
float &grad_q, int xyz, const int8_t axis_sign=1) {
// get rotation vector:
rot_backward_translation(
make_float3(cumul_mat[0 + xyz], cumul_mat[4 + xyz], cumul_mat[8 + xyz]),
&cumul_mat[0], &l_pos[0], loc_grad, grad_q);
&cumul_mat[0], &l_pos[0], loc_grad, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void
x_rot_backward_translation(float *cumul_mat, float *l_pos, float3 &loc_grad,
psum_t &grad_q) {
psum_t &grad_q, const int8_t axis_sign=1) {
// get rotation vector:
rot_backward_translation(
make_float3(cumul_mat[0], cumul_mat[4], cumul_mat[8]), &cumul_mat[0],
&l_pos[0], loc_grad, grad_q);
&l_pos[0], loc_grad, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void
y_rot_backward_translation(float *cumul_mat, float *l_pos, float3 &loc_grad,
psum_t &grad_q) {
psum_t &grad_q, const int8_t axis_sign=1) {
// get rotation vector:
rot_backward_translation(
make_float3(cumul_mat[1], cumul_mat[5], cumul_mat[9]), &cumul_mat[0],
&l_pos[0], loc_grad, grad_q);
&l_pos[0], loc_grad, grad_q, axis_sign);
}
template <typename psum_t>
__device__ __forceinline__ void
z_rot_backward_translation(float *cumul_mat, float *l_pos, float3 &loc_grad,
psum_t &grad_q) {
psum_t &grad_q, const int8_t axis_sign=1) {
// get rotation vector:
rot_backward_translation(
make_float3(cumul_mat[2], cumul_mat[6], cumul_mat[10]), &cumul_mat[0],
&l_pos[0], loc_grad, grad_q);
&l_pos[0], loc_grad, grad_q, axis_sign);
}
// An optimized version of kin_fused_warp_kernel.
@@ -604,7 +595,7 @@ kin_fused_warp_kernel2(scalar_t *link_pos, // batchSize xz store_n_links x M x M
*(float4 *)&global_cumul_mat[batch * nlinks * 16 + col_idx * M] =
*(float4 *)&cumul_mat[matAddrBase + col_idx * M];
}
for (int8_t l = 1; l < nlinks; l++) // TODO: add base link transform
for (int8_t l = 1; l < nlinks; l++) //
{
// get one row of fixedTransform
@@ -616,21 +607,27 @@ kin_fused_warp_kernel2(scalar_t *link_pos, // batchSize xz store_n_links x M x M
// check joint type and use one of the helper functions:
float JM[M];
int j_type = jointMapType[l];
float angle = q[batch * njoints + jointMap[l]];
if (j_type == FIXED) {
if (j_type == FIXED)
{
fixed_joint_fn(&fixedTransform[ftAddrStart + col_idx], &JM[0]);
} else if (j_type <= Z_PRISM) {
prism_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0], j_type);
} else if (j_type == X_ROT) {
xrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Y_ROT) {
yrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Z_ROT) {
zrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else {
assert(jointMapType[l] > -2 &&
jointMapType[l] < 6); // joint type not supported
}
}
else
{
float angle = q[batch * njoints + jointMap[l]];
update_axis_direction(angle, j_type);
if (j_type <= Z_PRISM) {
prism_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0], j_type);
} else if (j_type == X_ROT) {
xrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Y_ROT) {
yrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Z_ROT) {
zrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else {
assert(j_type >= FIXED && j_type <= Z_ROT);
}
}
#pragma unroll 4
for (int i = 0; i < M; i++) {
@@ -652,7 +649,8 @@ kin_fused_warp_kernel2(scalar_t *link_pos, // batchSize xz store_n_links x M x M
int16_t read_cumul_idx = -1;
int16_t spheres_perthread = (nspheres + 3) / 4;
for (int16_t i = 0; i < spheres_perthread; i++) {
const int16_t sph_idx = col_idx * spheres_perthread + i;
//const int16_t sph_idx = col_idx * spheres_perthread + i;
const int16_t sph_idx = col_idx + i * 4;
// const int8_t sph_idx =
// i * 4 + col_idx; // different order such that adjacent
// spheres are in neighboring threads
@@ -749,23 +747,30 @@ __global__ void kin_fused_backward_kernel3(
int inAddrStart = matAddrBase + linkMap[l] * M * M;
int outAddrStart = matAddrBase + l * M * M; // + (t % M) * M;
float angle = q[batch * njoints + jointMap[l]];
int j_type = jointMapType[l];
if (j_type == FIXED) {
fixed_joint_fn(&fixedTransform[ftAddrStart + col_idx], &JM[0]);
} else if (j_type == Z_ROT) {
zrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type <= Z_PRISM) {
prism_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0], j_type);
} else if (j_type == X_ROT) {
xrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Y_ROT) {
yrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else {
assert(jointMapType[l] > -2 &&
jointMapType[l] < 6); // joint type not supported
}
}
else {
float angle = q[batch * njoints + jointMap[l]];
update_axis_direction(angle, j_type);
if (j_type <= Z_PRISM) {
prism_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0], j_type);
} else if (j_type == X_ROT) {
xrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Y_ROT) {
yrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else if (j_type == Z_ROT) {
zrot_fn(&fixedTransform[ftAddrStart], angle, col_idx, &JM[0]);
} else {
assert(j_type >= FIXED && j_type <= Z_ROT);
}
}
// fetch one row of cumul_mat, multiply with a column, which is in JM
cumul_mat[outAddrStart + elem_idx] =
dot(*(float4 *)&cumul_mat[inAddrStart + ((elem_idx / 4) * M)],
@@ -789,7 +794,8 @@ __global__ void kin_fused_backward_kernel3(
const int spheres_perthread = (nspheres + 15) / 16;
for (int i = 0; i < spheres_perthread; i++) {
const int sph_idx = elem_idx * spheres_perthread + i;
//const int sph_idx = elem_idx * spheres_perthread + i;
const int sph_idx = elem_idx + i * 16;
if (sph_idx >= nspheres) {
break;
}
@@ -808,40 +814,49 @@ __global__ void kin_fused_backward_kernel3(
// read cumul idx:
read_cumul_idx = linkSphereMap[sph_idx];
float spheres_mem[4];
transform_sphere(&cumul_mat[matAddrBase + read_cumul_idx * 16],
&robotSpheres[sphAddrs], &spheres_mem[0]);
for (int j = read_cumul_idx; j > -1; j--) {
// assuming this sphere only depends on links lower than this index
// This could be relaxed by making read_cumul_idx = number of links.
//const int16_t loop_max = read_cumul_idx;
const int16_t loop_max = nlinks - 1;
for (int j = loop_max; j > -1; j--) {
if (linkChainMap[read_cumul_idx * nlinks + j] == 0.0) {
continue;
}
int8_t axis_sign = 1;
int j_type = jointMapType[j];
if(j_type != FIXED)
{
update_joint_type_direction(j_type, axis_sign);
}
if (j_type == Z_ROT) {
float result = 0.0;
z_rot_backward_translation(&cumul_mat[matAddrBase + j * 16],
&spheres_mem[0], loc_grad_sphere, result);
&spheres_mem[0], loc_grad_sphere, result, axis_sign);
psum_grad[jointMap[j]] += (psum_t)result;
} else if (j_type >= X_PRISM && j_type <= Z_PRISM) {
float result = 0.0;
xyz_prism_backward_translation(&cumul_mat[matAddrBase + j * 16],
loc_grad_sphere, result, j_type);
loc_grad_sphere, result, j_type, axis_sign);
psum_grad[jointMap[j]] += (psum_t)result;
} else if (j_type == X_ROT) {
float result = 0.0;
x_rot_backward_translation(&cumul_mat[matAddrBase + j * 16],
&spheres_mem[0], loc_grad_sphere, result);
&spheres_mem[0], loc_grad_sphere, result, axis_sign);
psum_grad[jointMap[j]] += (psum_t)result;
} else if (j_type == Y_ROT) {
float result = 0.0;
y_rot_backward_translation(&cumul_mat[matAddrBase + j * 16],
&spheres_mem[0], loc_grad_sphere, result);
&spheres_mem[0], loc_grad_sphere, result, axis_sign);
psum_grad[jointMap[j]] += (psum_t)result;
}
}
}
// Instead of accumulating the sphere_grad and link_grad separately, we will
// accumulate them together once below.
//
@@ -867,7 +882,7 @@ __global__ void kin_fused_backward_kernel3(
float3 g_position = *(float3 *)&grad_nlinks_pos[batchAddrs * 3 + i * 3];
float4 g_orientation_t =
*(float4 *)&grad_nlinks_quat[batchAddrs * 4 + i * 4];
// TODO: sparisty check here:
// sparisty check here:
if (enable_sparsity_opt) {
if (g_position.x == 0 && g_position.y == 0 && g_position.z == 0 &&
g_orientation_t.y == 0 && g_orientation_t.z == 0 &&
@@ -879,6 +894,7 @@ __global__ void kin_fused_backward_kernel3(
make_float3(g_orientation_t.y, g_orientation_t.z, g_orientation_t.w);
const int16_t l_map = storeLinkMap[i];
float l_pos[3];
const int outAddrStart = matAddrBase + l_map * M * M;
@@ -886,34 +902,47 @@ __global__ void kin_fused_backward_kernel3(
l_pos[1] = cumul_mat[outAddrStart + M + 3];
l_pos[2] = cumul_mat[outAddrStart + M * 2 + 3];
int16_t joints_per_thread = (l_map + 15) / 16;
for (int16_t k = joints_per_thread; k >= 0; k--) {
int16_t j = k * M + elem_idx;
if (j > l_map || j < 0)
const int16_t max_lmap = nlinks - 1;
const int16_t joints_per_thread = (max_lmap + 15) / 16;
//for (int16_t k = joints_per_thread; k >= 0; k--)
for (int16_t k=0; k < joints_per_thread; k++)
{
//int16_t j = elem_idx * joints_per_thread + k;
int16_t j = elem_idx + k * 16;
//int16_t j = k * M + elem_idx;
if (j > max_lmap || j < 0)
continue;
// This can be spread across threads as they are not sequential?
if (linkChainMap[l_map * nlinks + j] == 0.0) {
continue;
}
int16_t j_idx = jointMap[j];
int8_t j_type = jointMapType[j];
int j_type = jointMapType[j];
int8_t axis_sign = 1;
if (j_type != FIXED)
{
update_joint_type_direction(j_type, axis_sign);
}
// get rotation vector:
if (j_type == Z_ROT) {
z_rot_backward(&cumul_mat[matAddrBase + (j)*M * M], &l_pos[0],
g_position, g_orientation, psum_grad[j_idx]);
g_position, g_orientation, psum_grad[j_idx], axis_sign);
} else if (j_type >= X_PRISM & j_type <= Z_PRISM) {
xyz_prism_backward_translation(&cumul_mat[matAddrBase + j * 16],
g_position, psum_grad[j_idx], j_type);
g_position, psum_grad[j_idx], j_type, axis_sign);
} else if (j_type == X_ROT) {
x_rot_backward(&cumul_mat[matAddrBase + (j)*M * M], &l_pos[0],
g_position, g_orientation, psum_grad[j_idx]);
g_position, g_orientation, psum_grad[j_idx], axis_sign);
} else if (j_type == Y_ROT) {
y_rot_backward(&cumul_mat[matAddrBase + (j)*M * M], &l_pos[0],
g_position, g_orientation, psum_grad[j_idx]);
g_position, g_orientation, psum_grad[j_idx], axis_sign);
}
}
}
if (PARALLEL_WRITE) {
// accumulate the partial sums across the 16 threads
@@ -931,7 +960,8 @@ __global__ void kin_fused_backward_kernel3(
#pragma unroll
for (int16_t j = 0; j < joints_per_thread; j++) {
const int16_t j_idx = elem_idx * joints_per_thread + j;
//const int16_t j_idx = elem_idx * joints_per_thread + j;
const int16_t j_idx = elem_idx + j * 16;
if (j_idx >= njoints) {
break;
}

8
src/curobo/curobolib/cpp/lbfgs_step_kernel.cu
View File

@@ -441,7 +441,7 @@ __global__ void reduce_kernel(
rho_buffer[threadIdx.x * batchsize + batch] = rho;
}
}
template <typename scalar_t, typename psum_t>
template <typename scalar_t, typename psum_t, bool rolled_ys>
__global__ void lbfgs_update_buffer_and_step_v1(
scalar_t *step_vec, // b x 175
scalar_t *rho_buffer, // m x b x 1
@@ -452,7 +452,6 @@ __global__ void lbfgs_update_buffer_and_step_v1(
scalar_t *grad_0, // b x 175
const scalar_t *grad_q, // b x 175
const float epsilon, const int batchsize, const int m, const int v_dim,
const bool rolled_ys = false,
const bool stable_mode =
false) // s_buffer and y_buffer are not rolled by default
{
@@ -485,6 +484,7 @@ __global__ void lbfgs_update_buffer_and_step_v1(
scalar_t s =
q[batch * v_dim + threadIdx.x] - x_0[batch * v_dim + threadIdx.x];
reduce_v1(y * s, v_dim, &data[0], &result);
//reduce(y * s, v_dim, &data[0], &result);
scalar_t numerator = result;
// scalar_t rho = 1.0/numerator;
@@ -824,14 +824,14 @@ lbfgs_cuda_fuse(torch::Tensor step_vec, torch::Tensor rho_buffer,
y_buffer.scalar_type(), "lbfgs_cuda_fuse_kernel_v1", [&] {
smemsize = 3 * m * threadsPerBlock * sizeof(scalar_t) +
m * batch_size * sizeof(scalar_t);
lbfgs_update_buffer_and_step_v1<scalar_t, scalar_t>
lbfgs_update_buffer_and_step_v1<scalar_t, scalar_t, false>
<<<blocksPerGrid, threadsPerBlock, smemsize, stream>>>(
step_vec.data_ptr<scalar_t>(),
rho_buffer.data_ptr<scalar_t>(),
y_buffer.data_ptr<scalar_t>(), s_buffer.data_ptr<scalar_t>(),
q.data_ptr<scalar_t>(), x_0.data_ptr<scalar_t>(),
grad_0.data_ptr<scalar_t>(), grad_q.data_ptr<scalar_t>(),
epsilon, batch_size, m, v_dim, false, stable_mode);
epsilon, batch_size, m, v_dim, stable_mode);
});
}

23
src/curobo/curobolib/cpp/pose_distance_kernel.cu
View File

@@ -164,10 +164,14 @@ compute_distance(float *distance_vec, float &distance, float &position_distance,
distance = 0;
if (rotation_distance > vec_convergence[0] * vec_convergence[0]) {
rotation_distance = sqrtf(rotation_distance);
//rotation_distance -= vec_convergence[0];
distance += rotation_weight * rotation_distance;
}
if (position_distance > vec_convergence[1] * vec_convergence[1]) {
position_distance = sqrtf(position_distance);
//position_distance -= vec_convergence[1];
distance += position_weight * position_distance;
}
}
@@ -202,9 +206,12 @@ __device__ __forceinline__ void compute_metric_distance(
distance = 0;
if (rotation_distance > vec_convergence[0] * vec_convergence[0]) {
rotation_distance = sqrtf(rotation_distance);
//rotation_distance -= vec_convergence[0];
distance += rotation_weight * log2f(coshf(r_alpha * rotation_distance));
}
if (position_distance > vec_convergence[1] * vec_convergence[1]) {
//position_distance -= vec_convergence[1];
position_distance = sqrtf(position_distance);
distance += position_weight * log2f(coshf(p_alpha * position_distance));
}
@@ -372,6 +379,14 @@ __global__ void goalset_pose_distance_kernel(
// write out pose distance:
out_distance[batch_idx * horizon + h_idx] = best_distance;
if (write_distance) {
if(position_weight == 0.0)
{
best_position_distance = 0.0;
}
if (rotation_weight == 0.0)
{
best_rotation_distance = 0.0;
}
out_position_distance[batch_idx * horizon + h_idx] = best_position_distance;
out_rotation_distance[batch_idx * horizon + h_idx] = best_rotation_distance;
}
@@ -522,6 +537,14 @@ __global__ void goalset_pose_distance_metric_kernel(
// write out pose distance:
out_distance[batch_idx * horizon + h_idx] = best_distance;
if (write_distance) {
if(position_weight == 0.0)
{
best_position_distance = 0.0;
}
if (rotation_weight == 0.0)
{
best_rotation_distance = 0.0;
}
out_position_distance[batch_idx * horizon + h_idx] = best_position_distance;
out_rotation_distance[batch_idx * horizon + h_idx] = best_rotation_distance;
}

4
src/curobo/curobolib/cpp/tensor_step_cuda.cpp
View File

@@ -79,7 +79,7 @@ std::vector<torch::Tensor> step_position_clique_wrapper(
const torch::Tensor traj_dt, const int batch_size, const int horizon,
const int dof) {
const at::cuda::OptionalCUDAGuard guard(u_position.device());
assert(false); // not supported
CHECK_INPUT(u_position);
CHECK_INPUT(out_position);
CHECK_INPUT(out_velocity);
@@ -155,7 +155,7 @@ std::vector<torch::Tensor> backward_step_position_clique_wrapper(
const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof) {
const at::cuda::OptionalCUDAGuard guard(grad_position.device());
assert(false); // not supported
CHECK_INPUT(out_grad_position);
CHECK_INPUT(grad_position);
CHECK_INPUT(grad_velocity);

138
src/curobo/curobolib/cpp/tensor_step_kernel.cu
View File

@@ -319,14 +319,16 @@ __device__ __forceinline__ void compute_central_difference(scalar_t *out_positio
const float dt = traj_dt[0]; // assume same dt across traj TODO: Implement variable dt
// dt here is actually 1/dt;
const float dt_inv = 1.0 / dt;
const float st_jerk = 0.0; // Note: start jerk can also be passed from global memory
// read start state:
float out_pos=0.0, out_vel=0.0, out_acc=0.0, out_jerk=0.0;
float st_pos=0.0, st_vel=0.0, st_acc = 0.0;
const int b_addrs = b_idx * horizon * dof;
const int b_addrs_action = b_idx * (horizon-4) * dof;
float in_pos[5]; // create a 5 value scalar
const float acc_scale = 1.0;
#pragma unroll 5
for (int i=0; i< 5; i ++){
in_pos[i] = 0.0;
@@ -337,92 +339,108 @@ __device__ __forceinline__ void compute_central_difference(scalar_t *out_positio
st_acc = start_acceleration[b_offset * dof + d_idx];
}
if (h_idx > 3 && h_idx < horizon - 4)
{
in_pos[0] = u_position[b_addrs + (h_idx - 4) * dof + d_idx];
in_pos[1] = u_position[b_addrs + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs + (h_idx - 2 ) * dof + d_idx];
in_pos[3] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
in_pos[4] = u_position[b_addrs + (h_idx) * dof + d_idx];
if (h_idx > 3 && h_idx < horizon - 4){
in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs_action + (h_idx - 2 ) * dof + d_idx];
in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
}
else if (h_idx == 0)
{
in_pos[0] = st_pos - 3 * dt * ( st_vel + (0.5 * st_acc * dt)); // start -1, start, u0, u1
in_pos[1] = st_pos - 2 * dt * ( st_vel + (0.5 * st_acc * dt));
in_pos[2] = st_pos - dt * ( st_vel + (0.5 * st_acc * dt));
in_pos[0] = (3.0f/2) * ( - 1 * st_acc * (dt_inv * dt_inv) - (dt_inv * dt_inv * dt_inv) * st_jerk ) - 3.0f * dt_inv * st_vel + st_pos;
in_pos[1] = -2.0f * st_acc * dt_inv * dt_inv - (4.0f/3) * dt_inv * dt_inv * dt_inv * st_jerk - 2.0 * dt_inv * st_vel + st_pos;
in_pos[2] = -(3.0f/2) * st_acc * dt_inv * dt_inv - (7.0f/6) * dt_inv * dt_inv * dt_inv * st_jerk - dt_inv * st_vel + st_pos;
in_pos[3] = st_pos;
in_pos[4] = u_position[b_addrs + (h_idx) * dof + d_idx];
in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
}
else if (h_idx == 1)
{
in_pos[0] = st_pos - 2 * dt * ( st_vel + (0.5 * st_acc * dt)); // start -1, start, u0, u1
in_pos[1] = st_pos - dt * ( st_vel + (0.5 * st_acc * dt));
in_pos[0] = -2.0f * st_acc * dt_inv * dt_inv - (4.0f/3) * dt_inv * dt_inv * dt_inv * st_jerk - 2.0 * dt_inv * st_vel + st_pos;
in_pos[1] = -(3.0f/2) * st_acc * dt_inv * dt_inv - (7.0f/6) * dt_inv * dt_inv * dt_inv * st_jerk - dt_inv * st_vel + st_pos;
in_pos[2] = st_pos;
in_pos[3] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
in_pos[4] = u_position[b_addrs + (h_idx) * dof + d_idx];
in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
}
else if (h_idx == 2)
{
in_pos[0] = st_pos - dt * ( st_vel + (0.5 * st_acc * dt)); // start -1, start, u0, u1
in_pos[0] = -(3.0f/2) * st_acc * dt_inv * dt_inv - (7.0f/6) * dt_inv * dt_inv * dt_inv * st_jerk - dt_inv * st_vel + st_pos;
in_pos[1] = st_pos;
in_pos[2] = u_position[b_addrs + (h_idx - 2) * dof + d_idx];
in_pos[3] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
in_pos[4] = u_position[b_addrs + (h_idx) * dof + d_idx];
in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
}
else if (h_idx == 3)
{
in_pos[0] = st_pos;
in_pos[1] = u_position[b_addrs + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs + (h_idx - 2 ) * dof + d_idx];
in_pos[3] = u_position[b_addrs + (h_idx - 1 ) * dof + d_idx];
in_pos[4] = u_position[b_addrs + (h_idx) * dof + d_idx];
in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs_action + (h_idx - 2 ) * dof + d_idx];
in_pos[3] = u_position[b_addrs_action + (h_idx - 1 ) * dof + d_idx];
in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
}
else if (h_idx == horizon - 4)
{
in_pos[0] = u_position[b_addrs + (h_idx -4) * dof + d_idx];
in_pos[1] = u_position[b_addrs + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs + (h_idx - 2 ) * dof + d_idx];
in_pos[3] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
in_pos[4] = in_pos[3];//in_pos[3]; //u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx];
in_pos[0] = u_position[b_addrs_action + (h_idx -4) * dof + d_idx];
in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs_action + (h_idx - 2 ) * dof + d_idx];
in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
in_pos[4] = in_pos[3];//in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof + d_idx];
}
else if (h_idx == horizon - 3)
{
in_pos[0] = u_position[b_addrs + (h_idx - 4) * dof + d_idx];
in_pos[1] = u_position[b_addrs + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs + (h_idx - 2) * dof + d_idx];
in_pos[3] = in_pos[2];//u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];
in_pos[4] = in_pos[2];//in_pos[3]; //u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx];
in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
in_pos[3] = in_pos[2];//u_position[b_addrs_action + (h_idx - 1 + 1) * dof + d_idx];
in_pos[4] = in_pos[2];//in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof + d_idx];
}
else if (h_idx == horizon - 2)
{
in_pos[0] = u_position[b_addrs + (h_idx - 4) * dof + d_idx];
in_pos[1] = u_position[b_addrs + (h_idx - 3) * dof + d_idx];
in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
in_pos[2] = in_pos[1];
in_pos[3] = in_pos[1];//u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];
in_pos[4] = in_pos[1];//in_pos[3]; //u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx];
in_pos[3] = in_pos[1];//u_position[b_addrs_action + (h_idx - 1 + 1) * dof + d_idx];
in_pos[4] = in_pos[1];//in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof + d_idx];
}
else if (h_idx == horizon -1)
else if (h_idx == horizon - 1)
{
in_pos[0] = u_position[b_addrs + (h_idx - 4) * dof + d_idx];
in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
in_pos[1] = in_pos[0];
in_pos[2] = in_pos[0];//u_position[b_addrs + (h_idx - 1 ) * dof + d_idx];
in_pos[3] = in_pos[0];//u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];
in_pos[4] = in_pos[0];//in_pos[3]; //u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx];
in_pos[2] = in_pos[0];//u_position[b_addrs_action + (h_idx - 1 ) * dof + d_idx];
in_pos[3] = in_pos[0];//u_position[b_addrs_action + (h_idx - 1 + 1) * dof + d_idx];
in_pos[4] = in_pos[0];//in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof + d_idx];
}
out_pos = in_pos[2];
// out_vel = (0.5 * in_pos[3] - 0.5 * in_pos[1]) * dt;
out_vel = ((0.083333333f) * in_pos[0] - (0.666666667f) * in_pos[1] + (0.666666667f) * in_pos[3] + (-0.083333333f) * in_pos[4]) * dt;
@@ -693,8 +711,9 @@ __global__ void backward_position_clique_loop_central_difference_kernel2(
return;
}
const int b_addrs = b_idx * horizon * dof;
const int b_addrs_action = b_idx * (horizon-4) * dof;
if (h_idx < 2 || h_idx > horizon - 2)
if (h_idx < 2 || h_idx >= horizon - 2)
{
return;
}
@@ -717,7 +736,7 @@ __global__ void backward_position_clique_loop_central_difference_kernel2(
g_jerk[i] = 0.0;
}
int hid = h_idx;
const int hid = h_idx;
g_pos[0] = grad_position[b_addrs + (hid)*dof + d_idx];
g_pos[1] = 0.0;
@@ -745,30 +764,35 @@ __global__ void backward_position_clique_loop_central_difference_kernel2(
(0.5f * g_jerk[0] - g_jerk[1] + g_jerk[3] - 0.5f * g_jerk[4]) * dt * dt * dt);
}
else if (hid == horizon -3)
else if (hid == horizon - 3)
{
//The below can cause oscilatory gradient steps.
/*
#pragma unroll
for (int i=0; i< 4; i++)
for (int i=0; i< 5; i++)
{
g_vel[i] = grad_velocity[b_addrs + ((hid - 2) + i)*dof + d_idx];
g_acc[i] = grad_acceleration[b_addrs + ((hid -2) + i)*dof + d_idx];
g_jerk[i] = grad_jerk[b_addrs + ((hid -2) + i)*dof + d_idx];
}
*/
g_pos[1] = grad_position[b_addrs + (hid + 1)*dof + d_idx];
g_pos[2] = grad_position[b_addrs + (hid + 2)*dof + d_idx];
out_grad = (g_pos[0] + g_pos[1] + g_pos[2] +
out_grad = (g_pos[0] + g_pos[1] + g_pos[2]);
/* +
//((0.5) * g_vel[1] + (0.5) * g_vel[2]) * dt +
((-0.083333333f) * g_vel[0] + (0.583333333f) * g_vel[1] + (0.583333333f) * g_vel[2] + (-0.083333333f) * g_vel[3]) * dt +
((-0.083333333f) * g_acc[0] + (1.25f) * g_acc[1] + (-1.25f) * g_acc[2] + (0.083333333f) * g_acc[3]) * dt * dt +
//( g_acc[1] - g_acc[2]) * dt * dt +
(0.5f * g_jerk[0] - 0.5f * g_jerk[1] -0.5f * g_jerk[2] + 0.5f * g_jerk[3]) * dt * dt * dt);
*/
}
// write out:
out_grad_position[b_addrs + (h_idx-2)*dof + d_idx] = out_grad;
out_grad_position[b_addrs_action + (h_idx-2)*dof + d_idx] = out_grad;
}
// for MPPI:
@@ -1389,6 +1413,8 @@ std::vector<torch::Tensor> step_position_clique2(
}
else if (mode == CENTRAL_DIFF)
{
assert(u_position.sizes()[1] == horizon - 4);
AT_DISPATCH_FLOATING_TYPES(
out_position.scalar_type(), "step_position_clique", ([&] {
position_clique_loop_kernel2<scalar_t, CENTRAL_DIFF>
@@ -1436,7 +1462,7 @@ std::vector<torch::Tensor> step_position_clique2_idx(
if (mode == BWD_DIFF)
{
assert(false);
AT_DISPATCH_FLOATING_TYPES(
out_position.scalar_type(), "step_position_clique", ([&] {
position_clique_loop_idx_kernel2<scalar_t, BWD_DIFF>
@@ -1455,6 +1481,8 @@ std::vector<torch::Tensor> step_position_clique2_idx(
else if (mode == CENTRAL_DIFF)
{
assert(u_position.sizes()[1] == horizon - 4);
AT_DISPATCH_FLOATING_TYPES(
out_position.scalar_type(), "step_position_clique", ([&] {
position_clique_loop_idx_kernel2<scalar_t, CENTRAL_DIFF>
@@ -1538,7 +1566,7 @@ std::vector<torch::Tensor> backward_step_position_clique2(
if (mode == BWD_DIFF)
{
assert(false); // not supported anymore
AT_DISPATCH_FLOATING_TYPES(
out_grad_position.scalar_type(), "backward_step_position_clique", ([&] {
backward_position_clique_loop_backward_difference_kernel2<scalar_t>
@@ -1554,7 +1582,7 @@ std::vector<torch::Tensor> backward_step_position_clique2(
}
else if (mode == CENTRAL_DIFF)
{
assert(out_grad_position.sizes()[1] == horizon - 4);
AT_DISPATCH_FLOATING_TYPES(
out_grad_position.scalar_type(), "backward_step_position_clique", ([&] {
backward_position_clique_loop_central_difference_kernel2<scalar_t>

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

@@ -848,13 +848,13 @@ class WorldPrimitiveCollision(WorldCollision):
trange = [h - l for l, h in zip(low, high)]
x = torch.linspace(
-bounds[0], bounds[0], int(trange[0] // voxel_size), device=self.tensor_args.device
-bounds[0], bounds[0], int(trange[0] // voxel_size) + 1, device=self.tensor_args.device
)
y = torch.linspace(
-bounds[1], bounds[1], int(trange[1] // voxel_size), device=self.tensor_args.device
-bounds[1], bounds[1], int(trange[1] // voxel_size) + 1, device=self.tensor_args.device
)
z = torch.linspace(
-bounds[2], bounds[2], int(trange[2] // voxel_size), device=self.tensor_args.device
-bounds[2], bounds[2], int(trange[2] // voxel_size) + 1, device=self.tensor_args.device
)
w, l, h = x.shape[0], y.shape[0], z.shape[0]
xyz = (
@@ -893,8 +893,11 @@ class WorldPrimitiveCollision(WorldCollision):
voxel_size: float = 0.02,
) -> Mesh:
voxels = self.get_voxels_in_bounding_box(cuboid, voxel_size)
# voxels = voxels.cpu().numpy()
# cuboids = [Cuboid(name="c_"+str(x), pose=[voxels[x,0],voxels[x,1],voxels[x,2], 1,0,0,0], dims=[voxel_size, voxel_size, voxel_size]) for x in range(voxels.shape[0])]
# mesh = WorldConfig(cuboid=cuboids).get_mesh_world(True).mesh[0]
mesh = Mesh.from_pointcloud(
voxels[:, :3].detach().cpu().numpy(),
pitch=voxel_size,
pitch=voxel_size * 1.1,
)
return mesh

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

@@ -76,8 +76,8 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_mapper = Mapper(
voxel_sizes=voxel_sizes,
integrator_types=integrator_types,
cuda_device_id=self.tensor_args.device.index,
free_on_destruction=False,
cuda_device_id=self.tensor_args.device.index,
)
self._blox_voxel_sizes = voxel_sizes
# load map from file if it exists:
@@ -412,6 +412,8 @@ class WorldBloxCollision(WorldMeshCollision):
index = self._blox_names.index(layer_name)
pose_mat = camera_observation.pose.get_matrix().view(4, 4)
if camera_observation.rgb_image is not None:
if camera_observation.rgb_image.shape[-1] != 4:
log_error("nvblox color should be of shape HxWx4")
with profiler.record_function("world_blox/add_color_frame"):
self._blox_mapper.add_color_frame(
camera_observation.rgb_image,
@@ -473,6 +475,11 @@ class WorldBloxCollision(WorldMeshCollision):
)
return mesh
def save_layer(self, layer_name: str, file_name: str) -> bool:
index = self._blox_names.index(layer_name)
status = self._blox_mapper.save_map(file_name, index)
return status
def decay_layer(self, layer_name: str):
index = self._blox_names.index(layer_name)
self._blox_mapper.decay_occupancy(mapper_id=index)

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

@@ -391,9 +391,9 @@ class Mesh(Obstacle):
vertex_normals=self.vertex_normals,
face_colors=self.face_colors,
)
# if self.scale is not None:
# m.vertices = np.ravel(self.scale) * m.vertices
if self.scale is not None:
m.vertices = np.ravel(self.scale) * m.vertices
self.scale = None
return m
def update_material(self):

38
src/curobo/graph/graph_base.py
View File

@@ -637,29 +637,13 @@ class GraphPlanBase(GraphConfig):
@torch.no_grad()
def find_paths(self, x_init, x_goal, interpolation_steps: Optional[int] = None) -> GraphResult:
start_time = time.time()
path = None
try:
path = self._find_paths(x_init, x_goal)
path.success = torch.as_tensor(
path.success, device=self.tensor_args.device, dtype=torch.bool
)
path.solve_time = time.time() - start_time
if self.interpolation_type is not None and torch.count_nonzero(path.success):
(
path.interpolated_plan,
path.path_buffer_last_tstep,
path.optimized_dt,
) = self.get_interpolated_trajectory(path.plan, interpolation_steps)
# path.js_interpolated_plan = self.rollout_fn.get_full_dof_from_solution(
# path.interpolated_plan
# )
if self.compute_metrics:
# compute metrics on interpolated plan:
path.metrics = self.get_metrics(path.interpolated_plan)
path.success = torch.logical_and(
path.success, torch.all(path.metrics.feasible, 1)
)
except ValueError as e:
log_info(e)
@@ -667,12 +651,29 @@ class GraphPlanBase(GraphConfig):
torch.cuda.empty_cache()
success = torch.zeros(x_init.shape[0], device=self.tensor_args.device, dtype=torch.bool)
path = GraphResult(success, x_init, x_goal)
return path
except RuntimeError as e:
log_warn(e)
self.reset_buffer()
torch.cuda.empty_cache()
success = torch.zeros(x_init.shape[0], device=self.tensor_args.device, dtype=torch.long)
path = GraphResult(success, x_init, x_goal)
return path
if self.interpolation_type is not None and (torch.count_nonzero(path.success) > 0):
(
path.interpolated_plan,
path.path_buffer_last_tstep,
path.optimized_dt,
) = self.get_interpolated_trajectory(path.plan, interpolation_steps)
# path.js_interpolated_plan = self.rollout_fn.get_full_dof_from_solution(
# path.interpolated_plan
# )
if self.compute_metrics:
# compute metrics on interpolated plan:
path.metrics = self.get_metrics(path.interpolated_plan)
path.success = torch.logical_and(path.success, torch.all(path.metrics.feasible, 1))
return path
@abstractmethod
@@ -907,10 +908,11 @@ class GraphPlanBase(GraphConfig):
dof = self.dof
i = self.i
if x_set is not None:
if x_set.shape[0] == 0:
log_warn("no valid configuration found")
return
if connect_mode == "radius":
raise NotImplementedError
scale_radius = self.neighbour_radius * (np.log(i) / i) ** (1 / dof)

19
src/curobo/opt/newton/lbfgs.py
View File

@@ -24,9 +24,16 @@ from curobo.util.logger import log_warn
# kernel for l-bfgs:
@torch.jit.script
# @torch.jit.script
def compute_step_direction(
alpha_buffer, rho_buffer, y_buffer, s_buffer, grad_q, m: int, epsilon, stable_mode: bool = True
alpha_buffer,
rho_buffer,
y_buffer,
s_buffer,
grad_q,
m: int,
epsilon: float,
stable_mode: bool = True,
):
# m = 15 (int)
# y_buffer, s_buffer: m x b x 175
@@ -70,12 +77,12 @@ class LBFGSOpt(NewtonOptBase, LBFGSOptConfig):
if config is not None:
LBFGSOptConfig.__init__(self, **vars(config))
NewtonOptBase.__init__(self)
if self.d_opt >= 1024 or self.history >= 512:
log_warn("LBFGS: Not using LBFGS Cuda Kernel as d_opt>1024 or history>=512")
if self.d_opt >= 1024 or self.history > 15:
log_warn("LBFGS: Not using LBFGS Cuda Kernel as d_opt>1024 or history>15")
self.use_cuda_kernel = False
if self.history > self.d_opt:
if self.history >= self.d_opt:
log_warn("LBFGS: history >= d_opt, reducing history to d_opt-1")
self.history = self.d_opt
self.history = self.d_opt - 1
@profiler.record_function("lbfgs/reset")
def reset(self):

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

@@ -72,7 +72,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
):
if config is not None:
NewtonOptConfig.__init__(self, **vars(config))
self.d_opt = self.horizon * self.d_action
self.d_opt = self.action_horizon * self.d_action
self.line_scale = self._create_box_line_search(self.line_search_scale)
Optimizer.__init__(self)
self.i = -1
@@ -84,8 +84,8 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self.reset()
# reshape action lows and highs:
self.action_lows = self.action_lows.repeat(self.horizon)
self.action_highs = self.action_highs.repeat(self.horizon)
self.action_lows = self.action_lows.repeat(self.action_horizon)
self.action_highs = self.action_highs.repeat(self.action_horizon)
self.action_range = self.action_highs - self.action_lows
self.action_step_max = self.step_scale * torch.abs(self.action_range)
self.c_1 = 1e-5
@@ -99,10 +99,13 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self.use_cuda_line_search_kernel = False
if self.use_temporal_smooth:
self._temporal_mat = build_fd_matrix(
self.horizon, order=2, device=self.tensor_args.device, dtype=self.tensor_args.dtype
self.action_horizon,
order=2,
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
).unsqueeze(0)
eye_mat = torch.eye(
self.horizon, device=self.tensor_args.device, dtype=self.tensor_args.dtype
self.action_horizon, device=self.tensor_args.device, dtype=self.tensor_args.dtype
).unsqueeze(0)
self._temporal_mat += eye_mat
@@ -130,9 +133,9 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self._shift(shift_steps)
# reshape q:
if self.store_debug:
self.debug.append(q.view(-1, self.horizon, self.d_action).clone())
self.debug.append(q.view(-1, self.action_horizon, self.d_action).clone())
with profiler.record_function("newton_base/init_opt"):
q = q.view(self.n_envs, self.horizon * self.d_action)
q = q.view(self.n_envs, self.action_horizon * self.d_action)
grad_q = q.detach() * 0.0
# run opt graph
if not self.cu_opt_init:
@@ -147,7 +150,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
if check_convergence(self.best_iteration, self.current_iteration, self.last_best):
break
best_q = best_q.view(self.n_envs, self.horizon, self.d_action)
best_q = best_q.view(self.n_envs, self.action_horizon, self.d_action)
return best_q
def reset(self):
@@ -166,8 +169,11 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self.i += 1
cost_n, q, grad_q = self._opt_step(q.detach(), grad_q.detach())
if self.store_debug:
self.debug.append(self.best_q.view(-1, self.horizon, self.d_action).clone())
self.debug.append(self.best_q.view(-1, self.action_horizon, self.d_action).clone())
self.debug_cost.append(self.best_cost.detach().view(-1, 1).clone())
# self.debug.append(q.view(-1, self.action_horizon, self.d_action).clone())
# self.debug_cost.append(cost_n.detach().view(-1, 1).clone())
# print(grad_q)
return self.best_q.detach(), self.best_cost.detach(), q.detach(), grad_q.detach()
@@ -186,9 +192,9 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
def scale_step_direction(self, dx):
if self.use_temporal_smooth:
dx_v = dx.view(-1, self.horizon, self.d_action)
dx_v = dx.view(-1, self.action_horizon, self.d_action)
dx_new = self._temporal_mat @ dx_v # 1,h,h x b, h, dof -> b, h, dof
dx = dx_new.view(-1, self.horizon * self.d_action)
dx = dx_new.view(-1, self.action_horizon * self.d_action)
dx_scaled = scale_action(dx, self.action_step_max)
return dx_scaled
@@ -216,11 +222,11 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
def _compute_cost_gradient(self, x):
x_n = x.detach().requires_grad_(True)
x_in = x_n.view(
self.n_envs * self.num_particles, self.rollout_fn.horizon, self.rollout_fn.d_action
self.n_envs * self.num_particles, self.action_horizon, self.rollout_fn.d_action
)
trajectories = self.rollout_fn(x_in) # x_n = (batch*line_search_scale) x horizon x d_action
cost = torch.sum(
trajectories.costs.view(self.n_envs, self.num_particles, self.rollout_fn.horizon),
trajectories.costs.view(self.n_envs, self.num_particles, self.horizon),
dim=-1,
keepdim=True,
)

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

@@ -43,6 +43,7 @@ class OptimizerConfig:
n_envs: int
sync_cuda_time: bool
use_coo_sparse: bool
action_horizon: int
def __post_init__(self):
object.__setattr__(self, "action_highs", self.tensor_args.to_device(self.action_highs))
@@ -68,6 +69,8 @@ class OptimizerConfig:
child_dict["rollout_fn"] = rollout_fn
child_dict["tensor_args"] = tensor_args
child_dict["horizon"] = rollout_fn.horizon
child_dict["action_horizon"] = rollout_fn.action_horizon
if "num_particles" not in child_dict:
child_dict["num_particles"] = 1
return child_dict

29
src/curobo/opt/particle/parallel_mppi.py
View File

@@ -89,7 +89,7 @@ class ParallelMPPIConfig(ParticleOptConfig):
child_dict["squash_fn"] = SquashType[child_dict["squash_fn"]]
child_dict["cov_type"] = CovType[child_dict["cov_type"]]
child_dict["sample_params"]["d_action"] = rollout_fn.d_action
child_dict["sample_params"]["horizon"] = child_dict["horizon"]
child_dict["sample_params"]["horizon"] = rollout_fn.action_horizon
child_dict["sample_params"]["tensor_args"] = tensor_args
child_dict["sample_params"] = SampleConfig(**child_dict["sample_params"])
@@ -112,7 +112,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
# initialize covariance types:
if self.cov_type == CovType.FULL_HA:
self.I = torch.eye(
self.horizon * self.d_action,
self.action_horizon * self.d_action,
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
@@ -124,7 +124,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
self.Z_seq = torch.zeros(
1,
self.horizon,
self.action_horizon,
self.d_action,
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
@@ -145,7 +145,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
self.mean_lib = HaltonSampleLib(
SampleConfig(
self.horizon,
self.action_horizon,
self.d_action,
tensor_args=self.tensor_args,
**{"fixed_samples": False, "seed": 2567, "filter_coeffs": None}
@@ -330,7 +330,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
(cat_list),
dim=-3,
)
act_seq = act_seq.reshape(self.total_num_particles, self.horizon, self.d_action)
act_seq = act_seq.reshape(self.total_num_particles, self.action_horizon, self.d_action)
act_seq = scale_ctrl(act_seq, self.action_lows, self.action_highs, squash_fn=self.squash_fn)
# if not copy_tensor(act_seq, self.act_seq):
@@ -399,7 +399,8 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
act_seq = self.mean_action # .clone() # [self.mean_idx]#.clone()
elif mode == SampleMode.SAMPLE:
delta = self.generate_noise(
shape=torch.Size((1, self.horizon)), base_seed=self.seed + 123 * self.num_steps
shape=torch.Size((1, self.action_horizon)),
base_seed=self.seed + 123 * self.num_steps,
)
act_seq = self.mean_action + torch.matmul(delta, self.full_scale_tril)
elif mode == SampleMode.BEST:
@@ -426,9 +427,11 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
Tensor: dimension is (d_action, d_action)
"""
if self.cov_type == CovType.SIGMA_I:
return self.scale_tril.unsqueeze(-2).unsqueeze(-2).expand(-1, -1, self.horizon, -1)
return (
self.scale_tril.unsqueeze(-2).unsqueeze(-2).expand(-1, -1, self.action_horizon, -1)
)
elif self.cov_type == CovType.DIAG_A:
return self.scale_tril.unsqueeze(-2).expand(-1, -1, self.horizon, -1) # .cl
return self.scale_tril.unsqueeze(-2).expand(-1, -1, self.action_horizon, -1) # .cl
elif self.cov_type == CovType.FULL_A:
return self.scale_tril
elif self.cov_type == CovType.FULL_HA:
@@ -486,10 +489,10 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
def full_scale_tril(self):
if self.cov_type == CovType.SIGMA_I:
return (
self.scale_tril.unsqueeze(-2).unsqueeze(-2).expand(-1, -1, self.horizon, -1)
self.scale_tril.unsqueeze(-2).unsqueeze(-2).expand(-1, -1, self.action_horizon, -1)
) # .cl
elif self.cov_type == CovType.DIAG_A:
return self.scale_tril.unsqueeze(-2).expand(-1, -1, self.horizon, -1) # .cl
return self.scale_tril.unsqueeze(-2).expand(-1, -1, self.action_horizon, -1) # .cl
elif self.cov_type == CovType.FULL_A:
return self.scale_tril
elif self.cov_type == CovType.FULL_HA:
@@ -504,7 +507,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
self.sample_lib = SampleLib(self.sample_params)
self.mean_lib = HaltonSampleLib(
SampleConfig(
self.horizon,
self.action_horizon,
self.d_action,
tensor_args=self.tensor_args,
**{"fixed_samples": False, "seed": 2567, "filter_coeffs": None}
@@ -530,7 +533,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
n_iters,
self.n_envs,
self.sampled_particles_per_env,
self.horizon,
self.action_horizon,
self.d_action,
)
.clone()
@@ -541,7 +544,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
base_seed=self.seed,
)
s_set = s_set.view(
n_iters, 1, self.sampled_particles_per_env, self.horizon, self.d_action
n_iters, 1, self.sampled_particles_per_env, self.action_horizon, self.d_action
)
s_set = s_set.repeat(1, self.n_envs, 1, 1, 1).clone()
s_set[:, :, -1, :, :] = 0.0

4
src/curobo/opt/particle/particle_opt_base.py
View File

@@ -259,7 +259,7 @@ class ParticleOptBase(Optimizer, ParticleOptConfig):
# generate random simulated trajectories
trajectory = self.generate_rollouts()
trajectory.actions = trajectory.actions.view(
self.n_envs, self.particles_per_env, self.horizon, self.d_action
self.n_envs, self.particles_per_env, self.action_horizon, self.d_action
)
trajectory.costs = trajectory.costs.view(
self.n_envs, self.particles_per_env, self.horizon
@@ -295,7 +295,7 @@ class ParticleOptBase(Optimizer, ParticleOptConfig):
if self.null_per_env > 0:
self.null_act_seqs = torch.zeros(
self.null_per_env,
self.horizon,
self.action_horizon,
self.d_action,
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,

16
src/curobo/opt/particle/particle_opt_utils.py
View File

@@ -64,7 +64,9 @@ def get_stomp_cov(
Coefficients from here: https://en.wikipedia.org/wiki/Finite_difference_coefficient
More info here: https://github.com/ros-industrial/stomp_ros/blob/7fe40fbe6ad446459d8d4889916c64e276dbf882/stomp_core/src/utils.cpp#L36
"""
cov, scale_tril, scaled_M = get_stomp_cov_jit(horizon, d_action, cov_mode)
cov, scale_tril, scaled_M = get_stomp_cov_jit(
horizon, d_action, cov_mode, device=tensor_args.device
)
cov = tensor_args.to_device(cov)
scale_tril = tensor_args.to_device(scale_tril)
if RETURN_M:
@@ -77,13 +79,16 @@ def get_stomp_cov_jit(
horizon: int,
d_action: int,
cov_mode: str = "acc",
device: torch.device = torch.device("cuda:0"),
):
# This function can lead to nans. There are checks to raise an error when nan occurs.
vel_fd_array = [0.0, 0.0, 1.0, -2.0, 1.0, 0.0, 0.0]
fd_array = vel_fd_array
A = torch.zeros(
(d_action * horizon, d_action * horizon),
dtype=torch.float64,
dtype=torch.float32,
device=device,
)
if cov_mode == "vel":
@@ -117,14 +122,17 @@ def get_stomp_cov_jit(
A[k * horizon + i, k * horizon + index] = fd_array[j + 3]
R = torch.matmul(A.transpose(-2, -1), A)
M = torch.inverse(R)
scaled_M = (1 / horizon) * M / (torch.max(torch.abs(M), dim=1)[0].unsqueeze(0))
cov = M / torch.max(torch.abs(M))
# also compute the cholesky decomposition:
# scale_tril = torch.zeros((d_action * horizon, d_action * horizon), **tensor_args)
scale_tril = torch.linalg.cholesky(cov)
if (cov == cov.T).all() and (torch.linalg.eigvals(cov).real >= 0).all():
scale_tril = torch.linalg.cholesky(cov)
else:
scale_tril = cov
"""
k = 0
act_cov_matrix = cov[k * horizon:k * horizon + horizon, k * horizon:k * horizon + horizon]

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

@@ -324,9 +324,9 @@ class ArmBase(RolloutBase, ArmBaseConfig):
# set start state:
start_state = torch.randn((1, self.dynamics_model.d_state), **vars(self.tensor_args))
self._start_state = JointState(
position=start_state[:, : self.dynamics_model.d_action],
velocity=start_state[:, : self.dynamics_model.d_action],
acceleration=start_state[:, : self.dynamics_model.d_action],
position=start_state[:, : self.dynamics_model.d_dof],
velocity=start_state[:, : self.dynamics_model.d_dof],
acceleration=start_state[:, : self.dynamics_model.d_dof],
)
self.update_cost_dt(self.dynamics_model.dt_traj_params.base_dt)
return RolloutBase._init_after_config_load(self)
@@ -680,6 +680,10 @@ class ArmBase(RolloutBase, ArmBaseConfig):
def horizon(self):
return self.dynamics_model.horizon
@property
def action_horizon(self):
return self.dynamics_model.action_horizon
def get_init_action_seq(self) -> torch.Tensor:
act_seq = self.dynamics_model.init_action_mean.unsqueeze(0).repeat(self.batch_size, 1, 1)
return act_seq

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

@@ -29,6 +29,7 @@ from curobo.types.base import TensorDeviceType
from curobo.types.robot import RobotConfig
from curobo.types.tensor import T_BValue_float
from curobo.util.helpers import list_idx_if_not_none
from curobo.util.logger import log_info
from curobo.util.tensor_util import cat_max, cat_sum
# Local Folder
@@ -79,6 +80,7 @@ class ArmReacherCostConfig(ArmCostConfig):
zero_acc_cfg: Optional[CostConfig] = None
zero_vel_cfg: Optional[CostConfig] = None
zero_jerk_cfg: Optional[CostConfig] = None
link_pose_cfg: Optional[PoseCostConfig] = None
@staticmethod
def _get_base_keys():
@@ -91,6 +93,7 @@ class ArmReacherCostConfig(ArmCostConfig):
"zero_acc_cfg": CostConfig,
"zero_vel_cfg": CostConfig,
"zero_jerk_cfg": CostConfig,
"link_pose_cfg": PoseCostConfig,
}
new_k.update(base_k)
return new_k
@@ -166,10 +169,17 @@ class ArmReacher(ArmBase, ArmReacherConfig):
self.dist_cost = DistCost(self.cost_cfg.cspace_cfg)
if self.cost_cfg.pose_cfg is not None:
self.goal_cost = PoseCost(self.cost_cfg.pose_cfg)
self._link_pose_costs = {}
if self.cost_cfg.link_pose_cfg is None:
log_info(
"Deprecated: Add link_pose_cfg to your rollout config. Using pose_cfg instead."
)
self.cost_cfg.link_pose_cfg = self.cost_cfg.pose_cfg
self._link_pose_costs = {}
if self.cost_cfg.link_pose_cfg is not None:
for i in self.kinematics.link_names:
if i != self.kinematics.ee_link:
self._link_pose_costs[i] = PoseCost(self.cost_cfg.pose_cfg)
self._link_pose_costs[i] = PoseCost(self.cost_cfg.link_pose_cfg)
if self.cost_cfg.straight_line_cfg is not None:
self.straight_line_cost = StraightLineCost(self.cost_cfg.straight_line_cfg)
if self.cost_cfg.zero_vel_cfg is not None:
@@ -192,12 +202,20 @@ class ArmReacher(ArmBase, ArmReacherConfig):
self.z_tensor = torch.tensor(
0, device=self.tensor_args.device, dtype=self.tensor_args.dtype
)
self._link_pose_convergence = {}
if self.convergence_cfg.pose_cfg is not None:
self.pose_convergence = PoseCost(self.convergence_cfg.pose_cfg)
self._link_pose_convergence = {}
if self.convergence_cfg.link_pose_cfg is None:
log_warn(
"Deprecated: Add link_pose_cfg to your rollout config. Using pose_cfg instead."
)
self.convergence_cfg.link_pose_cfg = self.convergence_cfg.pose_cfg
if self.convergence_cfg.link_pose_cfg is not None:
for i in self.kinematics.link_names:
if i != self.kinematics.ee_link:
self._link_pose_convergence[i] = PoseCost(self.convergence_cfg.pose_cfg)
self._link_pose_convergence[i] = PoseCost(self.convergence_cfg.link_pose_cfg)
if self.convergence_cfg.cspace_cfg is not None:
self.cspace_convergence = DistCost(self.convergence_cfg.cspace_cfg)
@@ -307,6 +325,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
# print(z_vel.shape)
cost_list.append(z_vel)
cost = cat_sum(cost_list)
# print(cost[:].T)
return cost
def convergence_fn(

2
src/curobo/rollout/cost/pose_cost.py
View File

@@ -691,8 +691,6 @@ class PoseCost(CostBase, PoseCostConfig):
)
cost = distance
# print(cost.shape)
return cost
def forward_pose(

44
src/curobo/rollout/dynamics_model/integration_utils.py
View File

@@ -340,9 +340,9 @@ class CliqueTensorStepKernel(torch.autograd.Function):
grad_out_a,
grad_out_j,
traj_dt,
out_grad_position.shape[0],
out_grad_position.shape[1],
out_grad_position.shape[2],
grad_out_p.shape[0],
grad_out_p.shape[1],
grad_out_p.shape[2],
)
return (
u_grad,
@@ -412,9 +412,9 @@ class CliqueTensorStepIdxKernel(torch.autograd.Function):
grad_out_a,
grad_out_j,
traj_dt,
out_grad_position.shape[0],
out_grad_position.shape[1],
out_grad_position.shape[2],
grad_out_p.shape[0],
grad_out_p.shape[1],
grad_out_p.shape[2],
)
return (
u_grad,
@@ -483,9 +483,9 @@ class CliqueTensorStepCentralDifferenceKernel(torch.autograd.Function):
grad_out_a.contiguous(),
grad_out_j.contiguous(),
traj_dt,
out_grad_position.shape[0],
out_grad_position.shape[1],
out_grad_position.shape[2],
grad_out_p.shape[0],
grad_out_p.shape[1],
grad_out_p.shape[2],
0,
)
return (
@@ -557,9 +557,9 @@ class CliqueTensorStepIdxCentralDifferenceKernel(torch.autograd.Function):
grad_out_a.contiguous(),
grad_out_j.contiguous(),
traj_dt,
out_grad_position.shape[0],
out_grad_position.shape[1],
out_grad_position.shape[2],
grad_out_p.shape[0],
grad_out_p.shape[1],
grad_out_p.shape[2],
0,
)
return (
@@ -626,9 +626,9 @@ class CliqueTensorStepCoalesceKernel(torch.autograd.Function):
grad_out_a.transpose(-1, -2).contiguous(),
grad_out_j.transpose(-1, -2).contiguous(),
traj_dt,
out_grad_position.shape[0],
out_grad_position.shape[1],
out_grad_position.shape[2],
grad_out_p.shape[0],
grad_out_p.shape[1],
grad_out_p.shape[2],
)
return (
u_grad.transpose(-1, -2).contiguous(),
@@ -890,16 +890,16 @@ def interpolate_kernel(h, int_steps, tensor_args: TensorDeviceType):
return mat
def action_interpolate_kernel(h, int_steps, tensor_args: TensorDeviceType):
def action_interpolate_kernel(h, int_steps, tensor_args: TensorDeviceType, offset: int = 4):
mat = torch.zeros(
((h - 1) * (int_steps), h), device=tensor_args.device, dtype=tensor_args.dtype
)
delta = torch.arange(0, int_steps - 2, device=tensor_args.device, dtype=tensor_args.dtype) / (
int_steps - 1.0 - 2
)
delta = torch.arange(
0, int_steps - offset + 1, device=tensor_args.device, dtype=tensor_args.dtype
) / (int_steps - offset)
for i in range(h - 1):
mat[i * int_steps : i * (int_steps) + int_steps - 1 - 2, i] = delta.flip(0)[1:]
mat[i * int_steps : i * (int_steps) + int_steps - 1 - 2, i + 1] = delta[1:]
mat[-3:, 1] = 1.0
mat[i * int_steps : i * (int_steps) + int_steps - offset, i] = delta.flip(0)[1:]
mat[i * int_steps : i * (int_steps) + int_steps - offset, i + 1] = delta[1:]
mat[-offset:, 1] = 1.0
return mat

62
src/curobo/rollout/dynamics_model/kinematic_model.py
View File

@@ -176,6 +176,7 @@ class KinematicModel(KinematicModelConfig):
self._use_clique_kernel = True
self.d_state = 4 * self.n_dofs # + 1
self.d_action = self.n_dofs
self.d_dof = self.n_dofs
# Variables for enforcing joint limits
self.joint_names = self.robot_model.joint_names
@@ -190,7 +191,7 @@ class KinematicModel(KinematicModelConfig):
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
),
dof=int(self.d_state / 3),
dof=int(self.d_dof),
)
self.Z = torch.tensor([0.0], device=self.tensor_args.device, dtype=self.tensor_args.dtype)
@@ -204,10 +205,18 @@ class KinematicModel(KinematicModelConfig):
# self._cmd_step_fn = TensorStepAcceleration(self.tensor_args, self.traj_dt)
self._rollout_step_fn = TensorStepAccelerationKernel(
self.tensor_args, self.traj_dt, self.n_dofs
self.tensor_args,
self.traj_dt,
self.n_dofs,
self.batch_size,
self.horizon,
)
self._cmd_step_fn = TensorStepAccelerationKernel(
self.tensor_args, self.traj_dt, self.n_dofs
self.tensor_args,
self.traj_dt,
self.n_dofs,
self.batch_size,
self.horizon,
)
elif self.control_space == StateType.VELOCITY:
raise NotImplementedError()
@@ -215,8 +224,12 @@ class KinematicModel(KinematicModelConfig):
raise NotImplementedError()
elif self.control_space == StateType.POSITION:
if self.teleport_mode:
self._rollout_step_fn = TensorStepPositionTeleport(self.tensor_args)
self._cmd_step_fn = TensorStepPositionTeleport(self.tensor_args)
self._rollout_step_fn = TensorStepPositionTeleport(
self.tensor_args, self.batch_size, self.horizon
)
self._cmd_step_fn = TensorStepPositionTeleport(
self.tensor_args, self.batch_size, self.horizon
)
else:
if self._use_clique:
if self._use_clique_kernel:
@@ -237,6 +250,8 @@ class KinematicModel(KinematicModelConfig):
filter_velocity=False,
filter_acceleration=False,
filter_jerk=False,
batch_size=self.batch_size,
horizon=self.horizon,
)
self._cmd_step_fn = TensorStepPositionCliqueKernel(
self.tensor_args,
@@ -246,17 +261,36 @@ class KinematicModel(KinematicModelConfig):
filter_velocity=False,
filter_acceleration=self.filter_robot_command,
filter_jerk=self.filter_robot_command,
batch_size=self.batch_size,
horizon=self.horizon,
)
else:
self._rollout_step_fn = TensorStepPositionClique(
self.tensor_args, self.traj_dt
self.tensor_args,
self.traj_dt,
batch_size=self.batch_size,
horizon=self.horizon,
)
self._cmd_step_fn = TensorStepPositionClique(
self.tensor_args,
self.traj_dt,
batch_size=self.batch_size,
horizon=self.horizon,
)
self._cmd_step_fn = TensorStepPositionClique(self.tensor_args, self.traj_dt)
else:
self._rollout_step_fn = TensorStepPosition(self.tensor_args, self.traj_dt)
self._cmd_step_fn = TensorStepPosition(self.tensor_args, self.traj_dt)
self._rollout_step_fn = TensorStepPosition(
self.tensor_args,
self.traj_dt,
batch_size=self.batch_size,
horizon=self.horizon,
)
self._cmd_step_fn = TensorStepPosition(
self.tensor_args,
self.traj_dt,
batch_size=self.batch_size,
horizon=self.horizon,
)
self.update_batch_size(self.batch_size)
self.state_filter = JointStateFilter(self.state_filter_cfg)
@@ -542,10 +576,10 @@ class KinematicModel(KinematicModelConfig):
# output should be d_action * horizon
if self.control_space == StateType.POSITION:
# use joint limits:
return self.retract_config.unsqueeze(0).repeat(self.horizon, 1)
return self.retract_config.unsqueeze(0).repeat(self.action_horizon, 1)
if self.control_space == StateType.VELOCITY or self.control_space == StateType.ACCELERATION:
# use joint limits:
return self.retract_config.unsqueeze(0).repeat(self.horizon, 1) * 0.0
return self.retract_config.unsqueeze(0).repeat(self.action_horizon, 1) * 0.0
@property
def retract_config(self):
@@ -567,6 +601,10 @@ class KinematicModel(KinematicModelConfig):
def max_jerk(self):
return self.get_state_bounds().jerk[1, 0].item()
@property
def action_horizon(self):
return self._rollout_step_fn.action_horizon
def get_state_bounds(self):
joint_limits = self.robot_model.get_joint_limits()
return joint_limits

77
src/curobo/rollout/dynamics_model/tensor_step.py
View File

@@ -49,12 +49,16 @@ class TensorStepType(Enum):
class TensorStepBase:
def __init__(self, tensor_args: TensorDeviceType) -> None:
def __init__(
self, tensor_args: TensorDeviceType, batch_size: int = 1, horizon: int = 1
) -> None:
self.batch_size = -1
self.horizon = -1
self.tensor_args = tensor_args
self._diag_dt = None
self._inv_dt_h = None
self.action_horizon = horizon
self.update_batch_size(batch_size, horizon)
def update_dt(self, dt: float):
self._dt_h[:] = dt
@@ -83,8 +87,14 @@ class TensorStepBase:
class TensorStepAcceleration(TensorStepBase):
def __init__(self, tensor_args: TensorDeviceType, dt_h: torch.Tensor) -> None:
super().__init__(tensor_args)
def __init__(
self,
tensor_args: TensorDeviceType,
dt_h: torch.Tensor,
batch_size: int = 1,
horizon: int = 1,
) -> None:
super().__init__(tensor_args, batch_size=batch_size, horizon=horizon)
self._dt_h = dt_h
self._diag_dt_h = torch.diag(self._dt_h)
self._integrate_matrix_pos = None
@@ -138,8 +148,13 @@ class TensorStepAcceleration(TensorStepBase):
class TensorStepPositionTeleport(TensorStepBase):
def __init__(self, tensor_args: TensorDeviceType) -> None:
super().__init__(tensor_args)
def __init__(
self,
tensor_args: TensorDeviceType,
batch_size: int = 1,
horizon: int = 1,
) -> None:
super().__init__(tensor_args, batch_size=batch_size, horizon=horizon)
def forward(
self,
@@ -153,8 +168,14 @@ class TensorStepPositionTeleport(TensorStepBase):
class TensorStepPosition(TensorStepBase):
def __init__(self, tensor_args: TensorDeviceType, dt_h: torch.Tensor) -> None:
super().__init__(tensor_args)
def __init__(
self,
tensor_args: TensorDeviceType,
dt_h: torch.Tensor,
batch_size: int = 1,
horizon: int = 1,
) -> None:
super().__init__(tensor_args, batch_size=batch_size, horizon=horizon)
self._dt_h = dt_h
# self._diag_dt_h = torch.diag(1 / self._dt_h)
@@ -185,7 +206,6 @@ class TensorStepPosition(TensorStepBase):
)
self._fd_matrix = torch.diag(1.0 / self._dt_h) @ self._fd_matrix
# self._fd_matrix = self._diag_dt_h @ self._fd_matrix
return super().update_batch_size(batch_size, horizon)
def forward(
@@ -205,8 +225,14 @@ class TensorStepPosition(TensorStepBase):
class TensorStepPositionClique(TensorStepBase):
def __init__(self, tensor_args: TensorDeviceType, dt_h: torch.Tensor) -> None:
super().__init__(tensor_args)
def __init__(
self,
tensor_args: TensorDeviceType,
dt_h: torch.Tensor,
batch_size: int = 1,
horizon: int = 1,
) -> None:
super().__init__(tensor_args, batch_size=batch_size, horizon=horizon)
self._dt_h = dt_h
self._inv_dt_h = 1.0 / dt_h
@@ -281,12 +307,20 @@ class TensorStepPositionClique(TensorStepBase):
class TensorStepAccelerationKernel(TensorStepBase):
def __init__(self, tensor_args: TensorDeviceType, dt_h: torch.Tensor, dof: int) -> None:
super().__init__(tensor_args)
def __init__(
self,
tensor_args: TensorDeviceType,
dt_h: torch.Tensor,
dof: int,
batch_size: int = 1,
horizon: int = 1,
) -> None:
self.dof = dof
super().__init__(tensor_args, batch_size=batch_size, horizon=horizon)
self._dt_h = dt_h
self._u_grad = None
self.dof = dof
def update_batch_size(
self,
@@ -363,13 +397,15 @@ class TensorStepPositionCliqueKernel(TensorStepBase):
filter_velocity: bool = False,
filter_acceleration: bool = False,
filter_jerk: bool = False,
batch_size: int = 1,
horizon: int = 1,
) -> None:
super().__init__(tensor_args)
self.dof = dof
self._fd_mode = finite_difference_mode
super().__init__(tensor_args, batch_size=batch_size, horizon=horizon)
self._dt_h = dt_h
self._inv_dt_h = 1.0 / dt_h
self._u_grad = None
self.dof = dof
self._fd_mode = finite_difference_mode
self._filter_velocity = filter_velocity
self._filter_acceleration = filter_acceleration
self._filter_jerk = filter_jerk
@@ -381,10 +417,6 @@ class TensorStepPositionCliqueKernel(TensorStepBase):
weights = kernel
self._sma_kernel = weights
# self._sma = torch.nn.AvgPool1d(kernel_size=5, stride=2, padding=1).to(
# device=self.tensor_args.device
# )
def update_batch_size(
self,
batch_size: Optional[int] = None,
@@ -392,8 +424,11 @@ class TensorStepPositionCliqueKernel(TensorStepBase):
force_update: bool = False,
) -> None:
if batch_size != self.batch_size or horizon != self.horizon:
self.action_horizon = horizon
if self._fd_mode == 0:
self.action_horizon = horizon - 4
self._u_grad = torch.zeros(
(batch_size, horizon, self.dof),
(batch_size, self.action_horizon, self.dof),
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)

19
src/curobo/rollout/rollout_base.py
View File

@@ -222,6 +222,21 @@ class Goal(Sequence):
links_goal_pose=links_goal_pose,
)
def clone(self):
return Goal(
goal_state=self.goal_state,
goal_pose=self.goal_pose,
current_state=self.current_state,
retract_state=self.retract_state,
batch_pose_idx=self.batch_pose_idx,
batch_world_idx=self.batch_world_idx,
batch_enable_idx=self.batch_enable_idx,
batch_current_state_idx=self.batch_current_state_idx,
batch_retract_state_idx=self.batch_retract_state_idx,
batch_goal_state_idx=self.batch_goal_state_idx,
links_goal_pose=self.links_goal_pose,
)
def _tensor_repeat_seeds(self, tensor, num_seeds):
return tensor_repeat_seeds(tensor, num_seeds)
@@ -498,6 +513,10 @@ class RolloutBase:
def horizon(self) -> int:
raise NotImplementedError
@property
def action_horizon(self) -> int:
return self.horizon
def update_start_state(self, start_state: torch.Tensor):
if self.start_state is None:
self.start_state = start_state

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

@@ -19,11 +19,13 @@ import torch
# CuRobo
from curobo.types.math import Pose
from curobo.util.logger import log_error, log_warn
@dataclass
class CameraObservation:
name: str = "camera_image"
#: rgb image format is BxHxWxchannels
rgb_image: Optional[torch.Tensor] = None
depth_image: Optional[torch.Tensor] = None
image_segmentation: Optional[torch.Tensor] = None

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

@@ -488,7 +488,7 @@ class JointState(State):
return new_js
def trim_trajectory(self, start_idx: int, end_idx: Optional[int] = None):
if end_idx is None:
if end_idx is None or end_idx == 0:
end_idx = self.position.shape[-2]
if len(self.position.shape) < 2:
raise ValueError("JointState does not have horizon")

4
src/curobo/util/sample_lib.py
View File

@@ -130,7 +130,6 @@ class HaltonSampleLib(BaseSampleLib):
else:
self.samples = self.filter_samples(self.samples)
if self.samples.shape[0] != sample_shape[0]:
print(self.samples.shape, sample_shape)
log_error("sampling failed")
return self.samples
@@ -312,6 +311,7 @@ class StompSampleLib(BaseSampleLib):
raise ValueError
# seed = self.seed if base_seed is None else base_seed
self.sample_shape = sample_shape
# self.seed = seed
# torch.manual_seed(self.seed)
halton_samples = self.halton_generator.get_gaussian_samples(sample_shape[0])
@@ -327,6 +327,8 @@ class StompSampleLib(BaseSampleLib):
halton_samples = halton_samples / torch.max(torch.abs(halton_samples))
# halton_samples[:, 0, :] = 0.0
halton_samples[:, -1:, :] = 0.0
if torch.any(torch.isnan(halton_samples)):
log_error("Nan values found in samplelib, installation could have been corrupted")
self.samples = halton_samples
return self.samples

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

@@ -130,6 +130,7 @@ def get_batch_interpolated_trajectory(
tensor_args: TensorDeviceType = TensorDeviceType(),
max_deviation: float = 0.1,
min_dt: float = 0.02,
optimize_dt: bool = True,
):
# compute dt across trajectory:
b, horizon, dof = raw_traj.position.shape # horizon
@@ -146,6 +147,7 @@ def get_batch_interpolated_trajectory(
raw_dt,
min_dt,
horizon,
optimize_dt,
)
# traj_steps contains the tsteps for each trajectory
assert steps_max > 0
@@ -208,6 +210,7 @@ def get_cpu_linear_interpolation(
interpolation_dt=interpolation_dt,
)
retimed_traj[k, tstep:, :] = retimed_traj[k, tstep - 1 : tstep, :]
out_traj_state.position[:] = retimed_traj.to(device=raw_traj.position.device)
return out_traj_state
@@ -417,7 +420,7 @@ def linear_smooth(
y = np.linspace(0, last_step + 3, x.shape[0] + 4)
x = np.concatenate((x, x[-1:], x[-1:], x[-1:], x[-1:]))
elif y is None:
step = float(last_step) / float(x.shape[0] - 1)
step = float(last_step - 1) / float(x.shape[0] - 1)
y = np.ravel([float(i) * step for i in range(x.shape[0])])
# y[-1] = np.floor(y[-1])
@@ -506,9 +509,12 @@ def calculate_tsteps(
raw_dt: float,
min_dt: float,
horizon: int,
optimize_dt: bool = True,
):
# compute scaled dt:
opt_dt = calculate_dt(vel, acc, jerk, max_vel, max_acc, max_jerk, raw_dt, interpolation_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)
return traj_steps, steps_max, opt_dt

98
src/curobo/util/usd_helper.py
View File

@@ -11,7 +11,7 @@
# Standard Library
import math
from typing import List, Optional, Union
from typing import Dict, List, Optional, Union
# Third Party
import numpy as np
@@ -35,7 +35,7 @@ from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
from curobo.util.logger import log_error, log_warn
from curobo.util.logger import log_error, log_info, log_warn
from curobo.util_file import (
file_exists,
get_assets_path,
@@ -149,7 +149,11 @@ def set_geom_mesh_attrs(mesh_geom: UsdGeom.Mesh, obs: Mesh, timestep=None):
primvar = primvarsapi.CreatePrimvar(
"displayColor", Sdf.ValueTypeNames.Color3f, interpolation="faceVarying"
)
primvar.Set([Gf.Vec3f(x[0] / 255.0, x[1] / 255.0, x[2] / 255.0) for x in obs.vertex_colors])
scale = 1.0
# color needs to be in range of 0-1. Hence converting if the color is in [0,255]
if max(np.ravel(obs.vertex_colors) > 1.0):
scale = 255.0
primvar.Set([Gf.Vec3f(x[0] / scale, x[1] / scale, x[2] / scale) for x in obs.vertex_colors])
# low = np.min(verts, axis=0)
# high = np.max(verts, axis=0)
@@ -800,9 +804,12 @@ class UsdHelper:
kin_model: Optional[CudaRobotModel] = None,
visualize_robot_spheres: bool = True,
robot_color: Optional[List[float]] = None,
flatten_usd: bool = False,
):
if kin_model is None:
config_file = load_yaml(join_path(get_robot_configs_path(), robot_model_file))
if "robot_cfg" not in config_file:
config_file["robot_cfg"] = config_file
config_file["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
robot_cfg = CudaRobotModelConfig.from_data_dict(
config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args
@@ -852,7 +859,7 @@ class UsdHelper:
usd_helper.create_obstacle_animation(
sphere_traj, base_frame=base_frame, obstacles_frame="curobo/robot_collision"
)
usd_helper.write_stage_to_file(save_path)
usd_helper.write_stage_to_file(save_path, flatten=flatten_usd)
@staticmethod
def load_robot(
@@ -890,19 +897,20 @@ class UsdHelper:
visualize_robot_spheres: bool = True,
robot_asset_prim_path=None,
robot_color: Optional[List[float]] = None,
flatten_usd: bool = False,
):
usd_exists = False
# if usd file doesn't exist, fall back to urdf animation script
if kin_model is None:
config_file = load_yaml(join_path(get_robot_configs_path(), robot_model_file))
if "robot_cfg" in config_file:
config_file = config_file["robot_cfg"]
config_file["kinematics"]["load_link_names_with_mesh"] = True
config_file["kinematics"]["use_usd_kinematics"] = True
robot_model_file = load_yaml(join_path(get_robot_configs_path(), robot_model_file))
if "robot_cfg" in robot_model_file:
robot_model_file = robot_model_file["robot_cfg"]
robot_model_file["kinematics"]["load_link_names_with_mesh"] = True
robot_model_file["kinematics"]["use_usd_kinematics"] = True
usd_exists = file_exists(
join_path(get_assets_path(), config_file["kinematics"]["usd_path"])
join_path(get_assets_path(), robot_model_file["kinematics"]["usd_path"])
)
else:
if kin_model.generator_config.usd_path is not None:
@@ -914,7 +922,8 @@ class UsdHelper:
)
usd_exists = False
if not usd_exists:
log_warn("robot usd not found, using urdf animation instead")
log_info("robot usd not found, using urdf animation instead")
robot_model_file["kinematics"]["use_usd_kinematics"] = False
return UsdHelper.write_trajectory_animation(
robot_model_file,
world_model,
@@ -929,10 +938,11 @@ class UsdHelper:
kin_model=kin_model,
visualize_robot_spheres=visualize_robot_spheres,
robot_color=robot_color,
flatten_usd=flatten_usd,
)
if kin_model is None:
robot_cfg = CudaRobotModelConfig.from_data_dict(
config_file["kinematics"], tensor_args=tensor_args
robot_model_file["kinematics"], tensor_args=tensor_args
)
kin_model = CudaRobotModel(robot_cfg)
@@ -978,7 +988,7 @@ class UsdHelper:
usd_helper.create_obstacle_animation(
sphere_traj, base_frame=base_frame, obstacles_frame="curobo/robot_collision"
)
usd_helper.write_stage_to_file(save_path)
usd_helper.write_stage_to_file(save_path, flatten=flatten_usd)
@staticmethod
def create_grid_usd(
@@ -994,6 +1004,7 @@ class UsdHelper:
dt: float = 0.02,
interpolation_steps: int = 1,
prefix_string: Optional[str] = None,
flatten_usd: bool = False,
):
# create stage:
usd_helper = UsdHelper()
@@ -1036,7 +1047,7 @@ class UsdHelper:
# write usd to disk:
usd_helper.write_stage_to_file(save_path)
usd_helper.write_stage_to_file(save_path, flatten=flatten_usd)
def load_robot_usd(
self,
@@ -1138,6 +1149,9 @@ class UsdHelper:
grid_space: float = 1.0,
write_robot_usd_path="assets/",
robot_asset_prim_path="/panda",
fps: int = 24,
link_poses: Optional[Dict[str, Pose]] = None,
flatten_usd: bool = False,
):
if goal_object is None:
log_warn("Using franka gripper as goal object")
@@ -1150,10 +1164,26 @@ class UsdHelper:
color=[0.0, 0.8, 0.1, 1.0],
pose=goal_pose.to_list(),
)
if goal_object is not None:
goal_object.pose = np.ravel(goal_pose.tolist()).tolist()
world_config = world_config.clone()
world_config.add_obstacle(goal_object)
if link_poses is not None:
link_goals = []
for k in link_poses.keys():
link_goals.append(
Mesh(
name="target_" + k,
file_path=join_path(
get_assets_path(),
"robot/franka_description/meshes/visual/hand.dae",
),
color=[0.0, 0.8, 0.1, 1.0],
pose=link_poses[k].to_list(),
)
)
world_config.add_obstacle(link_goals[-1])
kin_model = UsdHelper.load_robot(robot_file)
if link_spheres is not None:
kin_model.kinematics_config.link_spheres = link_spheres
@@ -1173,7 +1203,7 @@ class UsdHelper:
num_seeds, n_iters, dof = vis_traj.shape
usd_paths = []
for j in tqdm(range(num_seeds)):
current_traj = vis_traj[j].view(-1, dof)[:-1, :] # we remove last timestep
current_traj = vis_traj[j].view(-1, dof)[:, :] # we remove last timestep
usd_paths.append(save_prefix + "_ik_seed_" + str(j) + ".usd")
UsdHelper.write_trajectory_animation_with_robot_usd(
@@ -1181,13 +1211,14 @@ class UsdHelper:
world_config,
start_state,
JointState.from_position(current_traj),
dt=(1),
dt=(1 / fps),
save_path=usd_paths[-1],
base_frame="/world_base_" + str(j),
kin_model=kin_model,
visualize_robot_spheres=visualize_robot_spheres,
write_robot_usd_path=write_robot_usd_path,
robot_asset_prim_path=robot_asset_prim_path,
flatten_usd=flatten_usd,
)
UsdHelper.create_grid_usd(
@@ -1200,7 +1231,7 @@ class UsdHelper:
y_space=y_space,
x_per_row=int(math.floor(math.sqrt(len(usd_paths)))),
local_asset_path="",
dt=(1),
dt=(1.0 / fps),
)
if write_trajopt:
if "trajopt_result" not in result.debug_info:
@@ -1226,8 +1257,9 @@ class UsdHelper:
(start_state.position.view(1, 1, 1, -1).repeat(n1, n2, 1, 1), full_traj), dim=-2
)
usd_paths = []
finetune_usd_paths = []
for j in tqdm(range(num_seeds)):
current_traj = full_traj[j].view(-1, dof)[:-1, :] # we remove last timestep
current_traj = full_traj[j].view(-1, dof) # we remove last timestep
# add start state to current trajectory since it's not in the optimization:
usd_paths.append(save_prefix + "_trajopt_seed_" + str(j) + ".usd")
UsdHelper.write_trajectory_animation_with_robot_usd(
@@ -1235,13 +1267,14 @@ class UsdHelper:
world_config,
start_state,
JointState.from_position(current_traj),
dt=(1 / 24),
dt=(1.0 / fps),
save_path=usd_paths[-1],
base_frame="/world_base_" + str(j),
kin_model=kin_model,
visualize_robot_spheres=visualize_robot_spheres,
write_robot_usd_path=write_robot_usd_path,
robot_asset_prim_path=robot_asset_prim_path,
flatten_usd=flatten_usd,
)
# add finetuning step:
@@ -1257,6 +1290,7 @@ class UsdHelper:
full_traj = torch.cat(vis_traj, dim=0)
num_seeds, h, dof = vis_traj[-1].shape
n, _, _ = full_traj.shape # this will have iterations
# print(full_traj.shape)
full_traj = full_traj.view(-1, num_seeds, h, dof)
n1, n2, _, _ = full_traj.shape
@@ -1268,25 +1302,28 @@ class UsdHelper:
full_traj = full_traj.transpose(0, 1).contiguous() # n_seeds, n_steps, h, dof
for j in tqdm(range(num_seeds)):
current_traj = full_traj[j].view(-1, dof)[:-1, :] # we remove last timestep
current_traj = full_traj[j][-1].view(-1, dof)
usd_paths.append(save_prefix + "_trajopt_finetune_seed_" + str(j) + ".usd")
finetune_usd_paths.append(save_prefix + "_finetune_seed_" + str(j) + ".usd")
UsdHelper.write_trajectory_animation_with_robot_usd(
robot_file,
world_config,
start_state,
JointState.from_position(current_traj),
dt=(1 / 24),
save_path=usd_paths[-1],
dt=(1.0 / fps),
save_path=finetune_usd_paths[-1],
base_frame="/world_base_" + str(j),
kin_model=kin_model,
visualize_robot_spheres=visualize_robot_spheres,
write_robot_usd_path=write_robot_usd_path,
robot_asset_prim_path=robot_asset_prim_path,
flatten_usd=flatten_usd,
)
x_space = y_space = grid_space
if overlay_trajopt:
x_space = y_space = 0.0
if len(finetune_usd_paths) == 1:
usd_paths.append(finetune_usd_paths[0])
UsdHelper.create_grid_usd(
usd_paths,
save_prefix + "_grid_trajopt.usd",
@@ -1297,5 +1334,18 @@ class UsdHelper:
y_space=y_space,
x_per_row=int(math.floor(math.sqrt(len(usd_paths)))),
local_asset_path="",
dt=(1 / 24),
dt=(1.0 / fps),
)
if False and len(finetune_usd_paths) > 1:
UsdHelper.create_grid_usd(
finetune_usd_paths,
save_prefix + "_grid_finetune_trajopt.usd",
base_frame="/world",
max_envs=len(finetune_usd_paths),
max_timecode=n_steps * h,
x_space=x_space,
y_space=y_space,
x_per_row=int(math.floor(math.sqrt(len(finetune_usd_paths)))),
local_asset_path="",
dt=(1.0 / fps),
)

1
src/curobo/util_file.py
View File

@@ -168,7 +168,6 @@ def get_motion_gen_robot_list() -> List[str]:
"franka.yml",
"ur5e.yml",
"ur10e.yml",
"dual_ur10e.yml",
"tm12.yml",
"jaco7.yml",
"kinova_gen3.yml",

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

@@ -122,7 +122,7 @@ class MotionGenConfig:
interpolation_dt: float = 0.01
#: scale initial dt by this value to finetune trajectory optimization.
finetune_dt_scale: float = 1.05
finetune_dt_scale: float = 0.98
@staticmethod
@profiler.record_function("motion_gen_config/load_from_robot_config")
@@ -192,12 +192,13 @@ class MotionGenConfig:
smooth_weight: List[float] = None,
finetune_smooth_weight: Optional[List[float]] = None,
state_finite_difference_mode: Optional[str] = None,
finetune_dt_scale: float = 1.05,
finetune_dt_scale: float = 0.98,
maximum_trajectory_time: Optional[float] = None,
maximum_trajectory_dt: float = 0.1,
velocity_scale: Optional[Union[List[float], float]] = None,
acceleration_scale: Optional[Union[List[float], float]] = None,
jerk_scale: Optional[Union[List[float], float]] = None,
optimize_dt: bool = True,
):
"""Load motion generation configuration from robot and world configurations.
@@ -279,7 +280,11 @@ class MotionGenConfig:
"""
init_warp(tensor_args=tensor_args)
if js_trajopt_tsteps is not None:
log_warn("js_trajopt_tsteps is deprecated, use trajopt_tsteps instead.")
trajopt_tsteps = js_trajopt_tsteps
if trajopt_tsteps is not None:
js_trajopt_tsteps = trajopt_tsteps
if velocity_scale is not None and isinstance(velocity_scale, float):
velocity_scale = [velocity_scale]
@@ -318,6 +323,8 @@ class MotionGenConfig:
if isinstance(robot_cfg, str):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"]
elif isinstance(robot_cfg, Dict) and "robot_cfg" in robot_cfg.keys():
robot_cfg = robot_cfg["robot_cfg"]
if isinstance(robot_cfg, RobotConfig):
if ee_link_name is not None:
log_error("ee link cannot be changed after creating RobotConfig")
@@ -445,6 +452,7 @@ class MotionGenConfig:
state_finite_difference_mode=state_finite_difference_mode,
filter_robot_command=filter_robot_command,
minimize_jerk=minimize_jerk,
optimize_dt=optimize_dt,
)
trajopt_solver = TrajOptSolver(trajopt_cfg)
@@ -465,8 +473,6 @@ class MotionGenConfig:
self_collision_check=self_collision_check,
self_collision_opt=self_collision_opt,
grad_trajopt_iters=grad_trajopt_iters,
# num_seeds=num_trajopt_noisy_seeds,
# seed_ratio=trajopt_seed_ratio,
interpolation_dt=interpolation_dt,
use_particle_opt=trajopt_particle_opt,
traj_evaluator_config=traj_evaluator_config,
@@ -486,6 +492,7 @@ class MotionGenConfig:
state_finite_difference_mode=state_finite_difference_mode,
minimize_jerk=minimize_jerk,
filter_robot_command=filter_robot_command,
optimize_dt=optimize_dt,
)
js_trajopt_solver = TrajOptSolver(js_trajopt_cfg)
@@ -523,6 +530,7 @@ class MotionGenConfig:
trim_steps=trim_steps,
use_gradient_descent=use_gradient_descent,
filter_robot_command=filter_robot_command,
optimize_dt=optimize_dt,
)
finetune_trajopt_solver = TrajOptSolver(finetune_trajopt_cfg)
@@ -748,7 +756,9 @@ class MotionGenResult:
@property
def motion_time(self):
return self.optimized_dt * (self.optimized_plan.position.shape[-2] - 1)
# -2 as last three timesteps have the same value
# 0, 1 also have the same position value.
return self.optimized_dt * (self.optimized_plan.position.shape[-2] - 1 - 2 - 1)
@dataclass
@@ -762,7 +772,7 @@ class MotionGenPlanConfig:
enable_graph_attempt: Optional[int] = 3
disable_graph_attempt: Optional[int] = None
ik_fail_return: Optional[int] = None
partial_ik_opt: bool = True
partial_ik_opt: bool = False
num_ik_seeds: Optional[int] = None
num_graph_seeds: Optional[int] = None
num_trajopt_seeds: Optional[int] = None
@@ -770,6 +780,7 @@ class MotionGenPlanConfig:
fail_on_invalid_query: bool = True
#: enables retiming trajectory optimization, useful for getting low jerk trajectories.
enable_finetune_trajopt: bool = True
parallel_finetune: bool = False
def __post_init__(self):
if not self.enable_opt and not self.enable_graph:
@@ -1103,8 +1114,11 @@ class MotionGen(MotionGenConfig):
if plan_config.enable_graph:
raise ValueError("Graph Search / Geometric Planner not supported in batch_env mode")
if plan_config.enable_graph:
log_info("Batch mode enable graph is only supported with num_graph_seeds==1")
if plan_config.enable_graph or (
plan_config.enable_graph_attempt is not None
and plan_config.max_attempts >= plan_config.enable_graph_attempt
):
log_warn("Batch mode enable graph is only supported with num_graph_seeds==1")
plan_config.num_trajopt_seeds = 1
plan_config.num_graph_seeds = 1
solve_state.num_trajopt_seeds = 1
@@ -1316,7 +1330,6 @@ class MotionGen(MotionGenConfig):
trajopt_seed_traj = None
trajopt_seed_success = None
trajopt_newton_iters = None
graph_success = 0
if len(start_state.shape) == 1:
log_error("Joint state should be not a vector (dof) should be (bxdof)")
@@ -1330,6 +1343,7 @@ class MotionGen(MotionGenConfig):
plan_config.partial_ik_opt,
link_poses,
)
if not plan_config.enable_graph and plan_config.partial_ik_opt:
ik_result.success[:] = True
@@ -1364,7 +1378,7 @@ class MotionGen(MotionGenConfig):
if plan_config.enable_graph:
interpolation_steps = None
if plan_config.enable_opt:
interpolation_steps = self.trajopt_solver.traj_tsteps - 4
interpolation_steps = self.trajopt_solver.action_horizon
log_info("MG: running GP")
graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
trajopt_seed_success = graph_result.success
@@ -1378,6 +1392,8 @@ class MotionGen(MotionGenConfig):
result.used_graph = True
if plan_config.enable_opt:
# print(result.graph_plan.position.shape, interpolation_steps,
# graph_result.path_buffer_last_tstep)
trajopt_seed = (
result.graph_plan.position.view(
1, # solve_state.batch_size,
@@ -1389,12 +1405,11 @@ class MotionGen(MotionGenConfig):
.contiguous()
)
trajopt_seed_traj = torch.zeros(
(trajopt_seed.shape[0], 1, self.trajopt_solver.traj_tsteps, self._dof),
(trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof),
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
trajopt_seed_traj[:, :, :-4, :] = trajopt_seed
trajopt_seed_traj[:, :, -4:, :] = trajopt_seed_traj[:, :, -5:-4, :]
trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed
trajopt_seed_success = ik_result.success.clone()
trajopt_seed_success[ik_result.success] = graph_result.success
@@ -1497,7 +1512,7 @@ class MotionGen(MotionGenConfig):
trajopt_seed_traj = trajopt_seed_traj.view(
solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
solve_state.batch_size,
self.trajopt_solver.traj_tsteps,
self.trajopt_solver.action_horizon,
self._dof,
).contiguous()
if plan_config.enable_finetune_trajopt:
@@ -1511,31 +1526,27 @@ class MotionGen(MotionGenConfig):
trajopt_seed_traj,
num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
newton_iters=trajopt_newton_iters,
return_all_solutions=plan_config.parallel_finetune,
)
if False and not traj_result.success.item():
# pose_convergence = traj_result.position_error < self.
print(
traj_result.position_error.item(),
traj_result.rotation_error.item(),
torch.count_nonzero(~traj_result.metrics.feasible[0]).item(),
torch.count_nonzero(~traj_result.metrics.feasible[1]).item(),
traj_result.optimized_dt.item(),
)
if plan_config.enable_finetune_trajopt:
self.trajopt_solver.interpolation_type = og_value
# self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate)
if self.store_debug_in_result:
result.debug_info["trajopt_result"] = traj_result
# run finetune
if plan_config.enable_finetune_trajopt and traj_result.success[0].item():
if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
with profiler.record_function("motion_gen/finetune_trajopt"):
seed_traj = traj_result.solution.position.clone()
seed_traj = torch.roll(seed_traj, -2, dims=-2)
# seed_traj[..., -2:, :] = seed_traj[..., -3, :]
seed_traj = traj_result.raw_action.clone() # solution.position.clone()
seed_traj = seed_traj.contiguous()
og_solve_time = traj_result.solve_time
seed_override = 1
opt_dt = traj_result.optimized_dt
if plan_config.parallel_finetune:
opt_dt = torch.min(opt_dt[traj_result.success])
seed_override = solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
scaled_dt = torch.clamp(
traj_result.optimized_dt * self.finetune_dt_scale,
opt_dt * self.finetune_dt_scale,
self.trajopt_solver.interpolation_dt,
)
self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item())
@@ -1545,26 +1556,16 @@ class MotionGen(MotionGenConfig):
solve_state,
seed_traj,
trajopt_instance=self.finetune_trajopt_solver,
num_seeds_override=1,
num_seeds_override=seed_override,
)
if False and not traj_result.success.item():
print(
traj_result.position_error.item(),
traj_result.rotation_error.item(),
torch.count_nonzero(~traj_result.metrics.feasible).item(),
traj_result.optimized_dt.item(),
)
# if not traj_result.success.item():
# #print(traj_result.metrics.constraint)
# print(traj_result.position_error.item() * 100.0,
# traj_result.rotation_error.item() * 100.0)
result.finetune_time = traj_result.solve_time
traj_result.solve_time = og_solve_time
if self.store_debug_in_result:
result.debug_info["finetune_trajopt_result"] = traj_result
elif plan_config.enable_finetune_trajopt:
traj_result.success = traj_result.success[0:1]
result.solve_time += traj_result.solve_time + result.finetune_time
result.trajopt_time = traj_result.solve_time
result.trajopt_attempts = 1
@@ -1576,12 +1577,220 @@ class MotionGen(MotionGenConfig):
result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory(
0, traj_result.path_buffer_last_tstep[0]
)
# print(ik_result.position_error[ik_result.success] * 1000.0)
# print(traj_result.position_error * 1000.0)
# exit()
result.interpolation_dt = self.trajopt_solver.interpolation_dt
result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
result.position_error = traj_result.position_error
result.rotation_error = traj_result.rotation_error
result.optimized_dt = traj_result.optimized_dt
result.optimized_plan = traj_result.solution
return result
def _plan_js_from_solve_state(
self,
solve_state: ReacherSolveState,
start_state: JointState,
goal_state: JointState,
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
) -> MotionGenResult:
trajopt_seed_traj = None
trajopt_seed_success = None
trajopt_newton_iters = None
graph_success = 0
if len(start_state.shape) == 1:
log_error("Joint state should be not a vector (dof) should be (bxdof)")
result = MotionGenResult(cspace_error=torch.zeros((1), device=self.tensor_args.device))
# do graph search:
if plan_config.enable_graph:
start_config = torch.zeros(
(solve_state.num_graph_seeds, self.js_trajopt_solver.dof),
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
goal_config = start_config.clone()
start_config[:] = start_state.position
goal_config[:] = goal_state.position
interpolation_steps = None
if plan_config.enable_opt:
interpolation_steps = self.js_trajopt_solver.action_horizon
log_info("MG: running GP")
graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
trajopt_seed_success = graph_result.success
graph_success = torch.count_nonzero(graph_result.success).item()
result.graph_time = graph_result.solve_time
result.solve_time += graph_result.solve_time
if graph_success > 0:
result.graph_plan = graph_result.interpolated_plan
result.interpolated_plan = graph_result.interpolated_plan
result.used_graph = True
if plan_config.enable_opt:
trajopt_seed = (
result.graph_plan.position.view(
1, # solve_state.batch_size,
graph_success, # solve_state.num_trajopt_seeds,
interpolation_steps,
self._dof,
)
.transpose(0, 1)
.contiguous()
)
trajopt_seed_traj = torch.zeros(
(trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof),
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed
trajopt_seed_success = graph_result.success
trajopt_seed_success = trajopt_seed_success.view(
1, solve_state.num_trajopt_seeds
)
trajopt_newton_iters = self.graph_trajopt_iters
else:
_, idx = torch.topk(
graph_result.path_length[graph_result.success], k=1, largest=False
)
result.interpolated_plan = result.interpolated_plan[idx].squeeze(0)
result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt)
result.optimized_plan = result.interpolated_plan[
: graph_result.path_buffer_last_tstep[idx.item()]
]
idx = idx.view(-1) + self._batch_col
result.cspace_error = torch.zeros((1), device=self.tensor_args.device)
result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[
idx.item() : idx.item() + 1
]
result.success = torch.as_tensor([True], device=self.tensor_args.device)
return result
else:
result.success = torch.as_tensor([False], device=self.tensor_args.device)
result.status = "Graph Fail"
if not graph_result.valid_query:
result.valid_query = False
if self.store_debug_in_result:
result.debug_info["graph_debug"] = graph_result.debug_info
return result
if plan_config.need_graph_success:
return result
# do trajopt:
if plan_config.enable_opt:
with profiler.record_function("motion_gen/setup_trajopt_seeds"):
# self._trajopt_goal_config[:, :ik_success] = goal_config
goal = Goal(
current_state=start_state,
goal_state=goal_state,
)
if trajopt_seed_traj is None or graph_success < solve_state.num_trajopt_seeds * 1:
seed_goal = Goal(
current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
goal_state=goal_state.repeat_seeds(solve_state.num_trajopt_seeds),
)
if trajopt_seed_traj is not None:
trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
# batch, num_seeds, h, dof
if trajopt_seed_success.shape[1] < self.js_trajopt_solver.num_seeds:
trajopt_seed_success_new = torch.zeros(
(1, solve_state.num_trajopt_seeds),
device=self.tensor_args.device,
dtype=torch.bool,
)
trajopt_seed_success_new[
0, : trajopt_seed_success.shape[1]
] = trajopt_seed_success
trajopt_seed_success = trajopt_seed_success_new
# create seeds here:
trajopt_seed_traj = self.js_trajopt_solver.get_seed_set(
seed_goal,
trajopt_seed_traj, # batch, num_seeds, h, dof
num_seeds=1,
batch_mode=False,
seed_success=trajopt_seed_success,
)
trajopt_seed_traj = trajopt_seed_traj.view(
self.js_trajopt_solver.num_seeds * 1,
1,
self.trajopt_solver.action_horizon,
self._dof,
).contiguous()
if plan_config.enable_finetune_trajopt:
og_value = self.trajopt_solver.interpolation_type
self.js_trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
with profiler.record_function("motion_gen/trajopt"):
log_info("MG: running TO")
traj_result = self._solve_trajopt_from_solve_state(
goal,
solve_state,
trajopt_seed_traj,
num_seeds_override=solve_state.num_trajopt_seeds * 1,
newton_iters=trajopt_newton_iters,
return_all_solutions=plan_config.enable_finetune_trajopt,
trajopt_instance=self.js_trajopt_solver,
)
if plan_config.enable_finetune_trajopt:
self.trajopt_solver.interpolation_type = og_value
# self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate)
if self.store_debug_in_result:
result.debug_info["trajopt_result"] = traj_result
if torch.count_nonzero(traj_result.success) == 0:
result.status = "TrajOpt Fail"
# run finetune
if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
with profiler.record_function("motion_gen/finetune_trajopt"):
seed_traj = traj_result.raw_action.clone() # solution.position.clone()
seed_traj = seed_traj.contiguous()
og_solve_time = traj_result.solve_time
scaled_dt = torch.clamp(
torch.max(traj_result.optimized_dt[traj_result.success]),
self.trajopt_solver.interpolation_dt,
)
og_dt = self.js_trajopt_solver.solver_dt
self.js_trajopt_solver.update_solver_dt(scaled_dt.item())
traj_result = self._solve_trajopt_from_solve_state(
goal,
solve_state,
seed_traj,
trajopt_instance=self.js_trajopt_solver,
num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
)
self.js_trajopt_solver.update_solver_dt(og_dt)
result.finetune_time = traj_result.solve_time
traj_result.solve_time = og_solve_time
if self.store_debug_in_result:
result.debug_info["finetune_trajopt_result"] = traj_result
if torch.count_nonzero(traj_result.success) == 0:
result.status = "Finetune Fail"
elif plan_config.enable_finetune_trajopt:
traj_result.success = traj_result.success[0:1]
result.solve_time += traj_result.solve_time + result.finetune_time
result.trajopt_time = traj_result.solve_time
result.trajopt_attempts = 1
result.success = traj_result.success
result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory(
0, traj_result.path_buffer_last_tstep[0]
)
# print(ik_result.position_error[ik_result.success] * 1000.0)
# print(traj_result.position_error * 1000.0)
# exit()
result.interpolation_dt = self.trajopt_solver.interpolation_dt
result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
result.cspace_error = traj_result.cspace_error
result.optimized_dt = traj_result.optimized_dt
result.optimized_plan = traj_result.solution
return result
@@ -1644,7 +1853,7 @@ class MotionGen(MotionGenConfig):
if plan_config.enable_graph:
interpolation_steps = None
if plan_config.enable_opt:
interpolation_steps = self.trajopt_solver.traj_tsteps - 4
interpolation_steps = self.trajopt_solver.action_horizon
start_graph_state = start_state.repeat_seeds(ik_out_seeds)
start_config = start_graph_state.position[ik_result.success.view(-1)].view(
@@ -1662,23 +1871,17 @@ class MotionGen(MotionGenConfig):
result.used_graph = True
if plan_config.enable_opt:
trajopt_seed = (
result.graph_plan.position.view(
1, # solve_state.batch_size,
graph_success, # solve_state.num_trajopt_seeds,
interpolation_steps,
self._dof,
)
.transpose(0, 1)
.contiguous()
)
trajopt_seed = result.graph_plan.position.view(
graph_success, # solve_state.num_trajopt_seeds,
interpolation_steps,
self._dof,
).contiguous()
trajopt_seed_traj = torch.zeros(
(trajopt_seed.shape[0], 1, self.trajopt_solver.traj_tsteps, self._dof),
(1, trajopt_seed.shape[0], self.trajopt_solver.action_horizon, self._dof),
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
trajopt_seed_traj[:, :, :-4, :] = trajopt_seed
trajopt_seed_traj[:, :, -4:, :] = trajopt_seed_traj[:, :, -5:-4, :]
trajopt_seed_traj[0, :, :interpolation_steps, :] = trajopt_seed
trajopt_seed_success = ik_result.success.clone()
trajopt_seed_success[ik_result.success] = graph_result.success
@@ -1766,14 +1969,54 @@ class MotionGen(MotionGenConfig):
trajopt_seed_traj = trajopt_seed_traj.view(
solve_state.num_trajopt_seeds,
solve_state.batch_size,
self.trajopt_solver.traj_tsteps,
self.trajopt_solver.action_horizon,
self._dof,
).contiguous()
if plan_config.enable_finetune_trajopt:
og_value = self.trajopt_solver.interpolation_type
self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
traj_result = self._solve_trajopt_from_solve_state(
goal, solve_state, trajopt_seed_traj, newton_iters=trajopt_newton_iters
goal,
solve_state,
trajopt_seed_traj,
newton_iters=trajopt_newton_iters,
return_all_solutions=True,
)
# output of traj result will have 1 solution per batch
# run finetune opt on 1 solution per batch:
if plan_config.enable_finetune_trajopt:
self.trajopt_solver.interpolation_type = og_value
# self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate)
if self.store_debug_in_result:
result.debug_info["trajopt_result"] = traj_result
# run finetune
if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
with profiler.record_function("motion_gen/finetune_trajopt"):
seed_traj = traj_result.raw_action.clone() # solution.position.clone()
seed_traj = seed_traj.contiguous()
og_solve_time = traj_result.solve_time
scaled_dt = torch.clamp(
torch.max(traj_result.optimized_dt[traj_result.success])
* self.finetune_dt_scale,
self.trajopt_solver.interpolation_dt,
)
self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item())
traj_result = self._solve_trajopt_from_solve_state(
goal,
solve_state,
seed_traj,
trajopt_instance=self.finetune_trajopt_solver,
num_seeds_override=solve_state.num_trajopt_seeds,
)
result.finetune_time = traj_result.solve_time
traj_result.solve_time = og_solve_time
if self.store_debug_in_result:
result.debug_info["finetune_trajopt_result"] = traj_result
result.success = traj_result.success
result.interpolated_plan = traj_result.interpolated_solution
@@ -1837,6 +2080,7 @@ class MotionGen(MotionGenConfig):
batch: Optional[int] = None,
warmup_js_trajopt: bool = True,
batch_env_mode: bool = False,
parallel_finetune: bool = False,
):
log_info("Warmup")
@@ -1854,7 +2098,11 @@ class MotionGen(MotionGenConfig):
self.plan_single(
start_state,
retract_pose,
MotionGenPlanConfig(max_attempts=1, enable_finetune_trajopt=True),
MotionGenPlanConfig(
max_attempts=1,
enable_finetune_trajopt=True,
parallel_finetune=parallel_finetune,
),
link_poses=link_poses,
)
if enable_graph:
@@ -1867,7 +2115,10 @@ class MotionGen(MotionGenConfig):
start_state,
retract_pose,
MotionGenPlanConfig(
max_attempts=1, enable_finetune_trajopt=True, enable_graph=enable_graph
max_attempts=1,
enable_finetune_trajopt=True,
enable_graph=enable_graph,
parallel_finetune=parallel_finetune,
),
link_poses=link_poses,
)
@@ -1925,14 +2176,24 @@ class MotionGen(MotionGenConfig):
}
result = None
goal = Goal(goal_state=goal_state, current_state=start_state)
solve_state = ReacherSolveState(
ReacherSolveType.SINGLE,
num_ik_seeds=1,
num_trajopt_seeds=self.js_trajopt_solver.num_seeds,
num_graph_seeds=self.js_trajopt_solver.num_seeds,
batch_size=1,
n_envs=1,
n_goalset=1,
)
for n in range(plan_config.max_attempts):
traj_result = self.js_trajopt_solver.solve_single(goal)
traj_result = self._plan_js_from_solve_state(
solve_state, start_state, goal_state, plan_config=plan_config
)
time_dict["trajopt_time"] += traj_result.solve_time
time_dict["trajopt_attempts"] = n
if result is None:
result = MotionGenResult(success=traj_result.success)
result = traj_result
if traj_result.success.item():
break
@@ -1940,25 +2201,7 @@ class MotionGen(MotionGenConfig):
result.solve_time = time_dict["trajopt_time"]
if self.store_debug_in_result:
result.debug_info = {"trajopt_result": traj_result}
status = None
if not traj_result.success.item():
# print(traj_result.cspace_error, traj_result.success)
status = ""
if traj_result.cspace_error.item() >= self.js_trajopt_solver.cspace_threshold:
status += " Fail: C-SPACE Convergence"
if torch.count_nonzero(~traj_result.metrics.feasible).item() > 0:
status += " Fail: Constraints"
# print(traj_result.metrics.feasible)
result.status = status
result.position_error = traj_result.position_error
result.rotation_error = traj_result.rotation_error
result.cspace_error = traj_result.cspace_error
result.optimized_dt = traj_result.optimized_dt
result.interpolated_plan = traj_result.interpolated_solution
result.optimized_plan = traj_result.solution
result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
result.success = traj_result.success
return result
def plan(

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

@@ -77,6 +77,7 @@ class TrajOptSolverConfig:
use_cuda_graph_metrics: bool = False
trim_steps: Optional[List[int]] = None
store_debug_in_result: bool = False
optimize_dt: bool = True
@staticmethod
@profiler.record_function("trajopt_config/load_from_robot_config")
@@ -107,7 +108,7 @@ class TrajOptSolverConfig:
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE,
traj_evaluator_config: TrajEvaluatorConfig = TrajEvaluatorConfig(),
traj_evaluator: Optional[TrajEvaluator] = None,
minimize_jerk: Optional[bool] = None,
minimize_jerk: bool = True,
use_gradient_descent: bool = False,
collision_cache: Optional[Dict[str, int]] = None,
n_collision_envs: Optional[int] = None,
@@ -126,7 +127,9 @@ class TrajOptSolverConfig:
smooth_weight: Optional[List[float]] = None,
state_finite_difference_mode: Optional[str] = None,
filter_robot_command: bool = False,
optimize_dt: bool = True,
):
# NOTE: Don't have default optimize_dt, instead read from a configuration file.
# use default values, disable environment collision checking
if isinstance(robot_cfg, str):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"]
@@ -199,6 +202,7 @@ class TrajOptSolverConfig:
fixed_iters = True
grad_config_data["lbfgs"]["store_debug"] = store_debug
config_data["mppi"]["store_debug"] = store_debug
store_debug_in_result = True
if use_cuda_graph is not None:
config_data["mppi"]["use_cuda_graph"] = use_cuda_graph
@@ -332,6 +336,7 @@ class TrajOptSolverConfig:
use_cuda_graph_metrics=use_cuda_graph,
trim_steps=trim_steps,
store_debug_in_result=store_debug_in_result,
optimize_dt=optimize_dt,
)
return trajopt_cfg
@@ -354,6 +359,7 @@ class TrajResult(Sequence):
smooth_label: Optional[T_BValue_bool] = None
optimized_dt: Optional[torch.Tensor] = None
raw_solution: Optional[JointState] = None
raw_action: Optional[torch.Tensor] = None
def __getitem__(self, idx):
# position_error = rotation_error = cspace_error = path_buffer_last_tstep = None
@@ -392,17 +398,14 @@ class TrajOptSolver(TrajOptSolverConfig):
def __init__(self, config: TrajOptSolverConfig) -> None:
super().__init__(**vars(config))
self.dof = self.rollout_fn.d_action
self.delta_vec = action_interpolate_kernel(2, self.traj_tsteps, self.tensor_args).unsqueeze(
0
)
self.waypoint_delta_vec = interpolate_kernel(
3, int(self.traj_tsteps / 2), self.tensor_args
).unsqueeze(0)
self.waypoint_delta_vec = torch.roll(self.waypoint_delta_vec, -1, dims=1)
self.waypoint_delta_vec[:, -1, :] = self.waypoint_delta_vec[:, -2, :]
assert self.traj_tsteps / 2 != 0.0
self.solver.update_nenvs(self.num_seeds)
self.action_horizon = self.rollout_fn.action_horizon
self.delta_vec = interpolate_kernel(2, self.action_horizon, self.tensor_args).unsqueeze(0)
self.waypoint_delta_vec = interpolate_kernel(
3, int(self.action_horizon / 2), self.tensor_args
).unsqueeze(0)
assert self.action_horizon / 2 != 0.0
self.solver.update_nenvs(self.num_seeds)
self._max_joint_vel = (
self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[1, :].reshape(
1, 1, self.dof
@@ -410,7 +413,6 @@ class TrajOptSolver(TrajOptSolverConfig):
) - 0.02
self._max_joint_acc = self.rollout_fn.state_bounds.acceleration[1, :] - 0.02
self._max_joint_jerk = self.rollout_fn.state_bounds.jerk[1, :] - 0.02
# self._max_joint_jerk = self._max_joint_jerk * 0.0 + 10.0
self._num_seeds = -1
self._col = None
if self.traj_evaluator is None:
@@ -844,9 +846,9 @@ class TrajOptSolver(TrajOptSolverConfig):
result.metrics = self.interpolate_rollout.get_metrics(interpolated_trajs)
st_time = time.time()
feasible = torch.all(result.metrics.feasible, dim=-1)
# if self.num_seeds == 1:
# print(result.metrics)
if result.metrics.position_error is not None:
converge = torch.logical_and(
result.metrics.position_error[..., -1] <= self.position_threshold,
@@ -874,6 +876,7 @@ class TrajOptSolver(TrajOptSolverConfig):
cspace_error=result.metrics.cspace_error,
optimized_dt=opt_dt,
raw_solution=result.action,
raw_action=result.raw_action,
)
else:
# get path length:
@@ -904,7 +907,6 @@ class TrajOptSolver(TrajOptSolverConfig):
convergence_error = result.metrics.cspace_error[..., -1]
else:
raise ValueError("convergence check requires either goal_pose or goal_state")
error = convergence_error + smooth_cost
error[~success] += 10000.0
if batch_mode:
@@ -919,6 +921,7 @@ class TrajOptSolver(TrajOptSolverConfig):
success = success[idx : idx + 1]
best_act_seq = result.action[idx]
best_raw_action = result.raw_action[idx]
interpolated_traj = interpolated_trajs[idx]
position_error = rotation_error = cspace_error = None
if result.metrics.position_error is not None:
@@ -961,6 +964,7 @@ class TrajOptSolver(TrajOptSolverConfig):
smooth_label=smooth_label,
optimized_dt=opt_dt,
raw_solution=best_act_seq,
raw_action=best_raw_action,
)
return traj_result
@@ -1044,7 +1048,7 @@ class TrajOptSolver(TrajOptSolverConfig):
if goal.goal_state is not None and self.use_cspace_seed:
# get linear seed
seed_traj = self.get_seeds(goal.current_state, goal.goal_state, num_seeds=num_seeds)
# .view(batch_size, self.num_seeds, self.traj_tsteps, self.dof)
# .view(batch_size, self.num_seeds, self.action_horizon, self.dof)
else:
# get start repeat seed:
log_info("No goal state found, using current config to seed")
@@ -1063,7 +1067,7 @@ class TrajOptSolver(TrajOptSolverConfig):
seed_traj = torch.cat((seed_traj, new_seeds), dim=0) # n_seed, batch, h, dof
seed_traj = seed_traj.view(
total_seeds, self.traj_tsteps, self.dof
total_seeds, self.action_horizon, self.dof
) # n_seeds,batch, h, dof
return seed_traj
@@ -1079,27 +1083,27 @@ class TrajOptSolver(TrajOptSolverConfig):
if n_seeds["linear"] > 0:
linear_seed = self.get_linear_seed(start_state, goal_state)
linear_seeds = linear_seed.view(1, -1, self.traj_tsteps, self.dof).repeat(
linear_seeds = linear_seed.view(1, -1, self.action_horizon, self.dof).repeat(
1, n_seeds["linear"], 1, 1
)
seed_set.append(linear_seeds)
if n_seeds["bias"] > 0:
bias_seed = self.get_bias_seed(start_state, goal_state)
bias_seeds = bias_seed.view(1, -1, self.traj_tsteps, self.dof).repeat(
bias_seeds = bias_seed.view(1, -1, self.action_horizon, self.dof).repeat(
1, n_seeds["bias"], 1, 1
)
seed_set.append(bias_seeds)
if n_seeds["start"] > 0:
bias_seed = self.get_start_seed(start_state)
bias_seeds = bias_seed.view(1, -1, self.traj_tsteps, self.dof).repeat(
bias_seeds = bias_seed.view(1, -1, self.action_horizon, self.dof).repeat(
1, n_seeds["start"], 1, 1
)
seed_set.append(bias_seeds)
if n_seeds["goal"] > 0:
bias_seed = self.get_start_seed(goal_state)
bias_seeds = bias_seed.view(1, -1, self.traj_tsteps, self.dof).repeat(
bias_seeds = bias_seed.view(1, -1, self.action_horizon, self.dof).repeat(
1, n_seeds["goal"], 1, 1
)
seed_set.append(bias_seeds)
@@ -1142,6 +1146,7 @@ class TrajOptSolver(TrajOptSolverConfig):
tensor_args=self.tensor_args,
out_traj_state=self._interpolated_traj_buffer,
min_dt=self.traj_evaluator_config.min_dt,
optimize_dt=self.optimize_dt,
)
return state, last_tstep, opt_dt

16
src/curobo/wrap/reacher/types.py
View File

@@ -63,6 +63,22 @@ class ReacherSolveState:
if self.num_seeds is None:
self.num_seeds = self.num_mpc_seeds
def clone(self):
return ReacherSolveState(
solve_type=self.solve_type,
n_envs=self.n_envs,
batch_size=self.batch_size,
n_goalset=self.n_goalset,
batch_env=self.batch_env,
batch_retract=self.batch_retract,
batch_mode=self.batch_mode,
num_seeds=self.num_seeds,
num_ik_seeds=self.num_ik_seeds,
num_graph_seeds=self.num_graph_seeds,
num_trajopt_seeds=self.num_trajopt_seeds,
num_mpc_seeds=self.num_mpc_seeds,
)
def get_batch_size(self):
return self.num_seeds * self.batch_size

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

@@ -42,6 +42,7 @@ class WrapResult:
metrics: Optional[RolloutMetrics] = None
debug: Any = None
js_action: Optional[State] = None
raw_action: Optional[torch.Tensor] = None
def clone(self):
return WrapResult(
@@ -155,6 +156,7 @@ class WrapBase(WrapConfig):
solve_time=self.opt_dt,
metrics=metrics,
debug={"steps": self.get_debug_data(), "cost": self.get_debug_cost()},
raw_action=act_seq,
)
return result