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

Add re-timing, minimum dt robustness

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

51
CHANGELOG.md
View File

@@ -10,6 +10,44 @@ its affiliates is strictly prohibited.
-->
# Changelog
## Version 0.7.2
### New Features
- Significant improvements for generating slow trajectories. Added re-timing post processing to
slow down optimized trajectories. Use `MotionGenPlanConfig.time_dilation_factor<1.0` to slow down a
planned trajectory. This is more robust than setting `velocity_scale<1.0` and also allows for
changing the speed of trajectories between planning calls
- `curobo.util.logger` adds `logger_name` as an input, enabling use of logging api with other
packages.
### Changes in default behavior
- Move `CudaRobotModelState` from `curobo.cuda_robot_model.types` to
`curobo.cuda_robot_model.cuda_robot_model`
- Activation distance for bound cost in now a ratio instead of absolute value to account for very
small range of joint limits when `velocity_scale<0.1`.
- `TrajResult` is renamed to `TrajOptResult` to be consistent with other solvers.
- Order of inputs to `get_batch_interpolated_trajectory` has changed.
- `MpcSolverConfig.load_from_robot_config` uses `world_model` instead of `world_cfg` to be
consistent with other wrappers.
### BugFixes & Misc.
- Fix bug in `MotionGen.plan_batch_env` where graph planner was being set to True. This also fixes
isaac sim example `batch_motion_gen_reacher.py`.
- Add `min_dt` as a parameter to `MotionGenConfig` and `TrajOptSolverConfig` to improve readability
and allow for having smaller `interpolation_dt`.
- Add `epsilon` to `min_dt` to make sure after time scaling, joint temporal values are not exactly
at their limits.
- Remove 0.02 offset for `max_joint_vel` and `max_joint_acc` in `TrajOptSolver`
- Bound cost now scales the cost by `1/limit_range**2` when `limit_range<1.0` to be robust to small
joint limits.
- Added documentation for `curobo.util.logger`, `curobo.wrap.reacher.motion_gen`,
`curobo.wrap.reacher.mpc`, and `curobo.wrap.reacher.trajopt`.
- When interpolation buffer is smaller than required, a new buffer is created with a warning
instead of raising an exception.
- `torch.cuda.synchronize()` now only synchronizes specified cuda device with
`torch.cuda.synchronize(device=self.tensor_args.device)`
- Added python example for MPC.
## Version 0.7.1
### New Features
@@ -30,7 +68,8 @@ limits are used. Use `TrajEvaluatorConfig.from_basic()` to initialize similar to
### BugFixes & Misc.
- Fix bug in `WorldVoxelCollision` where `env_query_idx` was being overwritten.
- Fix bug in `WorldVoxelCollision` where parent collision types were not getting called in some cases.
- Fix bug in `WorldVoxelCollision` where parent collision types were not getting called in some
cases.
- Change voxelization dimensions to include 1 extra voxel per dim.
- Added `seed` parameter to `IKSolverConfig`.
- Added `sampler_seed` parameter `RolloutConfig`.
@@ -86,8 +125,8 @@ parallel problems in optimization.
- Add locked joint state update api in MotionGen class.
- Add goalset warmup padding to handle varied number of goals during goalset planning and also when
calling plan_single after warmup of goalset.
- Add new trajopt config to allow for smooth solutions at slow speeds (`velocity_scale<=0.25`). Also
add error when `velocity_scale<0.1`.
- Add new trajopt config to allow for smooth solutions at slow speeds (`velocity_scale<=0.25`).
Also add error when `velocity_scale<0.1`.
- Add experimental robot image segmentation module to enable robot removal in depth images.
- Add constrained planning mode to motion_gen.
- Use `torch.compile` to leverage better kernel fusion in place of `torch.jit.script`.
@@ -97,7 +136,8 @@ now only 2x slower than cuboid (from 5x slower). Optimization convergence is als
- 2x faster LBFGS kernel that allocates upto 68kb of shared memory, preventing use in CUDA devices
with compute capability ``<7.0``.
- On benchmarking Dataset, Planning time is now 42ms on average from 50ms. Higher quality solutions
are also obtained. See [benchmarks](https://curobo.org/source/getting_started/4_benchmarks.html) for more details.
are also obtained. See [benchmarks](https://curobo.org/source/getting_started/4_benchmarks.html)
for more details.
- Add ``WorldCollisionVoxel``, a new collision checking implementation that uses a voxel grid
of signed distances (SDF) to compute collision avoidance metrics. Documentation coming soon, see
``benchmark/curobo_voxel_benchmark.py`` for an example.
@@ -172,7 +212,8 @@ external directory.
### BugFixes & Misc.
- Update nvblox wrappers to work with v0.0.5 without segfaults. Significantly improves stability.
- Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and 2022.2.1
- Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and
2022.2.1
- Cleanup docker scripts. Use `build_docker.sh` instead of `build_dev_docker.sh`. Added isaac sim
development docker.
- Fixed bug in backward kinematics kernel, helped improve IK and TO pose reaching accuracy..

4
benchmark/curobo_benchmark.py
View File

@@ -209,7 +209,6 @@ def load_curobo(
collision_activation_distance=collision_activation_distance,
trajopt_dt=0.25,
finetune_dt_scale=finetune_dt_scale,
maximum_trajectory_dt=0.15,
high_precision=args.high_precision,
)
mg = MotionGen(motion_gen_config)
@@ -296,7 +295,6 @@ def benchmark_mb(
enable_graph_attempt=1,
disable_graph_attempt=10,
enable_finetune_trajopt=not args.no_finetune,
partial_ik_opt=False,
enable_graph=graph_mode or force_graph,
timeout=60,
enable_opt=not graph_mode,
@@ -572,6 +570,7 @@ def benchmark_mb(
if not args.kpi:
try:
# Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]
@@ -604,6 +603,7 @@ def benchmark_mb(
g_m = CuroboGroupMetrics.from_list(all_files)
try:
# Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]

225
benchmark/curobo_nvblox_benchmark.py
View File

@@ -11,6 +11,7 @@
# Standard Library
import argparse
import time
from copy import deepcopy
from typing import Optional
@@ -24,7 +25,8 @@ from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
from tqdm import tqdm
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType, WorldConfig
from curobo.geom.sdf.world import CollisionCheckerType, WorldCollisionConfig, WorldConfig
from curobo.geom.sdf.world_blox import WorldBloxCollision
from curobo.geom.types import Cuboid as curobo_Cuboid
from curobo.geom.types import Mesh
from curobo.types.base import TensorDeviceType
@@ -35,6 +37,7 @@ from curobo.types.state import JointState
from curobo.util.logger import setup_curobo_logger
from curobo.util.metrics import CuroboGroupMetrics, CuroboMetrics
from curobo.util_file import (
file_exists,
get_assets_path,
get_robot_configs_path,
get_world_configs_path,
@@ -130,8 +133,10 @@ def load_curobo(
):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = 0.0
robot_cfg["kinematics"]["collision_link_names"].remove("attached_object")
robot_cfg["kinematics"]["ee_link"] = "panda_hand"
ik_seeds = 30
ik_seeds = 32
if graph_mode:
trajopt_seeds = 4
if trajopt_seeds >= 14:
@@ -147,11 +152,30 @@ def load_curobo(
"world": {
"pose": [0, 0, 0, 1, 0, 0, 0],
"integrator_type": "tsdf",
"voxel_size": 0.02,
"voxel_size": 0.01,
}
}
}
)
world_nvblox_config = WorldCollisionConfig.load_from_dict(
{"cache": None, "checker_type": "BLOX"},
world_cfg,
TensorDeviceType(),
)
world_nvblox = WorldBloxCollision(world_nvblox_config)
world_cfg = WorldConfig.from_dict(
{
"voxel": {
"base": {
"dims": [2.4, 2.4, 2.4],
"pose": [0, 0, 0, 1, 0, 0, 0],
"voxel_size": 0.005,
"feature_dtype": torch.float8_e4m3fn,
},
}
}
)
interpolation_steps = 2000
if graph_mode:
interpolation_steps = 100
@@ -164,7 +188,7 @@ def load_curobo(
robot_cfg_instance,
world_cfg,
trajopt_tsteps=tsteps,
collision_checker_type=CollisionCheckerType.BLOX,
collision_checker_type=CollisionCheckerType.VOXEL,
use_cuda_graph=cuda_graph,
position_threshold=0.005, # 0.5 cm
rotation_threshold=0.05,
@@ -177,7 +201,6 @@ def load_curobo(
interpolation_steps=interpolation_steps,
collision_activation_distance=collision_activation_distance,
trajopt_dt=0.25,
finetune_dt_scale=0.9,
maximum_trajectory_dt=0.15,
finetune_trajopt_iters=200,
)
@@ -188,12 +211,15 @@ def load_curobo(
robot_cfg_instance,
"collision_table.yml",
collision_activation_distance=0.0,
collision_checker_type=CollisionCheckerType.PRIMITIVE,
collision_checker_type=CollisionCheckerType.MESH,
n_cuboids=50,
n_meshes=50,
)
mg.clear_world_cache()
robot_world = RobotWorld(config)
robot_world.clear_world_cache()
return mg, robot_cfg, robot_world
return mg, robot_cfg, robot_world, world_nvblox
def benchmark_mb(
@@ -208,7 +234,7 @@ def benchmark_mb(
# load dataset:
graph_mode = args.graph
interpolation_dt = 0.02
file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:2]
file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][2:]
enable_debug = save_log or plot_cost
all_files = []
@@ -216,8 +242,9 @@ def benchmark_mb(
if override_tsteps is not None:
og_tsteps = override_tsteps
og_trajopt_seeds = 12
og_collision_activation_distance = 0.03 # 0.03
og_trajopt_seeds = 4
og_collision_activation_distance = 0.025
count = [0, 0, 0, 0]
if args.graph:
og_trajopt_seeds = 4
for file_path in file_paths:
@@ -228,6 +255,7 @@ def benchmark_mb(
for key, v in tqdm(problems.items()):
scene_problems = problems[key]
m_list = []
count[3] += 1
i = -1
ik_fail = 0
trajopt_seeds = og_trajopt_seeds
@@ -236,32 +264,7 @@ def benchmark_mb(
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
if "cage_panda" in key:
trajopt_seeds = 8
else:
continue
if "table_under_pick_panda" in key:
tsteps = 44
trajopt_seeds = 16
finetune_dt_scale = 0.98
if "cubby_task_oriented" in key and "merged" not in key:
trajopt_seeds = 24
finetune_dt_scale = 0.95
collision_activation_distance = 0.035
if "dresser_task_oriented" in key:
trajopt_seeds = 24
finetune_dt_scale = 0.95
collision_activation_distance = 0.035 # 0.035
if "merged_cubby_task_oriented" in key:
collision_activation_distance = 0.025 # 0.035
if "tabletop_task_oriented" in key:
collision_activation_distance = 0.025 # 0.035
if key in ["dresser_neutral_goal"]:
trajopt_seeds = 24
collision_activation_distance = og_collision_activation_distance
mg, robot_cfg, robot_world = load_curobo(
mg, robot_cfg, robot_world, world_nvblox = load_curobo(
1,
enable_debug,
tsteps,
@@ -280,7 +283,6 @@ def benchmark_mb(
max_attempts=10,
enable_graph_attempt=1,
enable_finetune_trajopt=True,
partial_ik_opt=False,
enable_graph=graph_mode,
timeout=60,
enable_opt=not graph_mode,
@@ -298,14 +300,60 @@ def benchmark_mb(
mg.reset(reset_seed=False)
world = WorldConfig.from_dict(problem["obstacles"])
mg.world_coll_checker.clear_cache()
mg.world_coll_checker.update_blox_hashes()
world_nvblox.clear_cache()
world_nvblox.update_blox_hashes()
mg.clear_world_cache()
save_path = "benchmark/log/nvblox/" + key + "_" + str(i)
save_path = "benchmark/log/nvblox_640_new/" + key + "_" + str(i)
save_path = (
"/home/bala/code/raven_internship/data/render_mpinets_640/reformat/"
+ key
+ "_"
+ str(i)
)
# save_path = "/home/bala/code/raven_internship/data/render_26k/_0_8/reformat/" + key + "_" + str(i)
if not file_exists(save_path):
continue
m_dataset = Sun3dDataset(save_path)
tensor_args = mg.tensor_args
if i == 0:
for j in tqdm(range(min(10, len(m_dataset))), leave=False):
data = m_dataset[j]
cam_obs = CameraObservation(
rgb_image=tensor_args.to_device(data["rgba"])
.squeeze()
.to(dtype=torch.uint8)
.permute(1, 2, 0), # data[rgba]: 4 x H x W -> H x W x 4
depth_image=tensor_args.to_device(data["depth"]),
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=mg.tensor_args.device)),
)
cam_obs = cam_obs
torch.cuda.synchronize()
st_int_time = time.time()
world_nvblox.add_camera_frame(cam_obs, "world")
torch.cuda.synchronize()
world_nvblox.process_camera_frames("world", False)
torch.cuda.synchronize()
world_nvblox.update_blox_hashes()
# get esdf grid:
esdf = world_nvblox.get_esdf_in_bounding_box(
curobo_Cuboid(
name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2.4, 2.4, 2.4]
),
voxel_size=0.005,
dtype=torch.float32,
)
mg.world_coll_checker.update_voxel_data(esdf)
world_nvblox.clear_cache()
world_nvblox.update_blox_hashes()
mg.clear_world_cache()
int_time = 0
for j in tqdm(range(min(1, len(m_dataset))), leave=False):
data = m_dataset[j]
cam_obs = CameraObservation(
@@ -318,18 +366,26 @@ def benchmark_mb(
pose=Pose.from_matrix(data["pose"].to(device=mg.tensor_args.device)),
)
cam_obs = cam_obs
mg.add_camera_frame(cam_obs, "world")
mg.process_camera_frames("world", False)
torch.cuda.synchronize()
mg.world_coll_checker.update_blox_hashes()
st_int_time = time.time()
world_nvblox.add_camera_frame(cam_obs, "world")
torch.cuda.synchronize()
# if i > 2:
# mg.world_coll_checker.clear_cache()
# mg.world_coll_checker.update_blox_hashes()
int_time += time.time() - st_int_time
st_time = time.time()
world_nvblox.process_camera_frames("world", False)
torch.cuda.synchronize()
world_nvblox.update_blox_hashes()
# mg.world_coll_checker.save_layer("world", "test.nvblx")
# get esdf grid:
esdf = world_nvblox.get_esdf_in_bounding_box(
curobo_Cuboid(name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2.4, 2.4, 2.4]),
voxel_size=0.005,
dtype=torch.float32,
)
mg.world_coll_checker.update_voxel_data(esdf)
torch.cuda.synchronize()
perception_time = time.time() - st_time + int_time
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single(
@@ -372,6 +428,31 @@ def benchmark_mb(
),
log_scale=False,
)
if save_log or write_usd:
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
# nvblox_obs = world_nvblox.get_mesh_from_blox_layer(
# "world",
# )
# nvblox_obs.vertex_colors = None
# if nvblox_obs.vertex_colors is not None:
# nvblox_obs.vertex_colors = nvblox_obs.vertex_colors.cpu().numpy()
# else:
# nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
# nvblox_obs.name = "nvblox_mesh_world"
# world.add_obstacle(nvblox_obs)
coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
curobo_Cuboid(name="new_test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2, 2, 2]),
voxel_size=0.01,
)
coll_mesh.color = [0.0, 0.8, 0.8, 0.8]
coll_mesh.name = "nvblox_voxel_world"
world.add_obstacle(coll_mesh)
if result.success.item():
q_traj = result.get_interpolated_plan()
problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist()
@@ -425,10 +506,10 @@ def benchmark_mb(
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
world_coll = WorldConfig.from_dict(problem["obstacles"]).get_obb_world()
world_coll = WorldConfig.from_dict(problem["obstacles"]).get_mesh_world()
# check if path is collision free w.r.t. ground truth mesh:
# robot_world.world_model.clear_cache()
robot_world.world_model.clear_cache()
robot_world.update_world(world_coll)
q_int_traj = result.get_interpolated_plan().position.unsqueeze(0)
@@ -448,6 +529,9 @@ def benchmark_mb(
# if not d_mask:
# write_usd = True
# #print(torch.max(d_world).item(), problem_name)
if d_int_mask:
count[0] += 1
current_metrics = CuroboMetrics(
skip=False,
success=True,
@@ -465,37 +549,11 @@ def benchmark_mb(
motion_time=result.motion_time.item(),
solve_time=result.solve_time,
jerk=torch.max(torch.abs(result.optimized_plan.jerk)).item(),
perception_time=perception_time,
)
if save_log or write_usd:
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
nvblox_obs = mg.world_coll_checker.get_mesh_from_blox_layer(
"world",
)
# nvblox_obs.vertex_colors = None
if nvblox_obs.vertex_colors is not None:
nvblox_obs.vertex_colors = nvblox_obs.vertex_colors.cpu().numpy()
else:
nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
nvblox_obs.name = "nvblox_mesh_world"
world.add_obstacle(nvblox_obs)
coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
curobo_Cuboid(
name="test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.5, 1.5, 1]
),
voxel_size=0.005,
)
coll_mesh.color = [0.0, 0.8, 0.8, 0.8]
coll_mesh.name = "nvblox_voxel_world"
world.add_obstacle(coll_mesh)
# run planner
if write_usd:
if write_usd and not d_mask:
# CuRobo
from curobo.util.usd_helper import UsdHelper
@@ -512,7 +570,7 @@ def benchmark_mb(
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
# flatten_usd=True,
flatten_usd=True,
)
# write_usd = False
# exit()
@@ -537,7 +595,9 @@ def benchmark_mb(
m_list.append(current_metrics)
all_groups.append(current_metrics)
count[1] += 1
elif result.valid_query:
count[1] += 1
current_metrics = CuroboMetrics()
debug = {
"used_graph": result.used_graph,
@@ -555,6 +615,7 @@ def benchmark_mb(
m_list.append(current_metrics)
all_groups.append(current_metrics)
else:
count[2] += 1
# print("invalid: " + problem_name)
debug = {
"used_graph": result.used_graph,
@@ -583,7 +644,7 @@ def benchmark_mb(
write_trajopt=True,
visualize_robot_spheres=True,
grid_space=2,
# flatten_usd=True,
flatten_usd=True,
)
exit()
g_m = CuroboGroupMetrics.from_list(m_list)
@@ -599,6 +660,7 @@ def benchmark_mb(
g_m.cspace_path_length.percent_98,
g_m.motion_time.percent_98,
g_m.perception_interpolated_success,
g_m.perception_time.mean,
# g_m.orientation_error.median,
)
print(g_m.attempts)
@@ -633,6 +695,7 @@ def benchmark_mb(
print("ST: ", g_m.solve_time)
print("accuracy: ", g_m.position_error, g_m.orientation_error)
print("Jerk: ", g_m.jerk)
print(count)
if __name__ == "__main__":

5
benchmark/curobo_voxel_benchmark.py
View File

@@ -185,7 +185,7 @@ def load_curobo(
finetune_trajopt_iters=200,
)
mg = MotionGen(motion_gen_config)
mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=True)
mg.warmup(enable_graph=True, warmup_js_trajopt=False)
# create a ground truth collision checker:
world_model = WorldConfig.from_dict(
{
@@ -305,6 +305,7 @@ def benchmark_mb(
voxel_size=0.005,
dtype=torch.float32,
)
# esdf.feature_tensor[esdf.feature_tensor < -1.0] = -1000.0
world_voxel_collision = mg.world_coll_checker
world_voxel_collision.update_voxel_data(esdf)
@@ -578,6 +579,7 @@ def benchmark_mb(
print(g_m.attempts)
g_m = CuroboGroupMetrics.from_list(all_groups)
try:
# Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]
@@ -611,6 +613,7 @@ def benchmark_mb(
g_m = CuroboGroupMetrics.from_list(all_files)
print("######## FULL SET ############")
try:
# Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]

4
benchmark/generate_nvblox_images.py
View File

@@ -36,7 +36,7 @@ np.random.seed(0)
def generate_images():
# load dataset:
file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:]
file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:2]
for file_path in file_paths:
problems = file_path()
@@ -57,7 +57,7 @@ def generate_images():
# generate images and write to disk:
MeshDataset(
None, n_frames=50, image_size=640, save_data_dir=save_path, trimesh_mesh=mesh
None, n_frames=1, image_size=640, save_data_dir=save_path, trimesh_mesh=mesh
)

1
benchmark/ik_benchmark.py
View File

@@ -185,6 +185,7 @@ if __name__ == "__main__":
print("Reported errors are 98th percentile")
df.to_csv(join_path(args.save_path, args.file_name + ".csv"))
try:
# Third Party
from tabulate import tabulate
print(tabulate(df, headers="keys", tablefmt="grid"))

2
docker/user.dockerfile
View File

@@ -15,7 +15,7 @@ FROM curobo_docker:${IMAGE_TAG}
# Set variables
ARG USERNAME
ARG USER_ID
ARG CACHE_DATE=2023-04-11
ARG CACHE_DATE=2024-04-25
# Set environment variables

5
docker/x86.dockerfile
View File

@@ -78,7 +78,7 @@ ENV TORCH_CUDA_ARCH_LIST "7.0+PTX"
ENV LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}"
# Add cache date to avoid using cached layers older than this
ARG CACHE_DATE=2024-04-11
ARG CACHE_DATE=2024-04-25
RUN pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"
@@ -125,3 +125,6 @@ RUN python -m pip install pyrealsense2 opencv-python transforms3d
# install benchmarks:
RUN python -m pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"
# update ucx path: https://github.com/openucx/ucc/issues/476
RUN export LD_LIBRARY_PATH=/opt/hpcx/ucx/lib:$LD_LIBRARY_PATH

13
examples/isaac_sim/batch_motion_gen_reacher.py
View File

@@ -144,17 +144,12 @@ def main():
robot_cfg,
world_cfg_list,
tensor_args,
trajopt_tsteps=32,
collision_checker_type=CollisionCheckerType.MESH,
use_cuda_graph=True,
num_trajopt_seeds=12,
num_graph_seeds=12,
interpolation_dt=0.03,
collision_cache={"obb": 6, "mesh": 6},
collision_cache={"obb": 10, "mesh": 10},
collision_activation_distance=0.025,
self_collision_check=True,
maximum_trajectory_dt=0.25,
fixed_iters_trajopt=True,
)
motion_gen = MotionGen(motion_gen_config)
j_names = robot_cfg["kinematics"]["cspace"]["joint_names"]
@@ -162,7 +157,9 @@ def main():
print("warming up...")
motion_gen.warmup(
enable_graph=False, warmup_js_trajopt=False, batch=n_envs, batch_env_mode=True
batch=n_envs,
batch_env_mode=True,
warmup_js_trajopt=False,
)
add_extensions(simulation_app, args.headless_mode)
@@ -176,7 +173,7 @@ def main():
x_sph[..., 3] = radius
env_query_idx = torch.arange(n_envs, device=tensor_args.device, dtype=torch.int32)
plan_config = MotionGenPlanConfig(
enable_graph=False, enable_graph_attempt=4, max_attempts=2, enable_finetune_trajopt=True
enable_graph=False, max_attempts=2, enable_finetune_trajopt=True
)
prev_goal = None
cmd_plan = [None, None]

64
examples/motion_gen_example.py
View File

@@ -25,7 +25,7 @@ from curobo.util_file import get_robot_configs_path, get_world_configs_path, joi
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
def plot_traj(trajectory, dt):
def plot_traj(trajectory, dt, file_name="test.png"):
# Third Party
import matplotlib.pyplot as plt
@@ -42,7 +42,7 @@ def plot_traj(trajectory, dt):
axs[3].plot(timesteps, qddd[:, i], label=str(i))
plt.legend()
plt.savefig("test.png")
plt.savefig(file_name)
plt.close()
# plt.show()
@@ -91,6 +91,53 @@ def plot_iters_traj_3d(trajectory, d_id=1, dof=7, seed=0):
plt.show()
def demo_motion_gen_simple():
world_config = {
"mesh": {
"base_scene": {
"pose": [10.5, 0.080, 1.6, 0.043, -0.471, 0.284, 0.834],
"file_path": "scene/nvblox/srl_ur10_bins.obj",
},
},
"cuboid": {
"table": {
"dims": [5.0, 5.0, 0.2], # x, y, z
"pose": [0.0, 0.0, -0.1, 1, 0, 0, 0.0], # x, y, z, qw, qx, qy, qz
},
},
}
motion_gen_config = MotionGenConfig.load_from_robot_config(
"ur5e.yml",
world_config,
interpolation_dt=0.01,
)
motion_gen = MotionGen(motion_gen_config)
motion_gen.warmup()
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.rollout_fn.compute_kinematics(
JointState.from_position(retract_cfg.view(1, -1))
)
goal_pose = Pose.from_list([-0.4, 0.0, 0.4, 1.0, 0.0, 0.0, 0.0]) # x, y, z, qw, qx, qy, qz
start_state = JointState.from_position(
torch.zeros(1, 6).cuda(),
joint_names=[
"shoulder_pan_joint",
"shoulder_lift_joint",
"elbow_joint",
"wrist_1_joint",
"wrist_2_joint",
"wrist_3_joint",
],
)
result = motion_gen.plan_single(start_state, goal_pose, MotionGenPlanConfig(max_attempts=1))
traj = result.get_interpolated_plan() # result.optimized_dt has the dt between timesteps
print("Trajectory Generated: ", result.success)
def demo_motion_gen_mesh():
PLOT = False
tensor_args = TensorDeviceType()
@@ -149,7 +196,7 @@ def demo_motion_gen(js=False):
)
motion_gen = MotionGen(motion_gen_config)
motion_gen.warmup(parallel_finetune=True)
motion_gen.warmup()
# motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js, parallel_finetune=True)
# robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
@@ -169,16 +216,21 @@ def demo_motion_gen(js=False):
result = motion_gen.plan_single_js(
start_state,
goal_state,
MotionGenPlanConfig(max_attempts=1, parallel_finetune=True),
MotionGenPlanConfig(max_attempts=1, time_dilation_factor=0.5),
)
else:
result = motion_gen.plan_single(
start_state,
retract_pose,
MotionGenPlanConfig(
max_attempts=1, parallel_finetune=True, enable_finetune_trajopt=True
max_attempts=1,
timeout=5,
time_dilation_factor=0.5,
),
)
new_result = result.clone()
new_result.retime_trajectory(0.5, create_interpolation_buffer=True)
print(new_result.optimized_dt, new_result.motion_time, result.motion_time)
print(
"Trajectory Generated: ",
result.success,
@@ -190,6 +242,7 @@ def demo_motion_gen(js=False):
traj = result.get_interpolated_plan()
plot_traj(traj, result.interpolation_dt)
plot_traj(new_result.get_interpolated_plan(), new_result.interpolation_dt, "test_slow.png")
# plt.save("test.png")
# plt.close()
# traj = result.debug_info["opt_solution"].position.view(-1,7)
@@ -431,6 +484,7 @@ def demo_motion_gen_batch_env(n_envs: int = 10):
if __name__ == "__main__":
setup_curobo_logger("error")
demo_motion_gen(js=False)
# demo_motion_gen_simple()
# demo_motion_gen_batch()
# demo_motion_gen_goalset()
# n = [2, 10]

2
setup.cfg
View File

@@ -120,7 +120,7 @@ doc =
# NOTE (roflaherty): Flake8 doesn't support pyproject.toml configuration yet.
[flake8]
max-line-length = 100
max-line-length = 99
docstring-convention = google
exclude = .git,build,deprecated,dist,venv
ignore =

6
src/curobo/__init__.py
View File

@@ -18,9 +18,10 @@ planning.
High-level APIs:
- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`.
- Inverse Kinematics: :mod:`curobo.wrap.reacher.ik_solver`.
- Model Predictive Control: :mod:`curobo.wrap.reacher.mpc`.
- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`.
- Trajectory Optimization: :mod:`curobo.wrap.reacher.trajopt`.
cuRobo package is split into several modules:
@@ -36,7 +37,8 @@ cuRobo package is split into several modules:
- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of
collision-free reacher and batched robot world collision checking.
- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
:meth:`types.state.JointState`, :meth:`types.camera.CameraObservation`, :meth:`types.math.Pose`.
:py:meth:`~types.state.JointState`, :py:meth:`~types.camera.CameraObservation`,
:py:meth:`~types.math.Pose`.
"""

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

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

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

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

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

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

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

@@ -19,9 +19,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
@@ -59,7 +59,7 @@ cost:
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk
activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -79,17 +79,17 @@ cost:
lbfgs:
n_iters: 300
n_iters: 400
inner_iters: 25
cold_start_n_iters: null
min_iters: 25
line_search_scale: [0.1, 0.3, 0.7,1.0] # #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 1.0
cost_relative_threshold: 0.999 #0.999
cost_convergence: 0.001
cost_delta_threshold: 0.0001
cost_relative_threshold: 0.99
epsilon: 0.01
history: 27 #15
history: 27
use_cuda_graph: True
n_problems: 1
store_debug: False

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

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

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

@@ -40,21 +40,19 @@ cost:
terminal: False
use_metric: True
run_weight: 1.0
cspace_cfg:
weight: 0.000
bound_cfg:
weight: 500.0
activation_distance: [0.1]
null_space_weight: [1.0]
weight: 5000.0
activation_distance: [0.001]
null_space_weight: [0.1]
primitive_collision_cfg:
weight: 5000.0
use_sweep: False
classify: False
activation_distance: 0.01
self_collision_cfg:
weight: 500.0
weight: 5000.0
classify: False

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

@@ -19,9 +19,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
@@ -48,17 +48,17 @@ cost:
use_metric: True
cspace_cfg:
weight: 1000.0
weight: 10000.0
terminal: True
run_weight: 0.0
bound_cfg:
weight: [50000.0, 50000.0, 50000.0,50000.0]
weight: [50000.0, 50000.0, 500.0,500.0]
smooth_weight: [0.0,10000.0,5.0, 0.0]
run_weight_velocity: 0.00
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk
activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:

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

@@ -18,9 +18,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_dt: 0.15
base_ratio: 1.0
max_dt: 0.2
max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
teleport_mode: False

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

File diff suppressed because it is too large Load Diff

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

@@ -8,7 +8,31 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
"""
This module contains :meth:`MpcSolver` that provides a high-level interface to for model
predictive control (MPC) for reaching Cartesian poses and also joint configurations while
avoiding obstacles. The solver uses Model Predictive Path Integral (MPPI) optimization as the
solver. MPC only optimizes locally so the robot can get stuck near joint limits or behind
obstacles. To generate global trajectories, use
:py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`.
A python example is available at :ref:`python_mpc_example`.
.. note::
Gradient-based MPC is also implemented with L-BFGS but is highly experimental and not
recommended for real robots.
.. raw:: html
<p>
<video autoplay="True" loop="True" muted="True" preload="auto" width="100%"><source src="../videos/mpc_clip.mp4" type="video/mp4"></video>
</p>
"""
# Standard Library
import time
@@ -19,6 +43,7 @@ from typing import Dict, Optional, Union
import torch
# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.geom.sdf.utils import create_collision_checker
from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig
from curobo.geom.types import WorldConfig
@@ -44,22 +69,31 @@ from curobo.wrap.wrap_mpc import WrapConfig, WrapMpc
@dataclass
class MpcSolverConfig:
"""Configuration dataclass for MPC."""
#: MPC Solver.
solver: WrapMpc
#: World Collision Checker.
world_coll_checker: Optional[WorldCollision] = None
#: Numeric precision and device to run computations.
tensor_args: TensorDeviceType = TensorDeviceType()
#: Capture full step in MPC as a single CUDA graph. This is not supported currently.
use_cuda_graph_full_step: bool = False
@staticmethod
def load_from_robot_config(
robot_cfg: Union[Union[str, dict], RobotConfig],
world_cfg: Union[Union[str, dict], WorldConfig],
world_model: Union[Union[str, dict], WorldConfig],
base_cfg: Optional[dict] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
compute_metrics: bool = True,
use_cuda_graph: Optional[bool] = None,
particle_opt_iters: Optional[int] = None,
self_collision_check: bool = True,
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE,
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH,
use_es: Optional[bool] = None,
es_learning_rate: Optional[float] = 0.01,
use_cuda_graph_metrics: bool = False,
@@ -74,9 +108,56 @@ class MpcSolverConfig:
use_lbfgs: bool = False,
use_mppi: bool = True,
):
"""Create an MPC solver configuration from robot and world configuration.
Args:
robot_cfg: Robot configuration. Can be a path to a YAML file or a dictionary or
an instance of :class:`~curobo.types.robot.RobotConfig`.
world_model: World configuration. Can be a path to a YAML file or a dictionary or
an instance of :class:`~curobo.geom.types.WorldConfig`.
base_cfg: Base configuration for the solver. This file is used to check constraints
and convergence. If None, the default configuration from ``base_cfg.yml`` is used.
tensor_args: Numeric precision and device to run computations.
compute_metrics: Compute metrics on MPC step.
use_cuda_graph: Use CUDA graph for the optimization step.
particle_opt_iters: Number of iterations for the particle optimization.
self_collision_check: Enable self-collision check during MPC optimization.
collision_checker_type: Type of collision checker to use. See :ref:`world_collision`.
use_es: Use Evolution Strategies (ES) solver for MPC. Highly experimental.
es_learning_rate: Learning rate for ES solver.
use_cuda_graph_metrics: Use CUDA graph for computing metrics.
store_rollouts: Store rollouts information for debugging. This will also store the
trajectory taken by the end-effector across the horizon.
use_cuda_graph_full_step: Capture full step in MPC as a single CUDA graph. This is
experimental and might not work reliably.
sync_cuda_time: Synchronize CUDA device with host using
:py:func:`torch.cuda.synchronize` before calculating compute time.
collision_cache: Cache of obstacles to create to load obstacles between planning calls.
An example: ``{"obb": 10, "mesh": 10}``, to create a cache of 10 cuboids and 10
meshes.
n_collision_envs: Number of collision environments to create for batched planning
across different environments. Only used for :py:meth:`MpcSolver.setup_solve_batch_env`
and :py:meth:`MpcSolver.setup_solve_batch_env_goalset`.
collision_activation_distance: Distance in meters to activate collision cost. A good
value to start with is 0.01 meters. Increase the distance if the robot needs to
stay further away from obstacles.
world_coll_checker: Instance of world collision checker to use for MPC. Leaving this to
None will create a new instance of world collision checker using the provided
:attr:`world_model`.
step_dt: Time step to use between each step in the trajectory. If None, the default
time step from the configuration~(`particle_mpc.yml` or `gradient_mpc.yml`)
is used. This dt should match the control frequency at which you are sending
commands to the robot. This dt should also be greater than than the compute
time for a single step.
use_lbfgs: Use L-BFGS solver for MPC. Highly experimental.
use_mppi: Use MPPI solver for MPC.
Returns:
MpcSolverConfig: Configuration for the MPC solver.
"""
if use_cuda_graph_full_step:
log_error("use_cuda_graph_full_step currently is not supported")
raise ValueError("use_cuda_graph_full_step currently is not supported")
task_file = "particle_mpc.yml"
config_data = load_yaml(join_path(get_task_configs_path(), task_file))
@@ -108,14 +189,14 @@ class MpcSolverConfig:
if isinstance(robot_cfg, dict):
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
if isinstance(world_cfg, str):
world_cfg = load_yaml(join_path(get_world_configs_path(), world_cfg))
if isinstance(world_model, str):
world_model = load_yaml(join_path(get_world_configs_path(), world_model))
if world_coll_checker is None and world_cfg is not None:
world_cfg = WorldCollisionConfig.load_from_dict(
base_cfg["world_collision_checker_cfg"], world_cfg, tensor_args
if world_coll_checker is None and world_model is not None:
world_model = WorldCollisionConfig.load_from_dict(
base_cfg["world_collision_checker_cfg"], world_model, tensor_args
)
world_coll_checker = create_collision_checker(world_cfg)
world_coll_checker = create_collision_checker(world_model)
grad_config_data = None
if use_lbfgs:
grad_config_data = load_yaml(join_path(get_task_configs_path(), "gradient_mpc.yml"))
@@ -134,7 +215,7 @@ class MpcSolverConfig:
base_cfg["constraint"],
base_cfg["convergence"],
base_cfg["world_collision_checker_cfg"],
world_cfg,
world_model,
world_coll_checker=world_coll_checker,
tensor_args=tensor_args,
)
@@ -172,7 +253,7 @@ class MpcSolverConfig:
base_cfg["constraint"],
base_cfg["convergence"],
base_cfg["world_collision_checker_cfg"],
world_cfg,
world_model,
world_coll_checker=world_coll_checker,
tensor_args=tensor_args,
)
@@ -201,13 +282,42 @@ class MpcSolverConfig:
class MpcSolver(MpcSolverConfig):
"""Model Predictive Control Solver for Arm Reacher task.
"""High-level interface for Model Predictive Control (MPC).
Args:
MpcSolverConfig: _description_
MPC can reach Cartesian poses and joint configurations while avoiding obstacles. The solver
uses Model Predictive Path Integral (MPPI) optimization as the solver. MPC only optimizes
locally so the robot can get stuck near joint limits or behind obstacles. To generate global
trajectories, use :py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`.
See :ref:`python_mpc_example` for an example. This MPC solver implementation can be used in the
following steps:
1. Create a :py:class:`~curobo.rollout.rollout_base.Goal` object with the target pose or joint
configuration.
2. Create a goal buffer for the problem type using :meth:`setup_solve_single`,
:meth:`setup_solve_goalset`, :meth:`setup_solve_batch`, :meth:`setup_solve_batch_goalset`,
:meth:`setup_solve_batch_env`, or :meth:`setup_solve_batch_env_goalset`. Pass the goal
object from the previous step to this function. This function will update the internal
solve state of MPC and also the goal for MPC. An augmented goal buffer is returned.
3. Call :meth:`step` with the current joint state to get the next action.
4. To change the goal, create a :py:class:`~curobo.types.math.Pose` object with new pose or
:py:class:`~curobo.types.state.JointState` object with new joint configuration. Then
copy the target into the augmented goal buffer using
``goal_buffer.goal_pose.copy_(new_pose)`` or ``goal_buffer.goal_state.copy_(new_state)``.
5. Call :meth:`update_goal` with the augmented goal buffer to update the goal for MPC.
6. Call :meth:`step` with the current joint state to get the next action.
To dynamically change the type of goal reached between pose and joint configuration targets,
create the goal object in step 1 with both targets and then use :meth:`enable_cspace_cost` and
:meth:`enable_pose_cost` to enable or disable reaching joint configuration cost and pose cost.
"""
def __init__(self, config: MpcSolverConfig) -> None:
"""Initializes the MPC solver.
Args:
config: Configuration parameters for MPC.
"""
super().__init__(**vars(config))
self.tensor_args = self.solver.rollout_fn.tensor_args
self._goal_buffer = Goal()
@@ -222,15 +332,326 @@ class MpcSolver(MpcSolverConfig):
self._cu_step_graph = None
self._cu_result = None
def _update_batch_size(self, batch_size):
if self.batch_size != batch_size:
self.batch_size = batch_size
def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a robot to reach target pose or joint configuration.
def update_goal_buffer(
Args:
goal: goal object containing target pose or joint configuration.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a robot to reach a pose in a set of poses.
Args:
goal: goal object containing target goalset or joint configuration.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.GOALSET,
num_mpc_seeds=num_seeds,
batch_size=1,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots to reach targets.
Args:
goal: goal object containing target poses or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots to reach a set of poses.
Args:
goal: goal object containing target goalset or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots in different collision worlds.
Args:
goal: goal object containing target poses or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
"""Creates a goal buffer to solve for a batch of robots in different collision worlds.
Args:
goal: goal object containing target goalset or joint configurations.
num_seeds: Number of seeds to use in the solver.
Returns:
Goal: Instance of augmented goal buffer.
"""
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
self.update_goal(goal_buffer)
return goal_buffer
def step(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
max_attempts: int = 1,
):
"""Solve for the next action given the current state.
Args:
current_state: Current joint state of the robot.
shift_steps: Number of steps to shift the trajectory.
seed_traj: Initial trajectory to seed the optimization. If None, the solver
uses the solution from the previous step.
max_attempts: Maximum number of attempts to solve the problem.
Returns:
WrapResult: Result of the optimization.
"""
converged = True
for _ in range(max_attempts):
result = self._step_once(current_state.clone(), shift_steps, seed_traj)
if (
torch.count_nonzero(torch.isnan(result.action.position)) == 0
and torch.count_nonzero(~result.metrics.feasible) == 0
):
converged = True
break
self.reset()
if not converged:
result.action.copy_(current_state)
log_warn("MPC didn't converge")
return result
def update_goal(self, goal: Goal):
"""Update the goal for MPC.
Args:
goal: goal object containing target pose or joint configuration. This goal instance
should be created using one of the setup_solve functions.
"""
self.solver.update_params(goal)
def reset(self):
"""Reset the solver."""
# reset warm start
self.solver.reset()
def reset_cuda_graph(self):
"""Reset captured CUDA graph. This does not work."""
self.solver.reset_cuda_graph()
def enable_cspace_cost(self, enable=True):
"""Enable or disable reaching joint configuration cost in the solver.
Args:
enable: Enable or disable reaching joint configuration cost. When False, cspace cost
is disabled.
"""
self.solver.safety_rollout.enable_cspace_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_cspace_cost(enable)
def enable_pose_cost(self, enable=True):
"""Enable or disable reaching pose cost in the solver.
Args:
enable: Enable or disable reaching pose cost. When False, pose cost is disabled.
"""
self.solver.safety_rollout.enable_pose_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_pose_cost(enable)
def get_active_js(
self,
in_js: JointState,
):
"""Get controlled joints indexed in MPC order from the input joint state.
Args:
in_js: Input joint state.
Returns:
JointState: Joint state with controlled joints.
"""
opt_jnames = self.rollout_fn.joint_names
opt_js = in_js.get_ordered_joint_state(opt_jnames)
return opt_js
def update_world(self, world: WorldConfig):
"""Update the collision world for the solver.
This allows for updating the world representation as long as the new world representation
does not have a larger number of obstacles than the :attr:`MpcSolver.collision_cache` as
created during initialization of :class:`MpcSolverConfig`.
Args:
world: New collision world configuration. See :ref:`world_collision` for more details.
"""
self.world_coll_checker.load_collision_model(world)
def get_visual_rollouts(self):
"""Get rollouts for debugging."""
return self.solver.optimizers[0].get_rollouts()
@property
def joint_names(self):
"""Get the ordered joint names of the robot."""
return self.rollout_fn.joint_names
@property
def collision_cache(self) -> Dict[str, int]:
"""Returns the collision cache created by the world collision checker."""
return self.world_coll_checker.cache
@property
def kinematics(self) -> CudaRobotModel:
"""Get kinematics instance of the robot."""
return self.solver.safety_rollout.dynamics_model.robot_model
@property
def world_collision(self) -> WorldCollision:
"""Get the world collision checker."""
return self.world_coll_checker
@property
def rollout_fn(self) -> ArmReacher:
"""Get the rollout function."""
return self.solver.safety_rollout
def _step_once(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
"""Solve for the next action given the current state.
Args:
current_state: Current joint state of the robot.
shift_steps: Number of steps to shift the trajectory.
seed_traj: Initial trajectory to seed the optimization. If None, the solver
uses the solution from the previous step.
Returns:
WrapResult: Result of the optimization.
"""
# Create cuda graph for whole solve step including computation of metrics
# Including updation of goal buffers
if self._solve_state is None:
log_error("Need to first setup solve state before calling solve()")
if self.use_cuda_graph_full_step:
st_time = time.time()
if not self._cu_step_init:
self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
self._cu_state_in.copy_(current_state)
if seed_traj is not None:
self._cu_seed.copy_(seed_traj)
self._cu_step_graph.replay()
result = self._cu_result.clone()
torch.cuda.synchronize(device=self.tensor_args.device)
result.solve_time = time.time() - st_time
else:
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
self._step_goal_buffer,
shift_steps,
seed_traj,
)
return result
def _update_solve_state_and_goal_buffer(
self,
solve_state: ReacherSolveState,
goal: Goal,
) -> Goal:
"""Update solve state and goal for MPC.
Args:
solve_state: New solve state.
goal: New goal buffer.
Returns:
Goal: Updated goal buffer.
"""
self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal(
goal,
self._solve_state,
@@ -250,71 +671,64 @@ class MpcSolver(MpcSolverConfig):
)
return self._goal_buffer
def step(
def _update_batch_size(self, batch_size: int):
"""Update the batch size of the solver.
Args:
batch_size: Number of problems to solve in parallel.
"""
if self.batch_size != batch_size:
self.batch_size = batch_size
def _solve_from_solve_state(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
max_attempts: int = 1,
):
converged = True
for _ in range(max_attempts):
result = self.step_once(current_state.clone(), shift_steps, seed_traj)
if (
torch.count_nonzero(torch.isnan(result.action.position)) == 0
and torch.max(torch.abs(result.action.position)) < 10.0
and torch.count_nonzero(~result.metrics.feasible) == 0
):
converged = True
break
self.reset()
if not converged:
result.action.copy_(current_state)
log_warn("NOT CONVERGED")
return result
def step_once(
self,
current_state: JointState,
solve_state: ReacherSolveState,
goal: Goal,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
# Create cuda graph for whole solve step including computation of metrics
# Including updation of goal buffers
"""Solve for the next action given the current state.
if self._solve_state is None:
log_error("Need to first setup solve state before calling solve()")
Args:
solve_state: solve state object containing information about the current MPC problem.
goal: goal object containing target pose or joint configuration.
shift_steps: Number of steps to shift the trajectory before optimization.
seed_traj: Initial trajectory to seed the optimization. If None, the solver
uses the solution from the previous step.
Returns:
WrapResult: Result of the optimization.
"""
if solve_state.batch_env:
if solve_state.batch_size > self.world_coll_checker.n_envs:
log_error("Batch Env is less that goal batch")
goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
if self.use_cuda_graph_full_step:
st_time = time.time()
if not self._cu_step_init:
self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
self._cu_state_in.copy_(current_state)
if seed_traj is not None:
self._cu_seed.copy_(seed_traj)
self._cu_step_graph.replay()
result = self._cu_result.clone()
torch.cuda.synchronize()
result.solve_time = time.time() - st_time
else:
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
self._step_goal_buffer,
shift_steps,
seed_traj,
)
self.solver.update_init_seed(seed_traj)
result = self.solver.solve(goal_buffer, seed_traj, shift_steps)
result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action)
return result
def _step(
def _mpc_step(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
):
"""One step function that is used to create a CUDA graph.
Args:
current_state: Current joint state of the robot.
shift_steps: Number of steps to shift the trajectory.
seed_traj: Initial trajectory to seed the optimization. If None, the solver
uses the solution from the previous step.
Returns:
WrapResult: Result of the optimization.
"""
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
@@ -331,6 +745,14 @@ class MpcSolver(MpcSolverConfig):
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
):
"""Create a CUDA graph for the full step of MPC.
Args:
current_state: Current joint state of the robot.
shift_steps: Number of steps to shift the trajectory.
seed_traj: Initial trajectory to seed the optimization. If None, the solver
uses the solution from the previous step.
"""
log_info("MpcSolver: Creating Cuda Graph")
self._cu_state_in = current_state.clone()
if seed_traj is not None:
@@ -339,7 +761,7 @@ class MpcSolver(MpcSolverConfig):
s.wait_stream(torch.cuda.current_stream(device=self.tensor_args.device))
with torch.cuda.stream(s):
for _ in range(3):
self._cu_result = self._step(
self._cu_result = self._mpc_step(
self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed
)
torch.cuda.current_stream(device=self.tensor_args.device).wait_stream(s)
@@ -350,142 +772,3 @@ class MpcSolver(MpcSolverConfig):
self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed
)
self._cu_step_init = True
def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.GOALSET,
num_mpc_seeds=num_seeds,
batch_size=1,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=1,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
solve_state = ReacherSolveState(
ReacherSolveType.BATCH_ENV_GOALSET,
num_mpc_seeds=num_seeds,
batch_size=goal.batch,
n_envs=1,
n_goalset=goal.n_goalset,
)
goal_buffer = self.update_goal_buffer(solve_state, goal)
return goal_buffer
def _solve_from_solve_state(
self,
solve_state: ReacherSolveState,
goal: Goal,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
if solve_state.batch_env:
if solve_state.batch_size > self.world_coll_checker.n_envs:
raise ValueError("Batch Env is less that goal batch")
goal_buffer = self.update_goal_buffer(solve_state, goal)
# NOTE: implement initialization from seed set here:
if seed_traj is not None:
self.solver.update_init_seed(seed_traj)
result = self.solver.solve(goal_buffer, seed_traj, shift_steps)
result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action)
return result
def fn(self):
# this will run one step of optimization and get new command
pass
def update_goal(self, goal: Goal):
self.solver.update_params(goal)
return True
def reset(self):
# reset warm start
self.solver.reset()
pass
def reset_cuda_graph(self):
self.solver.reset_cuda_graph()
@property
def rollout_fn(self):
return self.solver.safety_rollout
def enable_cspace_cost(self, enable=True):
self.solver.safety_rollout.enable_cspace_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_cspace_cost(enable)
def enable_pose_cost(self, enable=True):
self.solver.safety_rollout.enable_pose_cost(enable)
for opt in self.solver.optimizers:
opt.rollout_fn.enable_pose_cost(enable)
def get_active_js(
self,
in_js: JointState,
):
opt_jnames = self.rollout_fn.joint_names
opt_js = in_js.get_ordered_joint_state(opt_jnames)
return opt_js
@property
def joint_names(self):
return self.rollout_fn.joint_names
def update_world(self, world: WorldConfig):
self.world_coll_checker.load_collision_model(world)
return True
def get_visual_rollouts(self):
return self.solver.optimizers[0].get_rollouts()
@property
def kinematics(self):
return self.solver.safety_rollout.dynamics_model.robot_model
@property
def world_collision(self):
return self.world_coll_checker

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

File diff suppressed because it is too large Load Diff

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

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

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

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

9
tests/interpolation_test.py
View File

@@ -43,16 +43,21 @@ def test_linear_interpolation():
# create max_velocity buffer:
out_traj_gpu, _, _ = get_batch_interpolated_trajectory(
in_traj, int_dt, max_vel, max_acc=max_acc, max_jerk=max_jerk, raw_dt=raw_dt
in_traj,
raw_dt,
int_dt,
max_vel,
max_acc=max_acc,
max_jerk=max_jerk,
)
#
out_traj_gpu = out_traj_gpu.clone()
out_traj_cpu, _, _ = get_batch_interpolated_trajectory(
in_traj,
raw_dt,
int_dt,
max_vel,
raw_dt=raw_dt,
kind=InterpolateType.LINEAR,
max_acc=max_acc,
max_jerk=max_jerk,

330
tests/motion_gen_cuda_graph_test.py Normal file
View File

@@ -0,0 +1,330 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import pytest
import torch
# CuRobo
from curobo.geom.types import WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig
from curobo.util.trajectory import InterpolateType
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
@pytest.fixture(scope="function")
def motion_gen():
tensor_args = TensorDeviceType()
world_file = "collision_table.yml"
robot_file = "franka.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_file,
tensor_args,
use_cuda_graph=False,
)
motion_gen_instance = MotionGen(motion_gen_config)
return motion_gen_instance
@pytest.fixture(scope="function")
def motion_gen_batch_env():
tensor_args = TensorDeviceType()
world_files = ["collision_table.yml", "collision_test.yml"]
world_cfg = [
WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
for world_file in world_files
]
robot_file = "franka.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_cfg,
tensor_args,
use_cuda_graph=False,
)
motion_gen_instance = MotionGen(motion_gen_config)
return motion_gen_instance
@pytest.mark.parametrize(
"motion_gen_str,interpolation",
[
("motion_gen", InterpolateType.LINEAR),
("motion_gen", InterpolateType.CUBIC),
# ("motion_gen", InterpolateType.KUNZ_STILMAN_OPTIMAL),
("motion_gen", InterpolateType.LINEAR_CUDA),
],
)
def test_motion_gen_single(motion_gen_str, interpolation, request):
motion_gen = request.getfixturevalue(motion_gen_str)
motion_gen.update_interpolation_type(interpolation)
motion_gen.warmup()
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(False, True)
result = motion_gen.plan_single(start_state, goal_pose, m_config)
# get final solutions:
assert torch.count_nonzero(result.success) == 1
reached_state = motion_gen.compute_kinematics(result.optimized_plan[-1])
assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
def test_motion_gen_goalset(motion_gen):
motion_gen.warmup(n_goalset=2)
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(2, 1).view(1, -1, 3),
quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
)
goal_pose.position[0, 0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(False, True)
result = motion_gen.plan_goalset(start_state, goal_pose, m_config)
# get final solutions:
assert torch.count_nonzero(result.success) == 1
reached_state = motion_gen.compute_kinematics(result.optimized_plan[-1])
assert (
torch.min(
torch.norm(goal_pose.position[:, 0, :] - reached_state.ee_pos_seq),
torch.norm(goal_pose.position[:, 1, :] - reached_state.ee_pos_seq),
)
< 0.005
)
assert result.goalset_index is not None
assert (
torch.norm(goal_pose.position[:, result.goalset_index, :] - reached_state.ee_pos_seq)
< 0.005
)
def test_motion_gen_batch_goalset(motion_gen):
motion_gen.warmup(n_goalset=3, batch=3, warmup_js_trajopt=False, enable_graph=False)
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(6, 1).view(3, -1, 3).clone(),
quaternion=state.ee_quat_seq.repeat(6, 1).view(3, -1, 4).clone(),
)
goal_pose.position[0, 1, 1] = 0.2
goal_pose.position[1, 0, 1] = 0.2
goal_pose.position[2, 1, 1] = 0.2
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
m_config = MotionGenPlanConfig(False, True, max_attempts=1, enable_graph_attempt=None)
result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config)
# get final solutions:
assert torch.count_nonzero(result.success) == result.success.shape[0]
reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1))
#
goal_position = torch.cat(
[
goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
for x in range(len(result.goalset_index))
]
)
assert result.goalset_index is not None
assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
def test_motion_gen_batch(motion_gen):
motion_gen.warmup(batch=2)
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
).repeat_seeds(2)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2)
goal_pose.position[1, 0] -= 0.1
m_config = MotionGenPlanConfig(False, True)
result = motion_gen.plan_batch(start_state, goal_pose.clone(), m_config)
assert torch.count_nonzero(result.success) == 2
# get final solutions:
q = result.optimized_plan.trim_trajectory(-1).squeeze(1)
reached_state = motion_gen.compute_kinematics(q)
assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
@pytest.mark.parametrize(
"motion_gen_str,interpolation",
[
("motion_gen", InterpolateType.LINEAR),
("motion_gen", InterpolateType.CUBIC),
# ("motion_gen", InterpolateType.KUNZ_STILMAN_OPTIMAL),
("motion_gen", InterpolateType.LINEAR_CUDA),
],
)
def test_motion_gen_batch_graph(motion_gen_str: str, interpolation: InterpolateType, request):
motion_gen = request.getfixturevalue(motion_gen_str)
motion_gen.graph_planner.interpolation_type = interpolation
motion_gen.reset()
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
).repeat_seeds(5)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(5)
goal_pose.position[1, 0] -= 0.05
m_config = MotionGenPlanConfig(True, False)
result = motion_gen.plan_batch(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) > 0
# get final solutions:
q = result.interpolated_plan.trim_trajectory(-1).squeeze(1)
reached_state = motion_gen.compute_kinematics(q)
assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
def test_motion_gen_batch_env(motion_gen_batch_env):
motion_gen_batch_env.warmup(batch=2, batch_env_mode=True, enable_graph=False)
# motion_gen_batch_env.reset()
retract_cfg = motion_gen_batch_env.get_retract_config()
state = motion_gen_batch_env.compute_kinematics(
JointState.from_position(retract_cfg.view(1, -1))
)
goal_pose = Pose(
state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
).repeat_seeds(2)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2)
goal_pose.position[1, 0] -= 0.1
m_config = MotionGenPlanConfig(False, True, max_attempts=1)
result = motion_gen_batch_env.plan_batch_env(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) == 2
# get final solutions:
reached_state = motion_gen_batch_env.compute_kinematics(
result.optimized_plan.trim_trajectory(-1).squeeze(1)
)
assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
def test_motion_gen_batch_env_goalset(motion_gen_batch_env):
motion_gen_batch_env.warmup(batch=2, batch_env_mode=True, n_goalset=2, enable_graph=False)
retract_cfg = motion_gen_batch_env.get_retract_config()
state = motion_gen_batch_env.compute_kinematics(
JointState.from_position(retract_cfg.view(1, -1))
)
goal_pose = Pose(
state.ee_pos_seq.repeat(4, 1).view(2, -1, 3),
quaternion=state.ee_quat_seq.repeat(4, 1).view(2, -1, 4),
)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2)
goal_pose.position[1, 0] -= 0.2
m_config = MotionGenPlanConfig(False, True, enable_graph_attempt=None)
result = motion_gen_batch_env.plan_batch_env_goalset(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) > 0
# get final solutions:
reached_state = motion_gen_batch_env.compute_kinematics(
result.optimized_plan.trim_trajectory(-1).squeeze(1)
)
assert (
torch.min(
torch.norm(goal_pose.position[:, 0, :] - reached_state.ee_pos_seq),
torch.norm(goal_pose.position[:, 1, :] - reached_state.ee_pos_seq),
)
< 0.005
)
goal_position = torch.cat(
[
goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
for x in range(len(result.goalset_index))
]
)
assert result.goalset_index is not None
assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
@pytest.mark.parametrize(
"motion_gen_str,enable_graph",
[
("motion_gen", True),
("motion_gen", False),
],
)
def test_motion_gen_single_js(motion_gen_str, enable_graph, request):
motion_gen = request.getfixturevalue(motion_gen_str)
motion_gen.warmup(warmup_js_trajopt=True)
retract_cfg = motion_gen.get_retract_config()
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(enable_graph=enable_graph, max_attempts=2)
goal_state = start_state.clone()
goal_state.position -= 0.3
result = motion_gen.plan_single_js(start_state, goal_state, m_config)
assert torch.count_nonzero(result.success) == 1
reached_state = result.optimized_plan[-1]
assert torch.norm(goal_state.position - reached_state.position) < 0.05

141
tests/motion_gen_speed_test.py
View File

@@ -40,6 +40,24 @@ def motion_gen(request):
return motion_gen_instance
@pytest.fixture(scope="module")
def motion_gen_ur5e():
tensor_args = TensorDeviceType()
world_file = "collision_table.yml"
robot_file = "ur5e.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_file,
tensor_args,
interpolation_steps=10000,
interpolation_dt=0.05,
)
motion_gen_instance = MotionGen(motion_gen_config)
motion_gen_instance.warmup(warmup_js_trajopt=False, enable_graph=False)
return motion_gen_instance
@pytest.mark.parametrize(
"motion_gen",
[
@@ -66,3 +84,126 @@ def test_motion_gen_velocity_scale(motion_gen):
result = motion_gen.plan_single(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) == 1
@pytest.mark.parametrize(
"velocity_scale, acceleration_scale",
[
(1.0, 1.0),
(0.75, 1.0),
(0.5, 1.0),
(0.25, 1.0),
(0.15, 1.0),
(0.1, 1.0),
(1.0, 0.1),
(0.75, 0.1),
(0.5, 0.1),
(0.25, 0.1),
(0.15, 0.1),
(0.1, 0.1),
],
)
def test_pose_sequence_speed_ur5e_scale(velocity_scale, acceleration_scale):
# load ur5e motion gen:
world_file = "collision_table.yml"
robot_file = "ur5e.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_file,
interpolation_dt=(1.0 / 5.0),
velocity_scale=velocity_scale,
acceleration_scale=acceleration_scale,
)
motion_gen = MotionGen(motion_gen_config)
motion_gen.warmup(warmup_js_trajopt=False, enable_graph=False)
retract_cfg = motion_gen.get_retract_config()
start_state = JointState.from_position(retract_cfg.view(1, -1))
# poses for ur5e:
home_pose = [-0.431, 0.172, 0.348, 0, 1, 0, 0]
pose_1 = [0.157, -0.443, 0.427, 0, 1, 0, 0]
pose_2 = [0.126, -0.443, 0.729, 0, 0, 1, 0]
pose_3 = [-0.449, 0.339, 0.414, -0.681, -0.000, 0.000, 0.732]
pose_4 = [-0.449, 0.339, 0.414, 0.288, 0.651, -0.626, -0.320]
pose_5 = [-0.218, 0.508, 0.670, 0.529, 0.169, 0.254, 0.792]
pose_6 = [-0.865, 0.001, 0.411, 0.286, 0.648, -0.628, -0.321]
pose_list = [home_pose, pose_1, pose_2, pose_3, pose_4, pose_5, pose_6, home_pose]
trajectory = start_state
motion_time = 0
fail = 0
for i, pose in enumerate(pose_list):
goal_pose = Pose.from_list(pose, q_xyzw=False)
start_state = trajectory[-1].unsqueeze(0).clone()
start_state.velocity[:] = 0.0
start_state.acceleration[:] = 0.0
result = motion_gen.plan_single(
start_state.clone(),
goal_pose,
plan_config=MotionGenPlanConfig(
max_attempts=5,
),
)
if result.success.item():
plan = result.get_interpolated_plan()
trajectory = trajectory.stack(plan.clone())
motion_time += result.motion_time
else:
fail += 1
assert fail == 0
@pytest.mark.parametrize(
"motion_gen_str, time_dilation_factor",
[
("motion_gen_ur5e", 1.0),
("motion_gen_ur5e", 0.75),
("motion_gen_ur5e", 0.5),
("motion_gen_ur5e", 0.25),
("motion_gen_ur5e", 0.15),
("motion_gen_ur5e", 0.1),
("motion_gen_ur5e", 0.001),
],
)
def test_pose_sequence_speed_ur5e_time_dilation(motion_gen_str, time_dilation_factor, request):
# load ur5e motion gen:
motion_gen = request.getfixturevalue(motion_gen_str)
retract_cfg = motion_gen.get_retract_config()
start_state = JointState.from_position(retract_cfg.view(1, -1))
# poses for ur5e:
home_pose = [-0.431, 0.172, 0.348, 0, 1, 0, 0]
pose_1 = [0.157, -0.443, 0.427, 0, 1, 0, 0]
pose_2 = [0.126, -0.443, 0.729, 0, 0, 1, 0]
pose_3 = [-0.449, 0.339, 0.414, -0.681, -0.000, 0.000, 0.732]
pose_4 = [-0.449, 0.339, 0.414, 0.288, 0.651, -0.626, -0.320]
pose_5 = [-0.218, 0.508, 0.670, 0.529, 0.169, 0.254, 0.792]
pose_6 = [-0.865, 0.001, 0.411, 0.286, 0.648, -0.628, -0.321]
pose_list = [home_pose, pose_1, pose_2, pose_3, pose_4, pose_5, pose_6, home_pose]
trajectory = start_state
motion_time = 0
fail = 0
for i, pose in enumerate(pose_list):
goal_pose = Pose.from_list(pose, q_xyzw=False)
start_state = trajectory[-1].unsqueeze(0).clone()
start_state.velocity[:] = 0.0
start_state.acceleration[:] = 0.0
result = motion_gen.plan_single(
start_state.clone(),
goal_pose,
plan_config=MotionGenPlanConfig(
max_attempts=5,
time_dilation_factor=time_dilation_factor,
),
)
if result.success.item():
plan = result.get_interpolated_plan()
trajectory = trajectory.stack(plan.clone())
motion_time += result.motion_time
else:
fail += 1
assert fail == 0
assert motion_time < 15 * (1 / time_dilation_factor)

33
tests/trajopt_test.py
View File

@@ -57,15 +57,15 @@ def trajopt_solver_batch_env():
robot_cfg = RobotConfig.from_dict(
load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
)
# world_cfg = WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
trajopt_config = TrajOptSolverConfig.load_from_robot_config(
robot_cfg,
world_cfg,
tensor_args,
use_cuda_graph=False,
num_seeds=10,
num_seeds=4,
evaluate_interpolated_trajectory=True,
grad_trajopt_iters=200,
)
trajopt_solver = TrajOptSolver(trajopt_config)
@@ -73,7 +73,7 @@ def trajopt_solver_batch_env():
def test_trajopt_single_js(trajopt_solver):
q_start = trajopt_solver.retract_config
q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.2
goal_state = JointState.from_position(q_goal)
current_state = JointState.from_position(q_start)
@@ -88,7 +88,7 @@ def test_trajopt_single_js(trajopt_solver):
def test_trajopt_single_pose(trajopt_solver):
trajopt_solver.reset_seed()
q_start = trajopt_solver.retract_config
q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.1
kin_state = trajopt_solver.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
@@ -102,7 +102,7 @@ def test_trajopt_single_pose(trajopt_solver):
def test_trajopt_single_pose_no_seed(trajopt_solver):
trajopt_solver.reset_seed()
q_start = trajopt_solver.retract_config
q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.05
kin_state = trajopt_solver.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
@@ -116,7 +116,7 @@ def test_trajopt_single_pose_no_seed(trajopt_solver):
def test_trajopt_single_goalset(trajopt_solver):
# run goalset planning:
q_start = trajopt_solver.retract_config
q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.1
kin_state = trajopt_solver.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
@@ -133,7 +133,7 @@ def test_trajopt_single_goalset(trajopt_solver):
def test_trajopt_batch(trajopt_solver):
# run goalset planning:
q_start = trajopt_solver.retract_config.repeat(2, 1)
q_start = trajopt_solver.retract_config.clone().repeat(2, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -153,7 +153,7 @@ def test_trajopt_batch(trajopt_solver):
def test_trajopt_batch_js(trajopt_solver):
# run goalset planning:
q_start = trajopt_solver.retract_config.repeat(2, 1)
q_start = trajopt_solver.retract_config.clone().repeat(2, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -173,7 +173,7 @@ def test_trajopt_batch_js(trajopt_solver):
def test_trajopt_batch_goalset(trajopt_solver):
# run goalset planning:
q_start = trajopt_solver.retract_config.repeat(3, 1)
q_start = trajopt_solver.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -196,14 +196,12 @@ def test_trajopt_batch_goalset(trajopt_solver):
def test_trajopt_batch_env_js(trajopt_solver_batch_env):
# run goalset planning:
q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
q_start = trajopt_solver_batch_env.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal += 0.1
q_goal[2] += 0.1
q_goal[2][0] += 0.1
q_goal[1] -= 0.2
# q_goal[2, -1] += 0.1
kin_state = trajopt_solver_batch_env.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
goal_state = JointState.from_position(q_goal)
current_state = JointState.from_position(q_start)
@@ -213,14 +211,15 @@ def test_trajopt_batch_env_js(trajopt_solver_batch_env):
traj = result.solution.position
interpolated_traj = result.interpolated_solution.position
assert torch.count_nonzero(result.success) == 3
assert torch.linalg.norm((goal_state.position - traj[:, -1, :])).item() < 0.005
assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.005
error = torch.linalg.norm((goal_state.position - traj[:, -1, :]), dim=-1)
assert torch.max(error).item() < 0.05
assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.05
assert len(result) == 3
def test_trajopt_batch_env(trajopt_solver_batch_env):
# run goalset planning:
q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
q_start = trajopt_solver_batch_env.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -262,7 +261,7 @@ def test_trajopt_batch_env_goalset(trajopt_solver_batch_env):
def test_trajopt_batch_env(trajopt_solver):
# run goalset planning:
q_start = trajopt_solver.retract_config.repeat(3, 1)
q_start = trajopt_solver.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1