From 7362ccd4c2bd9e48539582356a5dbc3c2e630452 Mon Sep 17 00:00:00 2001 From: Balakumar Sundaralingam Date: Thu, 25 Apr 2024 12:24:17 -0700 Subject: [PATCH] Add re-timing, minimum dt robustness --- CHANGELOG.md | 149 +- benchmark/curobo_benchmark.py | 4 +- benchmark/curobo_nvblox_benchmark.py | 227 +- benchmark/curobo_voxel_benchmark.py | 5 +- benchmark/generate_nvblox_images.py | 4 +- benchmark/ik_benchmark.py | 1 + docker/user.dockerfile | 6 +- docker/x86.dockerfile | 9 +- .../isaac_sim/batch_motion_gen_reacher.py | 13 +- examples/motion_gen_example.py | 64 +- setup.cfg | 2 +- src/curobo/__init__.py | 16 +- src/curobo/content/configs/robot/franka.yml | 12 +- .../configs/robot/spheres/franka_mesh.yml | 8 +- .../configs/robot/ur5e_robotiq_2f_140.yml | 2 +- src/curobo/content/configs/task/base_cfg.yml | 4 +- .../content/configs/task/finetune_trajopt.yml | 34 +- .../configs/task/finetune_trajopt_slow.yml | 34 +- .../content/configs/task/gradient_ik.yml | 18 +- .../content/configs/task/gradient_trajopt.yml | 22 +- .../content/configs/task/particle_trajopt.yml | 24 +- .../cuda_robot_model/cuda_robot_model.py | 71 +- src/curobo/cuda_robot_model/types.py | 49 - src/curobo/curobolib/cpp/sphere_obb_kernel.cu | 17 +- src/curobo/geom/cv.py | 6 +- src/curobo/geom/sdf/world.py | 4 +- src/curobo/geom/sdf/world_blox.py | 12 +- src/curobo/geom/sdf/world_voxel.py | 6 +- src/curobo/geom/sphere_fit.py | 1 + src/curobo/geom/types.py | 4 +- src/curobo/opt/newton/newton_base.py | 6 +- src/curobo/opt/opt_base.py | 2 +- src/curobo/rollout/arm_base.py | 3 + src/curobo/rollout/arm_reacher.py | 3 +- src/curobo/rollout/cost/bound_cost.py | 119 +- .../rollout/cost/self_collision_cost.py | 1 - src/curobo/types/math.py | 2 +- src/curobo/types/robot.py | 2 + src/curobo/util/logger.py | 78 +- src/curobo/util/metrics.py | 3 + src/curobo/util/trajectory.py | 103 +- src/curobo/wrap/model/robot_segmenter.py | 7 +- src/curobo/wrap/model/robot_world.py | 3 +- src/curobo/wrap/reacher/evaluator.py | 9 +- src/curobo/wrap/reacher/ik_solver.py | 45 +- src/curobo/wrap/reacher/motion_gen.py | 3670 ++++++++++------- src/curobo/wrap/reacher/mpc.py | 701 +++- src/curobo/wrap/reacher/trajopt.py | 856 +++- src/curobo/wrap/wrap_base.py | 4 + src/curobo/wrap/wrap_mpc.py | 4 +- tests/interpolation_test.py | 9 +- tests/motion_gen_cuda_graph_test.py | 330 ++ tests/motion_gen_speed_test.py | 141 + tests/trajopt_test.py | 33 +- 54 files changed, 4773 insertions(+), 2189 deletions(-) create mode 100644 tests/motion_gen_cuda_graph_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 01d122f..e1350cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,27 +10,66 @@ 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 - Add mimic joint parsing and optimization support. Check `ur5e_robotiq_2f_140.yml`. -- Add `finetune_dt_scale` as a parameter to `MotionGenPlanConfig` to dynamically change the +- Add `finetune_dt_scale` as a parameter to `MotionGenPlanConfig` to dynamically change the time-optimal scaling on a per problem instance. -- `MotionGen.plan_single()` will now try finetuning in a for-loop, with larger and larger dt +- `MotionGen.plan_single()` will now try finetuning in a for-loop, with larger and larger dt until convergence. This also warm starts from previous failure. - Add `high_precision` mode to `MotionGenConfig` to support `<1mm` convergence. ### Changes in default behavior -- collision_sphere_buffer now supports having offset per link. Also, collision_sphere_buffer only -applies to world collision while self_collision_buffer applies for self collision. Previously, +- collision_sphere_buffer now supports having offset per link. Also, collision_sphere_buffer only +applies to world collision while self_collision_buffer applies for self collision. Previously, self_collision_buffer was added on top of collision_sphere_buffer. -- `TrajEvaluatorConfig` cannot be initialized without dof as now per-joint jerk and acceleration +- `TrajEvaluatorConfig` cannot be initialized without dof as now per-joint jerk and acceleration limits are used. Use `TrajEvaluatorConfig.from_basic()` to initialize similar to previous behavior. - `finetune_dt_scale` default value is 0.9 from 0.95. ### 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`. @@ -41,16 +80,16 @@ limits are used. Use `TrajEvaluatorConfig.from_basic()` to initialize similar to - Reduced branching in Kinematics kernels and added mimic joint computations. - Add init_cache to WorldVoxelCollision to create cache for Mesh and Cuboid obstacles. - `TrajEvaluator` now uses per-joint acceleration and jerk limits. -- Fixed regression in `batch_motion_gen_reacher.py` example where robot's position was not being +- Fixed regression in `batch_motion_gen_reacher.py` example where robot's position was not being set correctly. - Switched from smooth l2 to l2 for BoundCost as that gives better convergence. - `requires_grad` is explicitly stored in a varaible before `tensor.detach()` in warp kernel calls as this can get set to False in some instances. - Fix dt update in `MotionGen.plan_single_js()` where dt was not reset after finetunestep, causing joint space planner to fail often. -- Improve joint space planner success by changing smooth l2 distance cost to l2 distance. Also, +- Improve joint space planner success by changing smooth l2 distance cost to l2 distance. Also, added fallback to graph planner when linear path is not possible. -- Retuned weigths for IKSolver, now 98th percentile accuracy is 10 micrometers wtih 16 seeds +- Retuned weigths for IKSolver, now 98th percentile accuracy is 10 micrometers wtih 16 seeds (vs 24 seeds previously). - Switch float8 precision check from `const` to macro to avoid compile errors in older nvcc, this fixes docker build issues for isaac sim 2023.1.0. @@ -58,76 +97,77 @@ fixes docker build issues for isaac sim 2023.1.0. ## Version 0.7.0 ### Changes in default behavior - Increased default collision cache to 50 in RobotWorld. -- Changed `CSpaceConfig.position_limit_clip` default to 0 as previous default of 0.01 can make -default start state in examples be out of bounds. -- MotionGen uses parallel_finetune by default. To get previous motion gen behavior, pass +- Changed `CSpaceConfig.position_limit_clip` default to 0 as previous default of 0.01 can make +default start state in examples be out of bounds. +- MotionGen uses parallel_finetune by default. To get previous motion gen behavior, pass `warmup(parallel_finetune=False)` and `MotionGenPlanConfig(parallel_finetune=False)`. - MotionGen loads Mesh Collision checker instead of Primitive by default. -- UR10e and UR5e now don't have a collision sphere at tool frame for world collision checking. This +- UR10e and UR5e now don't have a collision sphere at tool frame for world collision checking. This sphere is only active for self collision avoidance. -- With torch>=2.0, cuRobo will use `torch.compile` instead of `torch.jit.script` to generate fused -kernels. This can take several seconds during the first run. To enable this feature, set +- With torch>=2.0, cuRobo will use `torch.compile` instead of `torch.jit.script` to generate fused +kernels. This can take several seconds during the first run. To enable this feature, set environment variable `export CUROBO_TORCH_COMPILE_DISABLE=0`. ### Breaking Changes - Renamed `copy_if_not_none` to `clone_if_not_none` to be more descriptive. Now `copy_if_not_none` will try to copy data into reference. -- Renamed `n_envs` in curobo.opt module to avoid confusion between parallel environments and +- Renamed `n_envs` in curobo.opt module to avoid confusion between parallel environments and parallel problems in optimization. - Added more inputs to pose distance kernels. Check `curobolib/geom.py`. - Pose cost `run_vec_weight` should now be `[0,0,0,0,0,0]` instead of `[1,1,1,1,1,1]` - ``max_distance`` is now tensor from ``float`` and is an input to collision kernels. -- Order of inputs to ``SweptSdfMeshWarpPy`` has changed. +- Order of inputs to ``SweptSdfMeshWarpPy`` has changed. ### New Features -- Add function to disable and enable collision for specific links in KinematicsTensorConfig. +- Add function to disable and enable collision for specific links in KinematicsTensorConfig. - Add goal index to reacher results to return index of goal reached when goalset planning. - 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`. +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 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`. -- Significantly improved collision computation for cuboids and meshes. Mesh collision checker is -now only 2x slower than cuboid (from 5x slower). Optimization convergence is also improved. +- Use `torch.compile` to leverage better kernel fusion in place of `torch.jit.script`. +- Significantly improved collision computation for cuboids and meshes. Mesh collision checker is +now only 2x slower than cuboid (from 5x slower). Optimization convergence is also improved. - LBFGS kernels now support ``history <= 31`` from ``history <= 15``. -- 2x faster LBFGS kernel that allocates upto 68kb of shared memory, preventing use in CUDA devices +- 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. -- Add ``WorldCollisionVoxel``, a new collision checking implementation that uses a voxel grid +- 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. +- 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. -- Add API for ESDF computation from world representations, see +- Add API for ESDF computation from world representations, see ``WorldCollision.get_esdf_in_bounding_box()``. -- Add partial support for isaac sim 2023.1.1. Most examples run for UR robots. `Franka Panda` is +- Add partial support for isaac sim 2023.1.1. Most examples run for UR robots. `Franka Panda` is unstable. ### BugFixes & Misc. - refactored wp.index() instances to `[]` to avoid errors in future releases of warp. - Fix bug in gaussian transformation to ensure values are not -1 or +1. - Fix bug in ik_solver loading ee_link_name from argument. -- Fix bug in batch_goalset planning, where pose cost was selected as GOALSET instead of +- Fix bug in batch_goalset planning, where pose cost was selected as GOALSET instead of BATCH_GOALSET. - Added package data to also export `.so` files. -- Fixed bug in transforming link visual mesh offset when reading from urdf. +- Fixed bug in transforming link visual mesh offset when reading from urdf. - Fixed bug in MotionGenPlanConfig.clone() that didn't clone the state of parallel_finetune. -- Increased weighting from 1.0 to 10.0 for optimized_dt in TrajEvaluator to select shorter +- Increased weighting from 1.0 to 10.0 for optimized_dt in TrajEvaluator to select shorter trajectories. - Improved determinism by setting global seed for random in `graph_nx.py`. - Added option to clear obstacles in WorldPrimitiveCollision. - Raise error when reference of tensors change in MotionGen, IKSolver, and TrajOpt when cuda graph is enabled. -- plan_single will get converted to plan_goalset when a plan_goalset was used to initialize cuda +- plan_single will get converted to plan_goalset when a plan_goalset was used to initialize cuda graph. - plan_goalset will pad for extra goals when called with less number of goal than initial creation. - Improved API documentation for Optimizer class. - Set `use_cuda_graph` to `True` as default from `None` in `MotionGenConfig.load_from_robot_config` -- Add batched mode to robot image segmentation, supports single robot multiple camera and batch -robot batch camera. +- Add batched mode to robot image segmentation, supports single robot multiple camera and batch +robot batch camera. - Add `log_warn` import to `arm_reacher.py` - Remove negative radius check in self collision kernel to allow for self collision checking with spheres of negative radius. @@ -135,9 +175,9 @@ spheres of negative radius. - Added UR5e robot with robotiq gripper (2f-140) with improved sphere model. - Fix bug in aarch64.dockerfile where curobo was cloned to wrong path. - Fix bug in aarch64.dockerfile where python was used instead of python3. -- Remove unused variables in kernels. +- Remove unused variables in kernels. - Added ``pybind11`` as a dependency as some pytorch dockers for Jetson do not have this installed. -- Fix incorrect dimensions in ``MotionGenResult.success`` in ``MotionGen.plan_batch()`` when +- Fix incorrect dimensions in ``MotionGenResult.success`` in ``MotionGen.plan_batch()`` when trajectory optimization fails. - Added unit tests for collision checking functions. - Fix bug in linear interpolation which was not reading the new ``optimized_dt`` to interpolate @@ -147,52 +187,53 @@ between trajopt and finetune_trajopt. ### Known Bugs (WIP) -- `Franka Panda` robot loading from urdf in isaac sim 2023.1.1 is unstable. +- `Franka Panda` robot loading from urdf in isaac sim 2023.1.1 is unstable. ## Version 0.6.2 ### New Features - Added support for actuated axis to be negative (i.e., urdf joints with `` are now natively supported). -- Improved gradient calculation to account for terminal state. Trajectory optimization can reach +- Improved gradient calculation to account for terminal state. Trajectory optimization can reach within 1mm of accuracy (median across 2600 problems at 0.017mm). -- Improved estimation of previous positions based on start velocity and acceleration. This enables +- Improved estimation of previous positions based on start velocity and acceleration. This enables Trajectory optimization to optimize from non-zero start velocity and accelerations. - Added graph planner and finetuning step to joint space planning (motion_gen.plan_single_js). This improves success and motion quality when planning to reach joint space targets. - Added finetuning across many seeds in motion_gen, improving success rate and motion quality. -- Add urdf support to usd helper to export optimization steps as animated usd files for debugging +- Add urdf support to usd helper to export optimization steps as animated usd files for debugging motion generation. Check `examples/usd_examples.py` for an example. -- Retuned weights for IK and Trajectory optimization. This (+ other fixes) significantly improves -pose reaching accuracy, IK accuracy improves by 100x (98th percentile < 10 micrometers) and motion +- Retuned weights for IK and Trajectory optimization. This (+ other fixes) significantly improves +pose reaching accuracy, IK accuracy improves by 100x (98th percentile < 10 micrometers) and motion generation median at 0.017mm (with). IK now solves most problems with 24 seeds (vs 30 seeds prev.). Run `benchmark/ik_benchmark.py` to get the latest results. -- Added `external_asset_path` to robot configuration to help in loading urdf and meshes from an +- Added `external_asset_path` to robot configuration to help in loading urdf and meshes from an 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 -- Cleanup docker scripts. Use `build_docker.sh` instead of `build_dev_docker.sh`. Added isaac sim +- 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.. +- Fixed bug in backward kinematics kernel, helped improve IK and TO pose reaching accuracy.. - Changed `panda_finger_joint2` from `` to `` in `franka_panda.urdf` to match real robot urdf as cuRobo now supports negative axis. -- Changed benchmarking scripts to use lock joint state of [0.025,0.025] for mpinets dataset. +- Changed benchmarking scripts to use lock joint state of [0.025,0.025] for mpinets dataset. - Added scaling of mesh to Mesh.get_trimesh_mesh() to help in debugging mesh world. - Improved stability and accuracy of MPPI for MPC. -- Added NaN checking in STOMP covariance computation to account for cases when cholesky decomp +- Added NaN checking in STOMP covariance computation to account for cases when cholesky decomp fails. - Added ground truth collision check validation in `benchmarks/curobo_nvblox_benchmark.py`. ### Performance Regressions - cuRobo now generates significantly shorter paths then previous version. E.g., cuRobo obtains -2.2 seconds 98th percentile motion time on the 2600 problems (`benchmark/curobo_benchmark.py`), -where previously it was at 3 seconds (1.36x quicker motions). This was obtained by retuning the -weights and slight reformulations of trajectory optimization. These changes have led to a slight +2.2 seconds 98th percentile motion time on the 2600 problems (`benchmark/curobo_benchmark.py`), +where previously it was at 3 seconds (1.36x quicker motions). This was obtained by retuning the +weights and slight reformulations of trajectory optimization. These changes have led to a slight degrade in planning time, 20ms slower on 4090 and 40ms on ORIN MAXN. We will address this slow down -in a later release. One way to avoid this regression is to set `finetune_dt_scale=1.05` in +in a later release. One way to avoid this regression is to set `finetune_dt_scale=1.05` in `MotionGenConfig.load_from_robot_config()`. diff --git a/benchmark/curobo_benchmark.py b/benchmark/curobo_benchmark.py index 5c582b1..8499969 100644 --- a/benchmark/curobo_benchmark.py +++ b/benchmark/curobo_benchmark.py @@ -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"] diff --git a/benchmark/curobo_nvblox_benchmark.py b/benchmark/curobo_nvblox_benchmark.py index 6217ef5..b80264e 100644 --- a/benchmark/curobo_nvblox_benchmark.py +++ b/benchmark/curobo_nvblox_benchmark.py @@ -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() + st_int_time = time.time() + world_nvblox.add_camera_frame(cam_obs, "world") + torch.cuda.synchronize() + int_time += time.time() - st_int_time + st_time = time.time() + world_nvblox.process_camera_frames("world", False) torch.cuda.synchronize() - mg.world_coll_checker.update_blox_hashes() - torch.cuda.synchronize() - # if i > 2: - # mg.world_coll_checker.clear_cache() - # mg.world_coll_checker.update_blox_hashes() + 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__": diff --git a/benchmark/curobo_voxel_benchmark.py b/benchmark/curobo_voxel_benchmark.py index 08171c3..191d500 100644 --- a/benchmark/curobo_voxel_benchmark.py +++ b/benchmark/curobo_voxel_benchmark.py @@ -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"] diff --git a/benchmark/generate_nvblox_images.py b/benchmark/generate_nvblox_images.py index 48681c0..994aae2 100644 --- a/benchmark/generate_nvblox_images.py +++ b/benchmark/generate_nvblox_images.py @@ -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 ) diff --git a/benchmark/ik_benchmark.py b/benchmark/ik_benchmark.py index 83216b8..6aaf40c 100644 --- a/benchmark/ik_benchmark.py +++ b/benchmark/ik_benchmark.py @@ -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")) diff --git a/docker/user.dockerfile b/docker/user.dockerfile index 1ad73d4..0b385e1 100644 --- a/docker/user.dockerfile +++ b/docker/user.dockerfile @@ -10,12 +10,12 @@ ## # Check architecture and load: -ARG IMAGE_TAG +ARG IMAGE_TAG 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 @@ -29,7 +29,7 @@ RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers # Set user -# RUN mkdir /home/$USERNAME && chown $USERNAME:$USERNAME /home/$USERNAME +# RUN mkdir /home/$USERNAME && chown $USERNAME:$USERNAME /home/$USERNAME USER $USERNAME WORKDIR /home/$USERNAME ENV USER=$USERNAME diff --git a/docker/x86.dockerfile b/docker/x86.dockerfile index b853afb..0bf970a 100644 --- a/docker/x86.dockerfile +++ b/docker/x86.dockerfile @@ -78,13 +78,13 @@ 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" # if you want to use a different version of curobo, create folder as docker/pkgs and put your -# version of curobo there. Then uncomment below line and comment the next line that clones from +# version of curobo there. Then uncomment below line and comment the next line that clones from # github # COPY pkgs /pkgs @@ -124,4 +124,7 @@ RUN cd /pkgs && git clone https://github.com/nvlabs/nvblox_torch.git && \ 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" \ No newline at end of file +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 \ No newline at end of file diff --git a/examples/isaac_sim/batch_motion_gen_reacher.py b/examples/isaac_sim/batch_motion_gen_reacher.py index 320bf19..9785dee 100644 --- a/examples/isaac_sim/batch_motion_gen_reacher.py +++ b/examples/isaac_sim/batch_motion_gen_reacher.py @@ -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] diff --git a/examples/motion_gen_example.py b/examples/motion_gen_example.py index ace6162..c01774a 100644 --- a/examples/motion_gen_example.py +++ b/examples/motion_gen_example.py @@ -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] diff --git a/setup.cfg b/setup.cfg index 3b8256c..f93dbfe 100644 --- a/setup.cfg +++ b/setup.cfg @@ -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 = diff --git a/src/curobo/__init__.py b/src/curobo/__init__.py index 403c5fa..85e12af 100644 --- a/src/curobo/__init__.py +++ b/src/curobo/__init__.py @@ -10,17 +10,18 @@ # """ -cuRobo provides accelerated modules for robotics which can be used to build high-performance +cuRobo provides accelerated modules for robotics which can be used to build high-performance robotics applications. The library has several modules for numerical optimization, robot kinematics, -geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for -performing tasks like collision-free inverse kinematics, model predictive control, and motion +geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for +performing tasks like collision-free inverse kinematics, model predictive control, and motion planning. High-level APIs: +- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`. - Inverse Kinematics: :mod:`curobo.wrap.reacher.ik_solver`. - Model Predictive Control: :mod:`curobo.wrap.reacher.mpc`. -- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`. +- Trajectory Optimization: :mod:`curobo.wrap.reacher.trajopt`. cuRobo package is split into several modules: @@ -33,10 +34,11 @@ cuRobo package is split into several modules: - :mod:`curobo.rollout` contains methods that map actions to costs. This class wraps instances of :mod:`curobo.cuda_robot_model` and :mod:`curobo.geom` to compute costs given trajectory of actions. - :mod:`curobo.util` contains utility methods. -- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of +- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of collision-free reacher and batched robot world collision checking. -- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including - :meth:`types.state.JointState`, :meth:`types.camera.CameraObservation`, :meth:`types.math.Pose`. +- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including + :py:meth:`~types.state.JointState`, :py:meth:`~types.camera.CameraObservation`, + :py:meth:`~types.math.Pose`. """ diff --git a/src/curobo/content/configs/robot/franka.yml b/src/curobo/content/configs/robot/franka.yml index c555f34..351bbd0 100644 --- a/src/curobo/content/configs/robot/franka.yml +++ b/src/curobo/content/configs/robot/franka.yml @@ -28,7 +28,7 @@ robot_cfg: "panda_finger_joint1": "Y", "panda_finger_joint2": "Y", } - + usd_flip_joint_limits: ["panda_finger_joint2"] urdf_path: "robot/franka_description/franka_panda.urdf" asset_root_path: "robot/franka_description" @@ -67,13 +67,13 @@ robot_cfg: "panda_hand": ["panda_leftfinger", "panda_rightfinger","attached_object"], "panda_leftfinger": ["panda_rightfinger", "attached_object"], "panda_rightfinger": ["attached_object"], - + } - + self_collision_buffer: { - "panda_link0": 0.1, - "panda_link1": 0.05, + "panda_link0": 0.1, + "panda_link1": 0.05, "panda_link2": 0.0, "panda_link3": 0.0, "panda_link4": 0.0, @@ -101,7 +101,7 @@ robot_cfg: "panda_rightfinger", ] lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04} - extra_links: {"attached_object":{"parent_link_name": "panda_hand" , + extra_links: {"attached_object":{"parent_link_name": "panda_hand" , "link_name": "attached_object", "fixed_transform": [0,0,0,1,0,0,0], "joint_type":"FIXED", "joint_name": "attach_joint" }} cspace: diff --git a/src/curobo/content/configs/robot/spheres/franka_mesh.yml b/src/curobo/content/configs/robot/spheres/franka_mesh.yml index 1dec717..b7ecae5 100644 --- a/src/curobo/content/configs/robot/spheres/franka_mesh.yml +++ b/src/curobo/content/configs/robot/spheres/franka_mesh.yml @@ -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 \ No newline at end of file + "radius": 0.011 \ No newline at end of file diff --git a/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml b/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml index e146757..1a377ec 100644 --- a/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml +++ b/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml @@ -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", diff --git a/src/curobo/content/configs/task/base_cfg.yml b/src/curobo/content/configs/task/base_cfg.yml index bd1eb42..27b79c2 100644 --- a/src/curobo/content/configs/task/base_cfg.yml +++ b/src/curobo/content/configs/task/base_cfg.yml @@ -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 diff --git a/src/curobo/content/configs/task/finetune_trajopt.yml b/src/curobo/content/configs/task/finetune_trajopt.yml index 6c1c66d..4ae1f61 100644 --- a/src/curobo/content/configs/task/finetune_trajopt.yml +++ b/src/curobo/content/configs/task/finetune_trajopt.yml @@ -19,12 +19,12 @@ model: acceleration: 1.0 enable: False dt_traj_params: - base_dt: 0.2 + base_dt: 0.15 base_ratio: 1.0 - max_dt: 0.2 + max_dt: 0.15 vel_scale: 1.0 control_space: 'POSITION' - state_finite_difference_mode: "CENTRAL" + state_finite_difference_mode: "CENTRAL" teleport_mode: False return_full_act_buffer: True @@ -35,7 +35,7 @@ cost: weight: [2000,500000.0,30,60] vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True - run_weight: 1.0 + run_weight: 1.0 use_metric: True link_pose_cfg: @@ -47,19 +47,19 @@ cost: run_weight: 1.0 use_metric: True - + cspace_cfg: weight: 10000.0 terminal: True run_weight: 0.00 #1 bound_cfg: - weight: [10000.0, 500000.0, 500.0, 500.0] - smooth_weight: [0.0,10000.0, 5.0, 0.0] + weight: [10000.0, 500000.0, 500.0, 500.0] + smooth_weight: [0.0,10000.0, 5.0, 0.0] run_weight_velocity: 0.0 run_weight_acceleration: 1.0 run_weight_jerk: 1.0 - activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk + activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk null_space_weight: [0.0] primitive_collision_cfg: @@ -75,21 +75,21 @@ cost: self_collision_cfg: weight: 5000.0 classify: False - - + + lbfgs: - n_iters: 300 - inner_iters: 25 - cold_start_n_iters: null + n_iters: 400 + inner_iters: 25 + cold_start_n_iters: null min_iters: 25 line_search_scale: [0.1, 0.3, 0.7,1.0] # # fixed_iters: True - cost_convergence: 0.01 - cost_delta_threshold: 1.0 - cost_relative_threshold: 0.999 #0.999 + cost_convergence: 0.001 + cost_delta_threshold: 0.0001 + cost_relative_threshold: 0.99 epsilon: 0.01 - history: 27 #15 + history: 27 use_cuda_graph: True n_problems: 1 store_debug: False diff --git a/src/curobo/content/configs/task/finetune_trajopt_slow.yml b/src/curobo/content/configs/task/finetune_trajopt_slow.yml index 383b903..8eb26c8 100644 --- a/src/curobo/content/configs/task/finetune_trajopt_slow.yml +++ b/src/curobo/content/configs/task/finetune_trajopt_slow.yml @@ -19,9 +19,9 @@ model: acceleration: 1.0 enable: False dt_traj_params: - base_dt: 0.2 + base_dt: 0.15 base_ratio: 1.0 - max_dt: 0.2 + max_dt: 0.15 vel_scale: 1.0 control_space: 'POSITION' state_finite_difference_mode: "CENTRAL" @@ -31,22 +31,21 @@ model: cost: pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps - run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position - weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40] + run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position + weight: [2000,500000.0,30,60] vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True - run_weight: 0.0 #0.05 + run_weight: 1.0 use_metric: True link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps - run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position - weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40] + run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position + weight: [2000,500000.0,30,60] #[150.0, 2000.0, 30, 40] vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True - run_weight: 0.001 + run_weight: 1.0 use_metric: True - cspace_cfg: weight: 10000.0 @@ -54,12 +53,12 @@ cost: run_weight: 0.00 #1 bound_cfg: - weight: [10000.0, 100000.0, 500.0, 500.0] - smooth_weight: [0.0,10000.0, 50.0, 0.0] + weight: [10000.0, 500000.0, 500.0, 500.0] + smooth_weight: [0.0,10000.0, 5.0, 0.0] run_weight_velocity: 0.0 run_weight_acceleration: 1.0 run_weight_jerk: 1.0 - activation_distance: [0.1,0.1,0.1,10.0] # for position, velocity, acceleration and jerk + activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk null_space_weight: [0.0] primitive_collision_cfg: @@ -75,21 +74,20 @@ cost: self_collision_cfg: weight: 5000.0 classify: False - - + lbfgs: - n_iters: 400 # 400 + n_iters: 400 inner_iters: 25 - cold_start_n_iters: 200 + cold_start_n_iters: null min_iters: 25 - line_search_scale: [0.01, 0.3, 0.7,1.0] # # + line_search_scale: [0.1, 0.3, 0.7,1.0] # # fixed_iters: True cost_convergence: 0.01 cost_delta_threshold: 1.0 cost_relative_threshold: 0.999 #0.999 epsilon: 0.01 - history: 15 #15 + history: 27 #15 use_cuda_graph: True n_problems: 1 store_debug: False diff --git a/src/curobo/content/configs/task/gradient_ik.yml b/src/curobo/content/configs/task/gradient_ik.yml index df3d90d..98354df 100644 --- a/src/curobo/content/configs/task/gradient_ik.yml +++ b/src/curobo/content/configs/task/gradient_ik.yml @@ -29,32 +29,30 @@ cost: pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] weight: [2000,10000,20,40] - vec_convergence: [0.0, 0.00] + vec_convergence: [0.0, 0.00] terminal: False use_metric: True run_weight: 1.0 link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] weight: [2000,10000,20,40] - vec_convergence: [0.00, 0.000] + vec_convergence: [0.00, 0.000] terminal: False use_metric: True run_weight: 1.0 - - cspace_cfg: weight: 0.000 bound_cfg: - weight: 500.0 - activation_distance: [0.1] - null_space_weight: [1.0] + weight: 5000.0 + activation_distance: [0.001] + null_space_weight: [0.1] primitive_collision_cfg: weight: 5000.0 use_sweep: False classify: False activation_distance: 0.01 self_collision_cfg: - weight: 500.0 + weight: 5000.0 classify: False @@ -81,10 +79,10 @@ lbfgs: use_cuda_line_search_kernel: True use_cuda_update_best_kernel: True sync_cuda_time: True - step_scale: 1.0 + step_scale: 1.0 use_coo_sparse: True last_best: 10 debug_info: visual_traj : null #'ee_pos_seq' - + diff --git a/src/curobo/content/configs/task/gradient_trajopt.yml b/src/curobo/content/configs/task/gradient_trajopt.yml index 506a25e..da82760 100644 --- a/src/curobo/content/configs/task/gradient_trajopt.yml +++ b/src/curobo/content/configs/task/gradient_trajopt.yml @@ -19,12 +19,12 @@ model: acceleration: 1.0 enable: False dt_traj_params: - base_dt: 0.2 + base_dt: 0.15 base_ratio: 1.0 - max_dt: 0.2 + max_dt: 0.15 vel_scale: 1.0 control_space: 'POSITION' - state_finite_difference_mode: "CENTRAL" + state_finite_difference_mode: "CENTRAL" teleport_mode: False return_full_act_buffer: True @@ -48,17 +48,17 @@ cost: use_metric: True cspace_cfg: - weight: 1000.0 + weight: 10000.0 terminal: True run_weight: 0.0 bound_cfg: - weight: [50000.0, 50000.0, 50000.0,50000.0] - smooth_weight: [0.0,10000.0,5.0, 0.0] + weight: [50000.0, 50000.0, 500.0,500.0] + smooth_weight: [0.0,10000.0,5.0, 0.0] run_weight_velocity: 0.00 run_weight_acceleration: 1.0 run_weight_jerk: 1.0 - activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk + activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk null_space_weight: [0.0] primitive_collision_cfg: @@ -75,13 +75,13 @@ cost: self_collision_cfg: weight: 5000.0 classify: False - - + + lbfgs: n_iters: 100 #175 - inner_iters: 25 - cold_start_n_iters: null + inner_iters: 25 + cold_start_n_iters: null min_iters: 25 line_search_scale: [0.1,0.3,0.7,1.0] fixed_iters: True diff --git a/src/curobo/content/configs/task/particle_trajopt.yml b/src/curobo/content/configs/task/particle_trajopt.yml index 268a613..0f2e288 100644 --- a/src/curobo/content/configs/task/particle_trajopt.yml +++ b/src/curobo/content/configs/task/particle_trajopt.yml @@ -18,9 +18,9 @@ model: acceleration: 1.0 enable: False dt_traj_params: - base_dt: 0.2 + base_dt: 0.15 base_ratio: 1.0 - max_dt: 0.2 + max_dt: 0.15 vel_scale: 1.0 control_space: 'POSITION' teleport_mode: False @@ -30,22 +30,22 @@ model: cost: pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight + run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight weight: [250.0, 5000.0, 20, 20] - vec_convergence: [0.0,0.0,1000.0,1000.0] + vec_convergence: [0.0,0.0,1000.0,1000.0] terminal: True run_weight: 0.0 use_metric: True - + link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight + run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight weight: [0.0, 5000.0, 40, 40] - vec_convergence: [0.0,0.0,1000.0,1000.0] + vec_convergence: [0.0,0.0,1000.0,1000.0] terminal: True - run_weight: 0.00 + run_weight: 0.00 use_metric: True - + cspace_cfg: weight: 10000.0 terminal: True @@ -74,8 +74,8 @@ cost: self_collision_cfg: weight: 500.0 classify: False - - + + mppi: @@ -89,7 +89,7 @@ mppi: alpha : 1 num_particles : 25 # 100 update_cov : True - cov_type : "DIAG_A" # + cov_type : "DIAG_A" # kappa : 0.001 null_act_frac : 0.0 sample_mode : 'MEAN' diff --git a/src/curobo/cuda_robot_model/cuda_robot_model.py b/src/curobo/cuda_robot_model/cuda_robot_model.py index 3135bb5..c7c99e3 100644 --- a/src/curobo/cuda_robot_model/cuda_robot_model.py +++ b/src/curobo/cuda_robot_model/cuda_robot_model.py @@ -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) diff --git a/src/curobo/cuda_robot_model/types.py b/src/curobo/cuda_robot_model/types.py index 8c4c20e..a213e5a 100644 --- a/src/curobo/cuda_robot_model/types.py +++ b/src/curobo/cuda_robot_model/types.py @@ -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 diff --git a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu index da250af..687344e 100644 --- a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu +++ b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu @@ -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); diff --git a/src/curobo/geom/cv.py b/src/curobo/geom/cv.py index b2b9da0..8d0e007 100644 --- a/src/curobo/geom/cv.py +++ b/src/curobo/geom/cv.py @@ -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 diff --git a/src/curobo/geom/sdf/world.py b/src/curobo/geom/sdf/world.py index 03d6550..fbe10eb 100644 --- a/src/curobo/geom/sdf/world.py +++ b/src/curobo/geom/sdf/world.py @@ -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]) diff --git a/src/curobo/geom/sdf/world_blox.py b/src/curobo/geom/sdf/world_blox.py index 3340406..0cba6d1 100644 --- a/src/curobo/geom/sdf/world_blox.py +++ b/src/curobo/geom/sdf/world_blox.py @@ -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, diff --git a/src/curobo/geom/sdf/world_voxel.py b/src/curobo/geom/sdf/world_voxel.py index 9eb9c62..0cdaad5 100644 --- a/src/curobo/geom/sdf/world_voxel.py +++ b/src/curobo/geom/sdf/world_voxel.py @@ -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"] ): diff --git a/src/curobo/geom/sphere_fit.py b/src/curobo/geom/sphere_fit.py index d8aa65d..2542aa0 100644 --- a/src/curobo/geom/sphere_fit.py +++ b/src/curobo/geom/sphere_fit.py @@ -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: diff --git a/src/curobo/geom/types.py b/src/curobo/geom/types.py index 8a15488..fe8b500 100644 --- a/src/curobo/geom/types.py +++ b/src/curobo/geom/types.py @@ -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. diff --git a/src/curobo/opt/newton/newton_base.py b/src/curobo/opt/newton/newton_base.py index ea9061d..590d18e 100644 --- a/src/curobo/opt/newton/newton_base.py +++ b/src/curobo/opt/newton/newton_base.py @@ -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() diff --git a/src/curobo/opt/opt_base.py b/src/curobo/opt/opt_base.py index f105685..1d41bb3 100644 --- a/src/curobo/opt/opt_base.py +++ b/src/curobo/opt/opt_base.py @@ -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 diff --git a/src/curobo/rollout/arm_base.py b/src/curobo/rollout/arm_base.py index f98279f..8560a42 100644 --- a/src/curobo/rollout/arm_base.py +++ b/src/curobo/rollout/arm_base.py @@ -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 diff --git a/src/curobo/rollout/arm_reacher.py b/src/curobo/rollout/arm_reacher.py index 77b970a..ac2cebc 100644 --- a/src/curobo/rollout/arm_reacher.py +++ b/src/curobo/rollout/arm_reacher.py @@ -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, diff --git a/src/curobo/rollout/cost/bound_cost.py b/src/curobo/rollout/cost/bound_cost.py index 3524cf0..0c91f75 100644 --- a/src/curobo/rollout/cost/bound_cost.py +++ b/src/curobo/rollout/cost/bound_cost.py @@ -21,6 +21,7 @@ import warp as wp from curobo.cuda_robot_model.types import JointLimits from curobo.types.robot import JointState from curobo.types.tensor import T_DOF +from curobo.util.logger import log_error from curobo.util.torch_utils import get_torch_jit_decorator from curobo.util.warp import init_warp @@ -53,6 +54,16 @@ class BoundCostConfig(CostConfig): self.joint_limits = bounds.clone() if teleport_mode: self.cost_type = BoundCostType.POSITION + if self.cost_type != BoundCostType.POSITION: + if torch.max(self.joint_limits.velocity[1] - self.joint_limits.velocity[0]) == 0.0: + log_error("Joint velocity limits is zero") + if ( + torch.max(self.joint_limits.acceleration[1] - self.joint_limits.acceleration[0]) + == 0.0 + ): + log_error("Joint acceleration limits is zero") + if torch.max(self.joint_limits.jerk[1] - self.joint_limits.jerk[0]) == 0.0: + log_error("Joint jerk limits is zero") def __post_init__(self): if isinstance(self.activation_distance, List): @@ -675,8 +686,14 @@ def forward_bound_pos_warp( target_p = retract_config[target_id] p_l = p_b[d_id] p_u = p_b[dof + d_id] + + p_range = p_u - p_l + eta_p = eta_p * (p_range) p_l += eta_p p_u -= eta_p + if p_range < 1.0: + w = (1.0 / p_range) * w + # compute cost: b_addrs = b_id * horizon * dof + h_id * dof + d_id @@ -690,33 +707,15 @@ def forward_bound_pos_warp( g_p = 2.0 * n_w * error # bound cost: - # if c_p < p_l: - # delta = p_l - c_p - # if (delta) > eta_p or eta_p == 0.0: - # c_total += w * (delta - 0.5 * eta_p) - # g_p += -w - # else: - # c_total += w * (0.5 / eta_p) * delta * delta - # g_p += -w * (1.0 / eta_p) * delta - # elif c_p > p_u: - # delta = c_p - p_u - # if (delta) > eta_p or eta_p == 0.0: - # c_total += w * (delta - 0.5 * eta_p) - # g_p += w - # else: - # c_total += w * (0.5 / eta_p) * delta * delta - # g_p += w * (1.0 / eta_p) * delta - # bound cost: + delta = 0.0 if c_p < p_l: delta = c_p - p_l - c_total += w * delta * delta - g_p += 2.0 * w * delta elif c_p > p_u: delta = c_p - p_u - c_total += w * delta * delta - g_p += 2.0 * w * delta + c_total += w * delta * delta + g_p += 2.0 * w * delta out_cost[b_addrs] = c_total # compute gradient @@ -807,16 +806,33 @@ def forward_bound_warp( c_j = jerk[b_addrs] - p_l = p_b[d_id] + eta_p - p_u = p_b[dof + d_id] - eta_p + p_l = p_b[d_id] + p_u = p_b[dof + d_id] + p_range = p_u - p_l + eta_p = eta_p * (p_range) + p_l += eta_p + p_u -= eta_p - v_l = v_b[d_id] + eta_v - v_u = v_b[dof + d_id] - eta_v - a_l = a_b[d_id] + eta_a - a_u = a_b[dof + d_id] - eta_a + v_l = v_b[d_id] + v_u = v_b[dof + d_id] + v_range = v_u - v_l + eta_v = eta_v * (v_range) + v_l += eta_v + v_u -= eta_v - j_l = j_b[d_id] + eta_j - j_u = j_b[dof + d_id] - eta_j + a_l = a_b[d_id] + a_u = a_b[dof + d_id] + a_range = a_u - a_l + eta_a = eta_a * (a_range) + a_l += eta_a + a_u -= eta_a + + j_l = j_b[d_id] + j_u = j_b[dof + d_id] + j_range = j_u - j_l + eta_j = eta_j * (j_range) + j_l += eta_j + j_u -= eta_j delta = float(0.0) if n_w > 0.0: @@ -826,6 +842,7 @@ def forward_bound_warp( # bound cost: delta = 0.0 + if c_p < p_l: delta = c_p - p_l elif c_p > p_u: @@ -995,19 +1012,47 @@ def forward_bound_smooth_warp( r_wj = run_weight_jerk[h_id] c_j = jerk[b_addrs] - p_l = p_b[d_id] + eta_p - p_u = p_b[dof + d_id] - eta_p + p_l = p_b[d_id] + p_u = p_b[dof + d_id] + p_range = p_u - p_l + eta_p = eta_p * (p_range) + p_l += eta_p + p_u -= eta_p - v_l = v_b[d_id] + eta_v - v_u = v_b[dof + d_id] - eta_v - a_l = a_b[d_id] + eta_a - a_u = a_b[dof + d_id] - eta_a + v_l = v_b[d_id] + v_u = v_b[dof + d_id] + v_range = v_u - v_l + eta_v = eta_v * (v_range) + v_l += eta_v + v_u -= eta_v - j_l = j_b[d_id] + eta_j - j_u = j_b[dof + d_id] - eta_j + a_l = a_b[d_id] + a_u = a_b[dof + d_id] + a_range = a_u - a_l + eta_a = eta_a * (a_range) + a_l += eta_a + a_u -= eta_a + + j_l = j_b[d_id] + j_u = j_b[dof + d_id] + j_range = j_u - j_l + eta_j = eta_j * (j_range) + j_l += eta_j + j_u -= eta_j delta = float(0.0) + if p_range < 1.0: + w = w * (1.0 / (p_range * p_range)) + + if v_range < 1.0: + b_wv = b_wv * (1.0 / (v_range * v_range)) + if a_range < 1.0: + b_wa = b_wa * (1.0 / (a_range * a_range)) + w_a = w_a * (1.0 / (a_range * a_range)) + if j_range < 1.0: + b_wj = b_wj * (1.0 / (j_range * j_range)) + w_j = w_j * (1.0 / (j_range * j_range)) # position: if n_w > 0.0: error = c_p - target_p diff --git a/src/curobo/rollout/cost/self_collision_cost.py b/src/curobo/rollout/cost/self_collision_cost.py index c4469f9..76dc949 100644 --- a/src/curobo/rollout/cost/self_collision_cost.py +++ b/src/curobo/rollout/cost/self_collision_cost.py @@ -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, ) diff --git a/src/curobo/types/math.py b/src/curobo/types/math.py index 98908e6..97f11dd 100644 --- a/src/curobo/types/math.py +++ b/src/curobo/types/math.py @@ -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: diff --git a/src/curobo/types/robot.py b/src/curobo/types/robot.py index 61d1ede..8095430 100644 --- a/src/curobo/types/robot.py +++ b/src/curobo/types/robot.py @@ -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( diff --git a/src/curobo/util/logger.py b/src/curobo/util/logger.py index 9b79a45..917de79 100644 --- a/src/curobo/util/logger.py +++ b/src/curobo/util/logger.py @@ -8,11 +8,26 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # +""" +This module provides logging API, wrapping :py:class:`logging.Logger`. These functions are used +to log messages in the cuRobo package. The functions can also be used in other packages by +creating a new logger (:py:meth:`setup_logger`) with the desired name. +""" # Standard Library import logging +import sys -def setup_curobo_logger(level="info"): +def setup_logger(level="info", logger_name: str = "curobo"): + """Set up logger level. + + Args: + level: Log level. Default is "info". Other options are "debug", "warning", "error". + logger_name: Name of the logger. Default is "curobo". + + Raises: + ValueError: If log level is not one of [info, debug, warning, error]. + """ FORMAT = "[%(levelname)s] [%(name)s] %(message)s" if level == "info": level = logging.INFO @@ -25,21 +40,64 @@ def setup_curobo_logger(level="info"): else: raise ValueError("Log level should be one of [info,debug, warn, error]") logging.basicConfig(format=FORMAT, level=level) - logger = logging.getLogger("curobo") + logger = logging.getLogger(logger_name) logger.setLevel(level=level) -def log_warn(txt: str, *args, **kwargs): - logger = logging.getLogger("curobo") +def setup_curobo_logger(level="info"): + """Set up logger level for curobo package. Deprecated. Use :py:meth:`setup_logger` instead.""" + return setup_logger(level, "curobo") + + +def log_warn(txt: str, logger_name: str = "curobo", *args, **kwargs): + """Log warning message. Also see :py:meth:`logging.Logger.warning`. + + Args: + txt: Warning message. + logger_name: Name of the logger. Default is "curobo". + """ + logger = logging.getLogger(logger_name) logger.warning(txt, *args, **kwargs) -def log_info(txt: str, *args, **kwargs): - logger = logging.getLogger("curobo") +def log_info(txt: str, logger_name: str = "curobo", *args, **kwargs): + """Log info message. Also see :py:meth:`logging.Logger.info`. + + Args: + txt: Info message. + logger_name: Name of the logger. Default is "curobo". + """ + logger = logging.getLogger(logger_name) logger.info(txt, *args, **kwargs) -def log_error(txt: str, exc_info=True, stack_info=True, *args, **kwargs): - logger = logging.getLogger("curobo") - logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs) - raise +def log_error( + txt: str, + logger_name: str = "curobo", + exc_info=True, + stack_info=False, + stacklevel: int = 2, + *args, + **kwargs +): + """Log error and raise ValueError. + + Args: + txt: Helpful message that conveys the error. + logger_name: Name of the logger. Default is "curobo". + exc_info: Add exception info to message. See :py:meth:`logging.Logger.error`. + stack_info: Add stacktracke to message. See :py:meth:`logging.Logger.error`. + stacklevel: See :py:meth:`logging.Logger.error`. Default value of 2 removes this function + from the stack trace. + + Raises: + ValueError: Error message with exception. + """ + logger = logging.getLogger(logger_name) + if sys.version_info.major == 3 and sys.version_info.minor <= 7: + logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs) + else: + logger.error( + txt, exc_info=exc_info, stack_info=stack_info, stacklevel=stacklevel, *args, **kwargs + ) + raise ValueError(txt) diff --git a/src/curobo/util/metrics.py b/src/curobo/util/metrics.py index 208eaf1..9c6a90c 100644 --- a/src/curobo/util/metrics.py +++ b/src/curobo/util/metrics.py @@ -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 diff --git a/src/curobo/util/trajectory.py b/src/curobo/util/trajectory.py index effc1d0..cb8f845 100644 --- a/src/curobo/util/trajectory.py +++ b/src/curobo/util/trajectory.py @@ -11,7 +11,7 @@ # Standard Library import math from enum import Enum -from typing import List, Optional +from typing import List, Optional, Tuple # Third Party import numpy as np @@ -121,44 +121,58 @@ def get_spline_interpolated_trajectory(raw_traj, des_horizon, degree=5): def get_batch_interpolated_trajectory( raw_traj: JointState, + raw_dt: torch.Tensor, interpolation_dt: float, - max_vel: torch.Tensor, - max_acc: torch.Tensor, - max_jerk: torch.Tensor, - raw_dt: float, + max_vel: Optional[torch.Tensor] = None, + max_acc: Optional[torch.Tensor] = None, + max_jerk: Optional[torch.Tensor] = None, kind: InterpolateType = InterpolateType.LINEAR_CUDA, out_traj_state: Optional[JointState] = None, tensor_args: TensorDeviceType = TensorDeviceType(), max_deviation: float = 0.1, min_dt: float = 0.02, + max_dt: float = 0.15, optimize_dt: bool = True, ): # compute dt across trajectory: + if len(raw_traj.shape) == 2: + raw_traj = raw_traj.unsqueeze(0) b, horizon, dof = raw_traj.position.shape # horizon # given the dt required to run trajectory at maximum velocity, # we find the number of timesteps required: - traj_steps, steps_max, opt_dt = calculate_tsteps( - raw_traj.velocity, - raw_traj.acceleration, - raw_traj.jerk, - interpolation_dt, - max_vel, - max_acc, - max_jerk, - raw_dt, - min_dt, - horizon, - optimize_dt, - ) + if optimize_dt: + traj_steps, steps_max, opt_dt = calculate_tsteps( + raw_traj.velocity, + raw_traj.acceleration, + raw_traj.jerk, + interpolation_dt, + max_vel, + max_acc, + max_jerk, + raw_dt, + min_dt, + max_dt, + horizon, + optimize_dt, + ) + else: + traj_steps, steps_max = calculate_traj_steps(raw_dt, interpolation_dt, horizon) + opt_dt = raw_dt # traj_steps contains the tsteps for each trajectory - assert steps_max > 0 - # To do linear interpolation, we + if steps_max <= 0: + log_error("Steps max is less than 0") + + if out_traj_state is not None and out_traj_state.position.shape[1] < steps_max: + log_warn( + "Interpolation buffer shape is smaller than steps_max," + + " creating new buffer of shape " + + str(steps_max) + ) + out_traj_state = None + if out_traj_state is None: out_traj_state = JointState.zeros([b, steps_max, dof], tensor_args) - if out_traj_state.position.shape[1] < steps_max: - log_error("Interpolation buffer shape is smaller than steps_max") - if kind in [InterpolateType.LINEAR, InterpolateType.CUBIC]: # plot and save: out_traj_state = get_cpu_linear_interpolation( @@ -187,7 +201,7 @@ def get_batch_interpolated_trajectory( ) else: raise ValueError("Unknown interpolation type") - # opt_dt = (raw_dt) / opt_dt + return out_traj_state, traj_steps, opt_dt @@ -448,7 +462,9 @@ def calculate_dt_fixed( max_acc: torch.Tensor, max_jerk: torch.Tensor, raw_dt: torch.Tensor, - interpolation_dt: float, + min_dt: float, + max_dt: float, + epsilon: float = 1e-6, ): # compute scaled dt: max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof @@ -465,9 +481,9 @@ def calculate_dt_fixed( dt_score_jerk = raw_dt * torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3) dt_score = torch.maximum(dt_score_vel, dt_score_acc) dt_score = torch.maximum(dt_score, dt_score_jerk) - dt_score = torch.clamp(dt_score, interpolation_dt, raw_dt) - # NOTE: this dt score is not dt, rather a scaling to convert velocity, acc, jerk that was - # computed with raw_dt to a new dt + + dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, max_dt) + return dt_score @@ -480,7 +496,8 @@ def calculate_dt( max_acc: torch.Tensor, max_jerk: torch.Tensor, raw_dt: float, - interpolation_dt: float, + min_dt: float, + epsilon: float = 1e-6, ): # compute scaled dt: max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof @@ -497,7 +514,8 @@ def calculate_dt( dt_score_jerk = raw_dt * torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3) dt_score = torch.maximum(dt_score_vel, dt_score_acc) dt_score = torch.maximum(dt_score, dt_score_jerk) - dt_score = torch.clamp(dt_score, interpolation_dt, raw_dt) + dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, raw_dt) + # NOTE: this dt score is not dt, rather a scaling to convert velocity, acc, jerk that was # computed with raw_dt to a new dt return dt_score @@ -511,6 +529,7 @@ def calculate_dt_no_clamp( max_vel: torch.Tensor, max_acc: torch.Tensor, max_jerk: torch.Tensor, + epsilon: float = 1e-6, ): # compute scaled dt: max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof @@ -527,9 +546,19 @@ def calculate_dt_no_clamp( dt_score_jerk = torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3) dt_score = torch.maximum(dt_score_vel, dt_score_acc) dt_score = torch.maximum(dt_score, dt_score_jerk) + dt_score = dt_score * (1.0 + epsilon) return dt_score +@get_torch_jit_decorator() +def calculate_traj_steps( + opt_dt: torch.Tensor, interpolation_dt: float, horizon: int +) -> Tuple[torch.Tensor, torch.Tensor]: + traj_steps = (torch.ceil((horizon - 1) * ((opt_dt) / interpolation_dt))).to(dtype=torch.int32) + steps_max = torch.max(traj_steps) + return traj_steps, steps_max + + @get_torch_jit_decorator() def calculate_tsteps( vel: torch.Tensor, @@ -541,15 +570,23 @@ def calculate_tsteps( max_jerk: torch.Tensor, raw_dt: torch.Tensor, min_dt: float, + max_dt: float, horizon: int, optimize_dt: bool = True, ): # compute scaled dt: opt_dt = calculate_dt_fixed( - vel, acc, jerk, max_vel, max_acc, max_jerk, raw_dt, interpolation_dt + vel, + acc, + jerk, + max_vel, + max_acc, + max_jerk, + raw_dt, + min_dt, + max_dt, ) if not optimize_dt: opt_dt[:] = raw_dt - traj_steps = (torch.ceil((horizon - 1) * ((opt_dt) / interpolation_dt))).to(dtype=torch.int32) - steps_max = torch.max(traj_steps) + traj_steps, steps_max = calculate_traj_steps(opt_dt, interpolation_dt, horizon) return traj_steps, steps_max, opt_dt diff --git a/src/curobo/wrap/model/robot_segmenter.py b/src/curobo/wrap/model/robot_segmenter.py index 173ff03..2948765 100644 --- a/src/curobo/wrap/model/robot_segmenter.py +++ b/src/curobo/wrap/model/robot_segmenter.py @@ -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: diff --git a/src/curobo/wrap/model/robot_world.py b/src/curobo/wrap/model/robot_world.py index 5900712..5073eb0 100644 --- a/src/curobo/wrap/model/robot_world.py +++ b/src/curobo/wrap/model/robot_world.py @@ -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 diff --git a/src/curobo/wrap/reacher/evaluator.py b/src/curobo/wrap/reacher/evaluator.py index 2dfc351..863b6a9 100644 --- a/src/curobo/wrap/reacher/evaluator.py +++ b/src/curobo/wrap/reacher/evaluator.py @@ -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 diff --git a/src/curobo/wrap/reacher/ik_solver.py b/src/curobo/wrap/reacher/ik_solver.py index 0e7120a..23ffd8e 100644 --- a/src/curobo/wrap/reacher/ik_solver.py +++ b/src/curobo/wrap/reacher/ik_solver.py @@ -9,9 +9,10 @@ # its affiliates is strictly prohibited. # -""" -This module contains :meth:`IKSolver` to solve inverse kinematics, along with -:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result. +""" +This module contains :meth:`IKSolver` to solve inverse kinematics, along with +:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result. A minimal +example to solve IK is available at :ref:`python_ik_example`. .. raw:: html @@ -20,6 +21,7 @@ This module contains :meth:`IKSolver` to solve inverse kinematics, along with

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.""" diff --git a/src/curobo/wrap/reacher/motion_gen.py b/src/curobo/wrap/reacher/motion_gen.py index 1468a5c..2edf67c 100644 --- a/src/curobo/wrap/reacher/motion_gen.py +++ b/src/curobo/wrap/reacher/motion_gen.py @@ -9,20 +9,17 @@ # its affiliates is strictly prohibited. # -""" -This module contains :meth:`MotionGen` class that provides a high-level interface for motion -generation. Motion Generation can take goals either as joint configurations +""" +This module contains :meth:`MotionGen` class that provides a high-level interface for motion +generation. Motion Generation can take goals either as joint configurations :meth:`MotionGen.plan_single_js` or as Cartesian poses :meth:`MotionGen.plan_single`. When Cartesian -pose is given as target, inverse kinematics is first done to generate seeds for trajectory -optimization. Motion generation fallback to using a graph planner when linear interpolated +pose is given as target, inverse kinematics is first done to generate seeds for trajectory +optimization. Motion generation fallback to using a graph planner when linear interpolated trajectory optimization seeds are not successful. Reaching one Cartesian pose in a goalset is also supported using :meth:`MotionGen.plan_goalset`. Batched planning in the same world environment ( -:meth:`MotionGen.plan_batch`) and different world environments (:meth:`MotionGen.plan_batch_env`) +:meth:`MotionGen.plan_batch`) and different world environments (:meth:`MotionGen.plan_batch_env`) is also provided. -A motion generation request can be configured using :meth:`MotionGenPlanConfig`. The result -of motion generation is returned as a :meth:`MotionGenResult`. - .. raw:: html @@ -30,6 +27,11 @@ of motion generation is returned as a :meth:`MotionGenResult`.

+ +A motion generation request can be configured using :meth:`MotionGenPlanConfig`. The result +of motion generation is returned as a :meth:`MotionGenResult`. A minimal example is availble at +:ref:`python_motion_gen_example`. + """ from __future__ import annotations @@ -64,10 +66,10 @@ from curobo.types.camera import CameraObservation from curobo.types.math import Pose from curobo.types.robot import RobotConfig from curobo.types.state import JointState -from curobo.types.tensor import T_BDOF, T_BValue_bool, T_BValue_float +from curobo.types.tensor import T_BDOF, T_DOF, T_BValue_float from curobo.util.logger import log_error, log_info, log_warn from curobo.util.tensor_util import tensor_repeat_seeds -from curobo.util.trajectory import InterpolateType +from curobo.util.trajectory import InterpolateType, get_batch_interpolated_trajectory from curobo.util.warp import init_warp from curobo.util_file import ( get_robot_configs_path, @@ -78,7 +80,7 @@ from curobo.util_file import ( ) from curobo.wrap.reacher.evaluator import TrajEvaluator, TrajEvaluatorConfig from curobo.wrap.reacher.ik_solver import IKResult, IKSolver, IKSolverConfig -from curobo.wrap.reacher.trajopt import TrajOptSolver, TrajOptSolverConfig +from curobo.wrap.reacher.trajopt import TrajOptResult, TrajOptSolver, TrajOptSolverConfig from curobo.wrap.reacher.types import ReacherSolveState, ReacherSolveType @@ -159,14 +161,14 @@ class MotionGenConfig: #: tried again until :attr:`MotionGenPlanConfig.finetune_attempts`. optimize_dt: bool = True - @staticmethod @profiler.record_function("motion_gen_config/load_from_robot_config") + @staticmethod def load_from_robot_config( robot_cfg: Union[Union[str, Dict], RobotConfig], world_model: Optional[Union[Union[str, Dict], WorldConfig]] = None, tensor_args: TensorDeviceType = TensorDeviceType(), num_ik_seeds: int = 32, - num_graph_seeds: int = 1, + num_graph_seeds: int = 4, num_trajopt_seeds: int = 4, num_batch_ik_seeds: int = 32, num_batch_trajopt_seeds: int = 1, @@ -228,8 +230,9 @@ class MotionGenConfig: finetune_smooth_weight: Optional[List[float]] = None, state_finite_difference_mode: Optional[str] = None, finetune_dt_scale: float = 0.9, + minimum_trajectory_dt: Optional[float] = None, maximum_trajectory_time: Optional[float] = None, - maximum_trajectory_dt: float = 0.1, + maximum_trajectory_dt: Optional[float] = None, velocity_scale: Optional[Union[List[float], float]] = None, acceleration_scale: Optional[Union[List[float], float]] = None, jerk_scale: Optional[Union[List[float], float]] = None, @@ -239,13 +242,249 @@ class MotionGenConfig: graph_seed: int = 1531, high_precision: bool = False, ): + """Create a motion generation configuration from robot and world configuration. + + Args: + robot_cfg: Robot configuration to use for motion generation. This can be a path to a + yaml file, a dictionary, or an instance of :class:`RobotConfig`. See + :ref:`available_robot_list` for a list of available robots. You can also create a + a configuration file for your robot using :ref:`tut_robot_configuration`. + world_model: World configuration to use for motion generation. This can be a path to a + yaml file, a dictionary, or an instance of :class:`WorldConfig`. See + :ref:`world_collision` for more details. + tensor_args: Numerical precision and compute device to use for motion generation. + num_ik_seeds: Number of seeds to use for solving inverse kinematics. Default of 32 is + found to be a good number for most cases. In sparse environments, a lower number of + 16 can also be used. + num_graph_seeds: Number of seeds to use for graph planner per problem query. When graph + planning is used to generate seeds for trajectory optimization, graph planner will + attempt to find collision-free paths from the start state to the many inverse + kinematics solutions. + num_trajopt_seeds: Number of seeds to use for trajectory optimization per problem + query. Default of 4 is found to be a good number for most cases. Increasing this + will increase memory usage. + num_batch_ik_seeds: Number of seeds to use for inverse kinematics during batched + planning. Default of 32 is found to be a good number for most cases. + num_batch_trajopt_seeds: Number of seeds to use for trajectory optimization during + batched planning. Using more than 1 will disable graph planning for batched + planning. + num_trajopt_noisy_seeds: Number of augmented trajectories to use per trajectory seed. + The augmentation is done by adding random noise to the trajectory. This + augmentation has not been found to be useful and it's recommended to keep this to + 1. The noisy seeds can also be used in conjunction with the trajopt_seed_ratio to + generate seeds that go through a bias point. + position_threshold: Position threshold in meters between reached position and target + position used to measure success. + rotation_threshold: Rotation threshold between reached orientation and target + orientation used to measure success. The metric is q^T * q, where q is the + quaternion difference between target and reached orientation. The metric is not + easy to interpret and a future release will provide a more intuitive metric. For + now, use 0.05 as a good value. + cspace_threshold: Joint space threshold in radians for revolute joints and meters for + linear joints between reached joint configuration and target joint configuration + used to measure success. Default of 0.05 has been found to be a good value for most + cases. + world_coll_checker: Instance of world collision checker to use for motion generation. + Leaving this to None will create a new instance of world collision checker using + the provided attr:`world_model`. + base_cfg_file: Base configuration file containing convergence and constraint criteria + to measure success. + particle_ik_file: Optimizer configuration file to use for particle-based optimization + during inverse kinematics. + gradient_ik_file: Optimizer configuration file to use for gradient-based optimization + during inverse kinematics. + graph_file: Configuration file to use for graph planner. + particle_trajopt_file: Optimizer configuration file to use for particle-based + optimization during trajectory optimization. + gradient_trajopt_file: Optimizer configuration file to use for gradient-based + optimization during trajectory optimization. + finetune_trajopt_file: Optimizer configuration file to use for finetuning trajectory + optimization. + trajopt_tsteps: Number of waypoints to use for trajectory optimization. Default of 32 + is found to be a good number for most cases. + interpolation_steps: Buffer size to use for storing interpolated trajectory. Default of + 5000 is found to be a good number for most cases. + interpolation_dt: Time step in seconds to use for generating interpolated trajectory + from optimized trajectory. Change this if you want to generate a trajectory with + a fixed timestep between waypoints. + interpolation_type: Interpolation type to use for generating dense waypoints from + optimized trajectory. Default of + :py:attr:`curobo.util.trajectory.InterpolateType.LINEAR_CUDA` is found to be a + good choice for most cases. Other suitable options for real robot execution are + :py:attr:`curobo.util.trajectory.InterpolateType.QUINTIC` and + :py:attr:`curobo.util.trajectory.InterpolateType.CUBIC`. + use_cuda_graph: Record compute ops as cuda graphs and replay recorded graphs where + implemented. This can speed up execution by upto 10x. Default of True is + recommended. Enabling this will prevent changing problem type or batch size + after the first call to the solver. + self_collision_check: Enable self collision checks for generated motions. Default of + True is recommended. Set this to False to debug planning failures. Setting this to + False will also set self_collision_opt to False. + self_collision_opt: Enable self collision cost during optimization (IK, TrajOpt). + Default of True is recommended. + grad_trajopt_iters: Number of iterations to run trajectory optimization. + trajopt_seed_ratio: Ratio of linear and bias seeds to use for trajectory optimization. + Linear seed will generate a linear interpolated trajectory from start state + to IK solutions. Bias seed will add a mid-waypoint through the retract + configuration. Default of 1.0 linear and 0.0 bias is recommended. This can be + changed to 0.5 linear and 0.5 bias, along with changing trajopt_noisy_seeds to 2. + ik_opt_iters: Number of iterations to run inverse kinematics. + ik_particle_opt: Enable particle-based optimization during inverse kinematics. Default + of True is recommended as particle-based optimization moves the random seeds to + a regions of local minima. + collision_checker_type: Type of collision checker to use for motion generation. Default + of CollisionCheckerType.MESH supports world represented by Cuboids and Meshes. See + :ref:`world_collision` for more details. + sync_cuda_time: Synchronize with host using :py:func:`torch.cuda.synchronize` before + measuring compute time. + trajopt_particle_opt: Enable particle-based optimization during trajectory + optimization. Default of True is recommended as particle-based optimization moves + the interpolated seeds away from bad local minima. + traj_evaluator_config: Configuration for trajectory evaluator. Default of None will + create a new instance of TrajEvaluatorConfig. After trajectory optimization across + many seeds, the best trajectory is selected based on this configuration. This + evaluator also checks if the optimized dt is within + :py:attr:`curobo.wrap.reacher.evaluator.TrajEvaluatorConfig.max_dt`. This check is + needed to measure smoothness of the optimized trajectory as bad trajectories can + have very high dt to fit within velocity, acceleration, and jerk limits. + traj_evaluator: Instance of trajectory evaluator to use for trajectory optimization. + Default of None will create a new instance of TrajEvaluator. In case you want to + use a custom evaluator, you can create a child instance of TrajEvaluator and + pass it. + minimize_jerk: Minimize jerk as regularization during trajectory optimizaiton. + filter_robot_command: Filter generated trajectory to remove finite difference + artifacts. Default of True is recommended. + use_gradient_descent: Use gradient descent instead of L-BFGS for trajectory + optimization. + 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:`MotionGen.plan_batch_env` + and :py:meth:`MotionGen.plan_batch_env_goalset`. + ee_link_name: End effector link/frame to use for reaching Cartesian poses. Default of + None will use the end effector link from the robot configuration. This cannot + be changed after creating the robot configuration. + use_es_ik: Use evolutionary strategy for as the particle-based optimizer for inverse + kinematics. Default of None will use MPPI as the optimization algorithm. ES is not + recommended as it's unstable and provided optimization parameters were not tuned. + use_es_trajopt: Use evolutionary strategy as the particle-based optimizer for + trajectory optimization. Default of None will use MPPI as the optimization + algorithm. ES is not recommended as it's unstable and provided optimization + parameters were not tuned. + es_ik_learning_rate: Learning rate to use for evolutionary strategy in IK. + es_trajopt_learning_rate: Learning rate to use for evolutionary strategy in TrajOpt. + use_ik_fixed_samples: Use fixed samples of noise during particle-based optimization + of IK. Default of None will use the setting from the optimizer configuration file + (``particle_ik.yml``). + use_trajopt_fixed_samples: Use fixed samples of noise during particle-based + optimization of trajectory optimization. Default of None will use the setting from + the optimizer configuration file (``particle_trajopt.yml``). + evaluate_interpolated_trajectory: Evaluate interpolated trajectory after optimization. + Default of True is recommended to ensure the optimized trajectory is not passing + through very thin obstacles. + partial_ik_iters: Number of iterations of L-BFGS to run inverse kinematics when + only partial IK is needed. + fixed_iters_trajopt: Use fixed number of iterations of L-BFGS for trajectory + optimization. Default of None will use the setting from the optimizer + configuration. In most cases, fixed iterations of solvers are run as current + solvers treat constraints as costs and there is no guarantee that the constraints + will be satisfied. Instead of checking constraints between iterations of a solver + and exiting, it's computationally cheaper to run a fixed number of iterations. In + addition, running fixed iterations of solvers is more robust to outlier problems. + store_ik_debug: Store debugging information such as values of optimization variables + at every iteration in IK result. Setting this to True will set + :attr:`use_cuda_graph` to False. + store_trajopt_debug: Store debugging information such as values of optimization + variables in TrajOpt result. Setting this to True will set :attr:`use_cuda_graph` + to False. + graph_trajopt_iters: Number of iterations to run trajectory optimization when seeded + from a graph plan. Default of None will use the same number of iterations as + linear seeded trajectory optimization. This can be set to a higher value of 200 + in case where trajectories obtained are not of requird quality. + collision_max_outside_distance: Maximum distance to check for collision outside a + obstacle. Increasing this value will slow down collision checks with Meshes as + closest point queries will be run up to this distance outside an obstacle. + 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. + trajopt_dt: Time step in seconds to use for trajectory optimization. A good value to + start with is 0.15 seconds. This value is used to compute velocity, acceleration, + and jerk values for waypoints through finite difference. + js_trajopt_dt: Time step in seconds to use for trajectory optimization when reaching + joint space targets. A good value to start with is 0.15 seconds. This value is used + to compute velocity, acceleration, and jerk values for waypoints through finite + difference. + js_trajopt_tsteps: Number of waypoints to use for trajectory optimization when reaching + joint space targets. Default of None will use the same number of waypoints as + Cartesian trajectory optimization. + trim_steps: Trim waypoints from optimized trajectory. The optimized trajectory will + contain the start state at index 0 and have the last two waypoints be the same + as T-2 as trajectory optimization implicitly optimizes for zero acceleration and + velocity at the last waypoint. An example: ``[1,-2]`` will trim the first waypoint + and last 3 waypoints from the optimized trajectory. + store_debug_in_result: Store debugging information in MotionGenResult. This value is + set to True if either store_ik_debug or store_trajopt_debug is set to True. + finetune_trajopt_iters: Number of iterations to run trajectory optimization for + finetuning after an initial collision-free trajectory is obtained. + smooth_weight: Override smooth weight for trajectory optimization. It's not recommended + to set this value for most cases. + finetune_smooth_weight: Override smooth weight for finetuning trajectory optimization. + It's not recommended to set this value for most cases. + state_finite_difference_mode: Finite difference mode to use for computing velocity, + acceleration, and jerk values. Default of None will use the setting from the + optimizer configuration file. The default finite difference method is a five + point stencil to compute the derivatives as this is accurate and provides + faster convergence compared to backward or central difference methods. + finetune_dt_scale: Scale initial estimated dt by this value to finetune trajectory + optimization. This is deprecated and will be removed in future releases. Use + :py:attr:`MotionGenPlanConfig.finetune_dt_scale` instead. + minimum_trajectory_dt: Minimum time step in seconds allowed for trajectory + optimization. + maximum_trajectory_time: Maximum time in seconds allowed for trajectory optimization. + maximum_trajectory_dt: Maximum time step in seconds allowed for trajectory + optimization. + velocity_scale: Scale velocity limits by this value. Default of None will not scale + the velocity limits. To generate slower trajectories, use + :py:attr:`MotionGenPlanConfig.time_dilation_factor` < 1.0 instead. Changing this + value is not recommended as it changes the scale of cost terms and they would + require retuning. + acceleration_scale: Scale acceleration limits by this value. Default of None will not + scale the acceleration limits. To generate slower trajectories, use + :py:attr:`MotionGenPlanConfig.time_dilation_factor` < 1.0 instead. Changing this + value is not recommended as it changes the scale of cost terms and they would + require retuning. + jerk_scale: Scale jerk limits by this value. Default of None will not scale the jerk + limits. To generate slower trajectories, use + :py:attr:`MotionGenPlanConfig.time_dilation_factor` < 1.0 instead. Changing this + value is not recommended as it changes the scale of cost terms and they would + require retuning. + optimize_dt: Optimize dt during trajectory optimization. Default of True is + recommended to find time-optimal trajectories. Setting this to False will use the + provided :attr:`trajopt_dt` for trajectory optimization. Setting to False is + required when optimizing from a non-static start state. + project_pose_to_goal_frame: Project pose to goal frame when calculating distance + between reached and goal pose. Use this to constrain motion to specific axes + either in the global frame or the goal frame. + ik_seed: Random seed to use for inverse kinematics. + graph_seed: Random seed to use for graph planner. + high_precision: Use high precision settings for motion generation. This will increase + the number of iterations for optimization solvers and reduce the thresholds for + position to 1mm and rotation to 0.025. Default of False is recommended for most + cases as standard motion generation settings reach within 0.5mm on most problems. + + Returns: + MotionGenConfig: Instance of motion generation configuration. + """ if position_threshold <= 0.001: high_precision = True if high_precision: finetune_trajopt_iters = ( 300 if finetune_trajopt_iters is None else max(300, finetune_trajopt_iters) ) - grad_trajopt_iters = 200 if grad_trajopt_iters is None else max(200, grad_trajopt_iters) + if grad_trajopt_iters is None: + grad_trajopt_iters = 200 + grad_trajopt_iters = max(200, grad_trajopt_iters) position_threshold = min(position_threshold, 0.001) rotation_threshold = min(rotation_threshold, 0.025) cspace_threshold = min(cspace_threshold, 0.01) @@ -256,28 +495,55 @@ class MotionGenConfig: if trajopt_tsteps is not None: js_trajopt_tsteps = trajopt_tsteps if velocity_scale is not None and isinstance(velocity_scale, float): + log_warn( + "To slow down trajectories, use MotionGenPlanConfig.time_dilation_factor" + + " instead of velocity_scale" + ) velocity_scale = [velocity_scale] if acceleration_scale is not None and isinstance(acceleration_scale, float): + log_warn( + "To slow down trajectories, use MotionGenPlanConfig.time_dilation_factor" + + " instead of acceleration_scale" + ) acceleration_scale = [acceleration_scale] if jerk_scale is not None and isinstance(jerk_scale, float): jerk_scale = [jerk_scale] if store_ik_debug or store_trajopt_debug: store_debug_in_result = True + + if ( + velocity_scale is not None + and min(velocity_scale) < 0.1 + and finetune_trajopt_file is None + and maximum_trajectory_dt is None + ): + log_error( + "velocity scale<0.1 requires a user determined maximum_trajectory_dt as" + + " default scaling will likely fail. A good value to start with would be 30" + + " seconds" + ) + + if maximum_trajectory_dt is None: + maximum_trajectory_dt = 0.15 + maximum_trajectory_dt_acc = maximum_trajectory_dt + maximum_trajectory_dt_vel = maximum_trajectory_dt if ( acceleration_scale is not None and min(acceleration_scale) < 1.0 - and maximum_trajectory_dt <= 0.1 + and maximum_trajectory_dt <= 0.2 ): - maximum_trajectory_dt = np.sqrt(1.0 / min(acceleration_scale)) * maximum_trajectory_dt - elif ( + maximum_trajectory_dt_acc = ( + np.sqrt(1.0 / min(acceleration_scale)) * maximum_trajectory_dt * 3 + ) + if ( velocity_scale is not None and min(velocity_scale) < 1.0 - and maximum_trajectory_dt <= 0.1 + and maximum_trajectory_dt <= 0.2 ): - maximum_trajectory_dt = (1.0 / min(velocity_scale)) * maximum_trajectory_dt - + maximum_trajectory_dt_vel = (1.0 / min(velocity_scale)) * maximum_trajectory_dt * 3 + maximum_trajectory_dt = max(maximum_trajectory_dt_acc, maximum_trajectory_dt_vel) if maximum_trajectory_dt is not None: if trajopt_dt is None: trajopt_dt = maximum_trajectory_dt @@ -288,16 +554,6 @@ class MotionGenConfig: if velocity_scale is not None and min(velocity_scale) < 0.5: fixed_iters_trajopt = True - if ( - velocity_scale is not None - and min(velocity_scale) < 0.1 - and finetune_trajopt_file is None - ): - log_error( - "velocity scale<0.1 is not supported with default finetune_trajopt.yml " - + "provide your own finetune_trajopt_file to override this error" - ) - if ( velocity_scale is not None and min(velocity_scale) <= 0.25 @@ -338,7 +594,10 @@ class MotionGenConfig: robot_cfg["kinematics"]["cspace"]["velocity_scale"] = velocity_scale robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args) - + if minimum_trajectory_dt is None: + minimum_trajectory_dt = interpolation_dt + elif minimum_trajectory_dt < interpolation_dt: + log_error("minimum_trajectory_dt cannot be lower than interpolation_dt") if traj_evaluator_config is None: if maximum_trajectory_dt is not None: max_dt = maximum_trajectory_dt @@ -347,7 +606,7 @@ class MotionGenConfig: if acceleration_scale is not None: max_dt = max_dt * (1.0 / np.sqrt(min(acceleration_scale))) traj_evaluator_config = TrajEvaluatorConfig.from_basic( - min_dt=interpolation_dt, max_dt=max_dt, dof=robot_cfg.kinematics.dof + min_dt=minimum_trajectory_dt, max_dt=max_dt, dof=robot_cfg.kinematics.dof ) traj_evaluator_config.max_acc = robot_cfg.kinematics.get_joint_limits().acceleration[1] @@ -362,7 +621,8 @@ class MotionGenConfig: if n_collision_envs is not None: base_config_data["world_collision_checker_cfg"]["n_envs"] = n_collision_envs if collision_max_outside_distance is not None: - assert collision_max_outside_distance >= 0.0 + if collision_max_outside_distance < 0.0: + log_error("collision_max_outside_distance cannot be negative") base_config_data["world_collision_checker_cfg"][ "max_distance" ] = collision_max_outside_distance @@ -394,7 +654,6 @@ class MotionGenConfig: grad_iters=ik_opt_iters, use_particle_opt=ik_particle_opt, sync_cuda_time=sync_cuda_time, - # use_gradient_descent=use_gradient_descent, use_es=use_es_ik, es_learning_rate=es_ik_learning_rate, use_fixed_samples=use_ik_fixed_samples, @@ -529,7 +788,7 @@ class MotionGenConfig: use_particle_opt=False, traj_evaluator_config=traj_evaluator_config, traj_evaluator=traj_evaluator, - evaluate_interpolated_trajectory=True, + evaluate_interpolated_trajectory=evaluate_interpolated_trajectory, fixed_iters=fixed_iters_trajopt, store_debug=store_trajopt_debug, collision_activation_distance=collision_activation_distance, @@ -653,7 +912,7 @@ class MotionGenPlanConfig: #: Run finetuning trajectory optimization after running 100 iterations of #: trajectory optimization. This will provide shorter and smoother trajectories. When - # :attr:`MotionGenConfig.optimize_dt` is True, this flag will also scale the trajectory + #: :attr:`MotionGenConfig.optimize_dt` is True, this flag will also scale the trajectory #: optimization by a new dt. Leave this to True for most cases. If you are not interested in #: finding time-optimal solutions and only want to use motion generation as a feasibility check, #: set this to False. Note that when set to False, the resulting trajectory is only guaranteed @@ -681,11 +940,21 @@ class MotionGenPlanConfig: #: fails to find a solution. This is only used when :attr:`MotionGenConfig.optimize_dt` is True. finetune_dt_decay: float = 1.025 + #: Slow down optimized trajectory by re-timing with a dilation factor. This is useful to + #: execute trajectories at a slower speed for debugging. Use this to generate slower + #: trajectories instead of reducing :attr:`MotionGenConfig.velocity_scale` or + #: :attr:`MotionGenConfig.acceleration_scale` as those parameters will require re-tuning + #: of the cost terms while :attr:`MotionGenPlanConfig.time_dilation_factor` will only + #: post-process the trajectory. + time_dilation_factor: Optional[float] = None + def __post_init__(self): + """Post initialization checks.""" if not self.enable_opt and not self.enable_graph: log_error("Graph search and Optimization are both disabled, enable one") def clone(self) -> MotionGenPlanConfig: + """Clone the current planning configuration.""" return MotionGenPlanConfig( enable_graph=self.enable_graph, enable_opt=self.enable_opt, @@ -710,6 +979,7 @@ class MotionGenPlanConfig: ), finetune_dt_scale=self.finetune_dt_scale, finetune_attempts=self.finetune_attempts, + time_dilation_factor=self.time_dilation_factor, ) @@ -814,6 +1084,7 @@ class MotionGenResult: goalset_index: Optional[torch.Tensor] = None def clone(self): + """Clone the current result.""" m = MotionGenResult( self.success.clone(), valid_query=self.valid_query, @@ -842,29 +1113,16 @@ class MotionGenResult: ) return m - @staticmethod - def _check_none_and_copy_idx( - current_tensor: Union[torch.Tensor, JointState, None], - source_tensor: Union[torch.Tensor, JointState, None], - idx: int, - ): - if source_tensor is not None: - if current_tensor is None: - current_tensor = source_tensor.clone() - else: - if isinstance(current_tensor, torch.Tensor) and isinstance( - source_tensor, torch.Tensor - ): - current_tensor[idx] = source_tensor[idx] - elif isinstance(current_tensor, JointState) and isinstance( - source_tensor, JointState - ): - source_state = source_tensor[idx] - current_tensor.copy_at_index(source_state, idx) - - return current_tensor - def copy_idx(self, idx: torch.Tensor, source_result: MotionGenResult): + """Copy data from source result to current result at index. + + Args: + idx: index to copy data at. + source_result: source result to copy data from. + + Returns: + MotionGenResult: copied result. + """ idx = idx.to(dtype=torch.long) # copy data from source result: self.success[idx] = source_result.success[idx] @@ -906,6 +1164,15 @@ class MotionGenResult: return self def get_paths(self) -> List[JointState]: + """Get interpolated trajectories from the result. Use for batched queries. + + This will return unsuccessful trajectories as well. Use + :meth:`MotionGenResult.get_successful_paths` to get only successful trajectories. + + Returns: + List[JointState]: Interpolated trajectories. Check + :attr:`MotionGenResult.interpolation_dt` for the time between steps. + """ path = [ self.interpolated_plan[x].trim_trajectory(0, self.path_buffer_last_tstep[x]) for x in range(len(self.interpolated_plan)) @@ -913,6 +1180,12 @@ class MotionGenResult: return path def get_successful_paths(self) -> List[torch.Tensor]: + """Get successful interpolated trajectories from the result. Use for batched queries. + + Returns: + List[JointState]: Interpolated trajectories. Check + :attr:`MotionGenResult.interpolation_dt` for the time between steps. + """ path = [] nz_i = torch.nonzero(self.success.view(-1)).view(-1) path = self.interpolated_plan[nz_i] @@ -928,26 +1201,133 @@ class MotionGenResult: return path_list def get_interpolated_plan(self) -> JointState: + """Get interpolated trajectory from the result. + + Returns: + JointState: Interpolated trajectory. Check :attr:`MotionGenResult.interpolation_dt` for + the time between steps. + """ if self.path_buffer_last_tstep is None: return self.interpolated_plan if len(self.path_buffer_last_tstep) > 1: - raise ValueError("only single result is supported") + log_error("only single result is supported") return self.interpolated_plan.trim_trajectory(0, self.path_buffer_last_tstep[0]) + def retime_trajectory( + self, + time_dilation_factor: float, + interpolate_trajectory: bool = True, + interpolation_dt: Optional[float] = None, + interpolation_kind: InterpolateType = InterpolateType.LINEAR_CUDA, + create_interpolation_buffer: bool = False, + ): + """Retime the optimized trajectory by a dilation factor. + + Args: + time_dilation_factor: Time factor to slow down the trajectory. Should be less than 1.0. + interpolate_trajectory: Interpolate the trajectory after retiming. + interpolation_dt: Time between steps in the interpolated trajectory. If None, + :attr:`MotionGenResult.interpolation_dt` is used. + interpolation_kind: Interpolation type to use. + create_interpolation_buffer: Create a new buffer for interpolated trajectory. Set this + to True if existing buffer is not large enough to store new trajectory. + """ + + if time_dilation_factor > 1.0: + log_error("time_dilation_factor should be less than 1.0") + if time_dilation_factor == 1.0: + return + if len(self.path_buffer_last_tstep) > 1: + log_error("only single result is supported") + + new_dt = self.optimized_dt * (1.0 / time_dilation_factor) + if len(self.optimized_plan.shape) == 3: + new_dt = new_dt.view(-1, 1, 1) + else: + new_dt = new_dt.view(-1, 1) + self.optimized_plan = self.optimized_plan.scale_by_dt(self.optimized_dt, new_dt) + self.optimized_dt = new_dt.view(-1) + if interpolate_trajectory: + if interpolation_dt is not None: + self.interpolation_dt = interpolation_dt + self.interpolated_plan, last_tstep, _ = get_batch_interpolated_trajectory( + self.optimized_plan, + self.optimized_dt, + self.interpolation_dt, + kind=interpolation_kind, + out_traj_state=self.interpolated_plan if not create_interpolation_buffer else None, + tensor_args=self.interpolated_plan.tensor_args, + optimize_dt=False, + ) + self.path_buffer_last_tstep = [last_tstep[i] for i in range(len(last_tstep))] + if len(self.optimized_plan.shape) == 2: + self.interpolated_plan = self.interpolated_plan.squeeze(0) + @property - def motion_time(self): + def motion_time(self) -> Union[float, torch.Tensor]: + """Compute motion time in seconds.""" + # -2 as last three timesteps have the same value # 0, 1 also have the same position value. return self.optimized_dt * (self.optimized_plan.position.shape[-2] - 1 - 2 - 1) + @staticmethod + def _check_none_and_copy_idx( + current_tensor: Union[torch.Tensor, JointState, None], + source_tensor: Union[torch.Tensor, JointState, None], + idx: int, + ) -> Union[torch.Tensor, JointState]: + """Helper function to copy data from source tensor to current tensor at index. + + Also supports copying from JointState to JointState. + + Args: + current_tensor: tensor to copy data to. + source_tensor: tensor to copy data from. + idx: index to copy data at. + + Returns: + Union[torch.Tensor, JointState]: copied tensor. + """ + if source_tensor is not None: + if current_tensor is None: + current_tensor = source_tensor.clone() + else: + if isinstance(current_tensor, torch.Tensor) and isinstance( + source_tensor, torch.Tensor + ): + current_tensor[idx] = source_tensor[idx] + elif isinstance(current_tensor, JointState) and isinstance( + source_tensor, JointState + ): + source_state = source_tensor[idx] + current_tensor.copy_at_index(source_state, idx) + + return current_tensor + class MotionGen(MotionGenConfig): + """Motion generation wrapper for generating collision-free trajectories. + + This module provides an interface to generate collision-free trajectories for manipulators. It + supports Cartesian Pose for end-effectors and joint space goals. When a Cartesian Pose is used + as target, it uses batched inverse kinematics to find multiple joint configuration solutions + for the Cartesian Pose and then runs trajectory optimization with a linear interpolation from + start configuration to each of the IK solutions. When trajectory optimization fails, it uses + a graph planner to find collision-free paths to the IK solutions and then uses these paths as + seeds for trajectory optimization. The module also supports batched queries for motion + generation. Use this module as the high-level API for generating collision-free trajectories. + """ + def __init__(self, config: MotionGenConfig): + """Initializes the motion generation module. + + Args: + config: Configuration for motion generation. + """ super().__init__(**vars(config)) - self.rollout_fn = ( - self.graph_planner.safety_rollout_fn - ) # TODO: compute_kinematics fn in rolloutbase + self.rollout_fn = self.graph_planner.safety_rollout_fn self._trajopt_goal_config = None self._dof = self.rollout_fn.d_action self._batch_graph_search_buffer = None @@ -960,6 +1340,12 @@ class MotionGen(MotionGenConfig): self.update_batch_size(seeds=self.trajopt_seeds) def update_batch_size(self, seeds=10, batch=1): + """Update the batch size for motion generation. + + Args: + seeds: Number of seeds for trajectory optimization and graph planner. + batch: Number of queries to run in batch mode. + """ if ( self._trajopt_goal_config is None or self._trajopt_goal_config.shape[0] != batch @@ -994,6 +1380,26 @@ class MotionGen(MotionGenConfig): use_nn_seed: bool = True, newton_iters: Optional[int] = None, ) -> IKResult: + """Solve inverse kinematics for a given Cartesian Pose or a batch of Poses. + + Use this only if your problem size is same as when using motion generation. If you want + to solve IK for a different problem size, create an instance of + :class:`curobo.wrap.reacher.ik_solver.IKSolver`. + + Args: + goal_pose: Goal Pose for the end-effector. + retract_config: Retract configuration for the end-effector. + seed_config: Seed configuration for inverse kinematics. + return_seeds: Number of solutions to return per problem query. + num_seeds: Number of seeds to use for solving inverse kinematics. + use_nn_seed: Use neural network seed for solving inverse kinematics. This is not + implemented. + newton_iters: Number of Newton iterations to run for solving inverse kinematics. + This is useful to override the default number of iterations. + + Returns: + IKResult: Result of inverse kinematics. + """ return self.ik_solver.solve( goal_pose, retract_config, @@ -1004,383 +1410,22 @@ class MotionGen(MotionGenConfig): newton_iters, ) - @profiler.record_function("motion_gen/ik") - def _solve_ik_from_solve_state( - self, - goal_pose: Pose, - solve_state: ReacherSolveState, - start_state: JointState, - use_nn_seed: bool, - partial_ik_opt: bool, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - newton_iters = None - if partial_ik_opt: - newton_iters = self.partial_ik_iters - ik_result = self.ik_solver.solve_any( - solve_state.solve_type, - goal_pose, - start_state.position.view(-1, self._dof), - start_state.position.view(-1, 1, self._dof), - solve_state.num_trajopt_seeds, - solve_state.num_ik_seeds, - use_nn_seed, - newton_iters, - link_poses, - ) - return ik_result - - @profiler.record_function("motion_gen/trajopt_solve") - def _solve_trajopt_from_solve_state( - self, - goal: Goal, - solve_state: ReacherSolveState, - act_seed: Optional[JointState] = None, - use_nn_seed: bool = False, - return_all_solutions: bool = False, - seed_success: Optional[torch.Tensor] = None, - newton_iters: Optional[int] = None, - trajopt_instance: Optional[TrajOptSolver] = None, - num_seeds_override: Optional[int] = None, - ) -> TrajOptResult: - if trajopt_instance is None: - trajopt_instance = self.trajopt_solver - if num_seeds_override is None: - num_seeds_override = solve_state.num_trajopt_seeds - traj_result = trajopt_instance.solve_any( - solve_state.solve_type, - goal, - act_seed, - use_nn_seed, - return_all_solutions, - num_seeds_override, - seed_success, - newton_iters=newton_iters, - ) - return traj_result - @profiler.record_function("motion_gen/graph_search") def graph_search( self, start_config: T_BDOF, goal_config: T_BDOF, interpolation_steps: Optional[int] = None ) -> GraphResult: + """Run graph search to find collision-free paths between start and goal configurations. + + Args: + start_config: Start joint configurations of the robot. + goal_config: Goal joint configurations of the robot. + interpolation_steps: Number of interpolation steps to interpolate obtained solutions. + + Returns: + GraphResult: Result of graph search. + """ return self.graph_planner.find_paths(start_config, goal_config, interpolation_steps) - def _get_solve_state( - self, - solve_type: ReacherSolveType, - plan_config: MotionGenPlanConfig, - goal_pose: Pose, - start_state: JointState, - ): - # TODO: for batch seeds - num_ik_seeds = ( - self.ik_seeds if plan_config.num_ik_seeds is None else plan_config.num_ik_seeds - ) - num_trajopt_seeds = ( - self.trajopt_seeds - if plan_config.num_trajopt_seeds is None - else plan_config.num_trajopt_seeds - ) - - num_graph_seeds = ( - self.trajopt_seeds if plan_config.num_graph_seeds is None else num_trajopt_seeds - ) - solve_state = None - if solve_type == ReacherSolveType.SINGLE: - solve_state = ReacherSolveState( - solve_type, - num_ik_seeds=num_ik_seeds, - num_trajopt_seeds=num_trajopt_seeds, - num_graph_seeds=num_graph_seeds, - batch_size=1, - n_envs=1, - n_goalset=1, - ) - elif solve_type == ReacherSolveType.GOALSET: - solve_state = ReacherSolveState( - solve_type, - num_ik_seeds=num_ik_seeds, - num_trajopt_seeds=num_trajopt_seeds, - num_graph_seeds=num_graph_seeds, - batch_size=1, - n_envs=1, - n_goalset=goal_pose.n_goalset, - ) - elif solve_type == ReacherSolveType.BATCH: - solve_state = ReacherSolveState( - solve_type, - num_ik_seeds=num_ik_seeds, - num_trajopt_seeds=num_trajopt_seeds, - num_graph_seeds=num_graph_seeds, - batch_size=goal_pose.batch, - n_envs=1, - n_goalset=1, - ) - elif solve_type == ReacherSolveType.BATCH_GOALSET: - solve_state = ReacherSolveState( - solve_type, - num_ik_seeds=num_ik_seeds, - num_trajopt_seeds=num_trajopt_seeds, - num_graph_seeds=num_graph_seeds, - batch_size=goal_pose.batch, - n_envs=1, - n_goalset=goal_pose.n_goalset, - ) - elif solve_type == ReacherSolveType.BATCH_ENV: - solve_state = ReacherSolveState( - solve_type, - num_ik_seeds=num_ik_seeds, - num_trajopt_seeds=num_trajopt_seeds, - num_graph_seeds=num_graph_seeds, - batch_size=goal_pose.batch, - n_envs=goal_pose.batch, - n_goalset=1, - ) - elif solve_type == ReacherSolveType.BATCH_ENV_GOALSET: - solve_state = ReacherSolveState( - solve_type, - num_ik_seeds=num_ik_seeds, - num_trajopt_seeds=num_trajopt_seeds, - num_graph_seeds=num_graph_seeds, - batch_size=goal_pose.batch, - n_envs=goal_pose.batch, - n_goalset=goal_pose.n_goalset, - ) - else: - raise ValueError("Unsupported Solve type") - return solve_state - - def _plan_attempts( - self, - solve_state: ReacherSolveState, - start_state: JointState, - goal_pose: Pose, - plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), - link_poses: List[Pose] = None, - ): - start_time = time.time() - if plan_config.pose_cost_metric is not None: - valid_query = self.update_pose_cost_metric( - plan_config.pose_cost_metric, start_state, goal_pose - ) - if not valid_query: - result = MotionGenResult( - success=torch.as_tensor([False], device=self.tensor_args.device), - valid_query=valid_query, - status="Invalid Hold Partial Pose", - ) - return result - # if plan_config.enable_opt: - self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size) - 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") - force_graph = plan_config.enable_graph - partial_ik = plan_config.partial_ik_opt - ik_fail_count = 0 - time_dict = { - "solve_time": 0, - "ik_time": 0, - "graph_time": 0, - "trajopt_time": 0, - "trajopt_attempts": 0, - } - best_status = 0 # 0== None, 1==IK Fail, 2== Graph Fail, 3==Opt Fail - if plan_config.finetune_dt_scale is None: - plan_config.finetune_dt_scale = self.finetune_dt_scale - for n in range(plan_config.max_attempts): - log_info("MG Iter: " + str(n)) - result = self._plan_from_solve_state( - solve_state, - start_state, - goal_pose, - plan_config, - link_poses, - ) - time_dict["solve_time"] += result.solve_time - time_dict["ik_time"] += result.ik_time - time_dict["graph_time"] += result.graph_time - time_dict["trajopt_time"] += result.trajopt_time - time_dict["trajopt_attempts"] += result.trajopt_attempts - if ( - result.status == "IK Fail" and plan_config.ik_fail_return is not None - ): # IF IK fails the first time, we exist assuming the goal is not reachable - ik_fail_count += 1 - best_status = max(best_status, 1) - - if ik_fail_count > plan_config.ik_fail_return: - break - if result.success[0].item(): - break - - if result.status == "Finetune Opt Fail": - plan_config.finetune_dt_scale *= ( - plan_config.finetune_dt_decay**plan_config.finetune_attempts - ) - plan_config.finetune_dt_scale = min(plan_config.finetune_dt_scale, 1.25) - if plan_config.enable_graph_attempt is not None and ( - n >= plan_config.enable_graph_attempt - 1 - and result.status in ["Opt Fail", "Finetune Opt Fail"] - and not plan_config.enable_graph - ): - plan_config.enable_graph = True - plan_config.partial_ik_opt = False - if plan_config.disable_graph_attempt is not None and ( - n >= plan_config.disable_graph_attempt - 1 - and result.status in ["Opt Fail", "Graph Fail", "Finetune Opt Fail"] - and not force_graph - ): - plan_config.enable_graph = False - plan_config.partial_ik_opt = partial_ik - if result.status in ["Opt Fail"]: - best_status = 3 - elif result.status in ["Graph Fail"]: - best_status = 2 - if time.time() - start_time > plan_config.timeout: - break - if not result.valid_query: - result.status = "Invalid Problem" - break - if n == 10: - self.reset_seed() - log_warn("Couldn't find solution with 10 attempts, resetting seeds") - - result.solve_time = time_dict["solve_time"] - result.ik_time = time_dict["ik_time"] - result.graph_time = time_dict["graph_time"] - result.trajopt_time = time_dict["trajopt_time"] - result.trajopt_attempts = time_dict["trajopt_attempts"] - result.attempts = n + 1 - torch.cuda.synchronize() - if plan_config.pose_cost_metric is not None: - self.update_pose_cost_metric(PoseCostMetric.reset_metric()) - result.total_time = time.time() - start_time - return result - - def _plan_batch_attempts( - self, - solve_state: ReacherSolveState, - start_state: JointState, - goal_pose: Pose, - plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), - ): - start_time = time.time() - goal_pose = goal_pose.clone() - if plan_config.pose_cost_metric is not None: - valid_query = self.update_pose_cost_metric( - plan_config.pose_cost_metric, start_state, goal_pose - ) - if not valid_query: - result = MotionGenResult( - success=torch.as_tensor( - [False for _ in solve_state.batch_size], - device=self.motion_gen.tensor_args.device, - ), - valid_query=valid_query, - status="Invalid Hold Partial Pose", - ) - return result - - 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") - if plan_config.enable_graph: - raise ValueError("Graph Search / Geometric Planner not supported in batch_env mode") - - if plan_config.enable_graph or ( - plan_config.enable_graph_attempt is not None - and plan_config.max_attempts >= plan_config.enable_graph_attempt - ): - log_warn("Batch mode enable graph is only supported with num_graph_seeds==1") - plan_config.num_trajopt_seeds = 1 - plan_config.num_graph_seeds = 1 - solve_state.num_trajopt_seeds = 1 - solve_state.num_graph_seeds = 1 - self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size) - - ik_fail_count = 0 - force_graph = plan_config.enable_graph - partial_ik = plan_config.partial_ik_opt - - time_dict = { - "solve_time": 0, - "ik_time": 0, - "graph_time": 0, - "trajopt_time": 0, - "trajopt_attempts": 0, - } - best_result = None - - for n in range(plan_config.max_attempts): - result = self._plan_from_solve_state_batch( - solve_state, start_state, goal_pose, plan_config - ) - - time_dict["solve_time"] += result.solve_time - time_dict["ik_time"] += result.ik_time - - time_dict["graph_time"] += result.graph_time - time_dict["trajopt_time"] += result.trajopt_time - time_dict["trajopt_attempts"] += result.trajopt_attempts - - # if not all have succeeded, store the successful ones and re attempt: - # TODO: update stored based on error - if best_result is None: - best_result = result.clone() - else: - # get success idx: - idx = torch.nonzero(result.success).reshape(-1) - if len(idx) > 0: - best_result.copy_idx(idx, result) - - if ( - result.status == "IK Fail" and plan_config.ik_fail_return is not None - ): # IF IK fails the first time, we exit assuming the goal is not reachable - ik_fail_count += 1 - if ik_fail_count > plan_config.ik_fail_return: - break - - if ( - torch.count_nonzero(best_result.success) - >= goal_pose.batch * plan_config.success_ratio - ): # we want 90% targets to succeed - best_result.status = None - break - if plan_config.enable_graph_attempt is not None and ( - n >= plan_config.enable_graph_attempt - 1 - and result.status != "IK Fail" - and not plan_config.enable_graph - ): - plan_config.enable_graph = True - plan_config.partial_ik_opt = False - - if plan_config.disable_graph_attempt is not None and ( - n >= plan_config.disable_graph_attempt - 1 - and result.status in ["Opt Fail", "Graph Fail"] - and not force_graph - ): - plan_config.enable_graph = False - plan_config.partial_ik_opt = partial_ik - - if plan_config.fail_on_invalid_query: - if not result.valid_query: - best_result.valid_query = False - best_result.status = "Invalid Problem" - break - if time.time() - start_time > plan_config.timeout: - break - best_result.solve_time = time_dict["solve_time"] - best_result.ik_time = time_dict["ik_time"] - best_result.graph_time = time_dict["graph_time"] - best_result.trajopt_time = time_dict["trajopt_time"] - best_result.trajopt_attempts = time_dict["trajopt_attempts"] - best_result.attempts = n + 1 - torch.cuda.synchronize() - if plan_config.pose_cost_metric is not None: - self.update_pose_cost_metric(PoseCostMetric.reset_metric()) - best_result.total_time = time.time() - start_time - return best_result - def plan_single( self, start_state: JointState, @@ -1388,6 +1433,27 @@ class MotionGen(MotionGenConfig): plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), link_poses: List[Pose] = None, ) -> MotionGenResult: + """Plan a single motion to reach a goal pose from a start joint state. + + This also supports reaching target poses for different links in the robot by providing + goal poses for each link in the link_poses argument. The link_poses argument should contain + a pose for each link specified in :attr:`MotionGen.kinematics`. Constrained planning is + supported by calling :meth:`MotionGen.update_pose_cost_metric` before calling this method. + See :ref:`tut_constrained_planning` for more details. + + Args: + start_state: Start joint state of the robot. When planning from a non-static state, i.e, + when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` to + False. + goal_pose: Goal pose for the end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for each link in the robot when planning for multiple links. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to see if the query was successful. + """ + log_info("Planning for Single Goal: " + str(goal_pose.batch)) solve_state = self._get_solve_state( ReacherSolveType.SINGLE, plan_config, goal_pose, start_state ) @@ -1408,6 +1474,25 @@ class MotionGen(MotionGenConfig): plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), link_poses: List[Pose] = None, ) -> MotionGenResult: + """Plan a single motion to reach a goal from set of poses, from a start joint state. + + Use this when planning to reach a grasp from a set of possible grasps. This method will + try to reach the closest goal pose from the set of goal poses at every iteration of + optimization during IK and trajectory optimization. This method only supports goalset for + main end-effector (i.e., `goal_pose`). Use single Pose target for links in `link_poses`. + + Args: + start_state: Start joint state of the robot. When planning from a non-static state, + i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` + to False. + goal_pose: Goal pose for the end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for each link in the robot when planning for multiple links. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to see if the query was successful. + """ solve_state = self._get_solve_state( ReacherSolveType.GOALSET, plan_config, goal_pose, start_state ) @@ -1426,7 +1511,22 @@ class MotionGen(MotionGenConfig): start_state: JointState, goal_pose: Pose, plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Dict[str, List[Pose]] = None, ) -> MotionGenResult: + """Plan motions to reach a batch of goal poses from a batch of start joint states. + + Args: + start_state: Start joint states of the robot. When planning from a non-static state, + i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` + to False. + goal_pose: Goal poses for the end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for each link in the robot when planning for multiple links. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to check which indices of the batch were successful. + """ solve_state = self._get_solve_state( ReacherSolveType.BATCH, plan_config, goal_pose, start_state ) @@ -1436,6 +1536,7 @@ class MotionGen(MotionGenConfig): start_state, goal_pose, plan_config, + link_poses=link_poses, ) return result @@ -1444,7 +1545,23 @@ class MotionGen(MotionGenConfig): start_state: JointState, goal_pose: Pose, plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Dict[str, List[Pose]] = None, ) -> MotionGenResult: + """Plan motions to reach a batch of poses (goalset) from a batch of start joint states. + + Args: + start_state: Start joint states of the robot. When planning from a non-static state, + i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` + to False. + goal_pose: Goal poses for the end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for each link in the robot when planning for multiple links. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to check which indices of the batch were successful. + """ + solve_state = self._get_solve_state( ReacherSolveType.BATCH_GOALSET, plan_config, goal_pose, start_state ) @@ -1454,24 +1571,7 @@ class MotionGen(MotionGenConfig): start_state, goal_pose, plan_config, - ) - return result - - def plan_batch_env_goalset( - self, - start_state: JointState, - goal_pose: Pose, - plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), - ) -> MotionGenResult: - solve_state = self._get_solve_state( - ReacherSolveType.BATCH_ENV_GOALSET, plan_config, goal_pose, start_state - ) - - result = self._plan_batch_attempts( - solve_state, - start_state, - goal_pose, - plan_config, + link_poses=link_poses, ) return result @@ -1480,10 +1580,85 @@ class MotionGen(MotionGenConfig): start_state: JointState, goal_pose: Pose, plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Dict[str, List[Pose]] = None, ) -> MotionGenResult: + """Plan motions to reach (batch) poses in different collision environments. + + Args: + start_state: Start joint states of the robot. When planning from a non-static state, + i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` + to False. + goal_pose: Goal poses for the end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for each link in the robot when planning for multiple links. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to check which indices of the batch were successful. + """ + if plan_config.enable_graph: + log_info( + "Batch env mode does not support graph search, setting " + + "MotionGenPlanConfig.enable_graph=False" + ) + plan_config.enable_graph = False + + if plan_config.enable_graph_attempt is not None: + log_info( + "Batch env mode does not support graph search, setting " + + "MotionGenPlanConfig.enable_graph_attempt=None" + ) + plan_config.enable_graph_attempt = None solve_state = self._get_solve_state( ReacherSolveType.BATCH_ENV, plan_config, goal_pose, start_state ) + result = self._plan_batch_attempts( + solve_state, + start_state, + goal_pose, + plan_config, + link_poses=link_poses, + ) + return result + + def plan_batch_env_goalset( + self, + start_state: JointState, + goal_pose: Pose, + plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Dict[str, List[Pose]] = None, + ) -> MotionGenResult: + """Plan motions to reach (batch) goalset poses in different collision environments. + + Args: + start_state: Start joint states of the robot. When planning from a non-static state, + i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` + to False. + goal_pose: Goal poses for the end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for each link in the robot when planning for multiple links. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to check which indices of the batch were successful. + """ + + if plan_config.enable_graph: + log_info( + "Batch env mode does not support graph search, setting " + + "MotionGenPlanConfig.enable_graph=False" + ) + plan_config.enable_graph = False + + if plan_config.enable_graph_attempt is not None: + log_info( + "Batch env mode does not support graph search, setting " + + "MotionGenPlanConfig.enable_graph_attempt=None" + ) + plan_config.enable_graph_attempt = None + solve_state = self._get_solve_state( + ReacherSolveType.BATCH_ENV_GOALSET, plan_config, goal_pose, start_state + ) result = self._plan_batch_attempts( solve_state, start_state, @@ -1492,786 +1667,85 @@ class MotionGen(MotionGenConfig): ) return result - def _plan_from_solve_state( - self, - solve_state: ReacherSolveState, - start_state: JointState, - goal_pose: Pose, - plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), - link_poses: Optional[Dict[str, Pose]] = None, - ) -> MotionGenResult: - trajopt_seed_traj = None - trajopt_seed_success = None - trajopt_newton_iters = None - graph_success = 0 - if len(start_state.shape) == 1: - log_error("Joint state should be not a vector (dof) should be (bxdof)") - # plan ik: - - ik_result = self._solve_ik_from_solve_state( - goal_pose, - solve_state, - start_state, - plan_config.use_nn_ik_seed, - plan_config.partial_ik_opt, - link_poses, - ) - - if not plan_config.enable_graph and plan_config.partial_ik_opt: - ik_result.success[:] = True - - # check for success: - result = MotionGenResult( - ik_result.success.view(-1)[0:1], # This is not true for batch mode - ik_time=ik_result.solve_time, - solve_time=ik_result.solve_time, - ) - - if self.store_debug_in_result: - result.debug_info = {"ik_result": ik_result} - ik_success = torch.count_nonzero(ik_result.success) - if ik_success == 0: - result.status = "IK Fail" - return result - - # do graph search: - with profiler.record_function("motion_gen/post_ik"): - ik_out_seeds = solve_state.num_trajopt_seeds - if plan_config.enable_graph: - ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success) - - goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof)[ - :ik_success - ] - start_config = tensor_repeat_seeds(start_state.position, ik_out_seeds) - if plan_config.enable_opt: - self._trajopt_goal_config[:] = ik_result.solution - - # do graph search: - if plan_config.enable_graph: - interpolation_steps = None - if plan_config.enable_opt: - interpolation_steps = self.trajopt_solver.action_horizon - log_info("MG: running GP") - graph_result = self.graph_search(start_config, goal_config, interpolation_steps) - trajopt_seed_success = graph_result.success - - graph_success = torch.count_nonzero(graph_result.success).item() - result.graph_time = graph_result.solve_time - result.solve_time += graph_result.solve_time - if graph_success > 0: - result.graph_plan = graph_result.interpolated_plan - result.interpolated_plan = graph_result.interpolated_plan - - result.used_graph = True - if plan_config.enable_opt: - trajopt_seed = ( - result.graph_plan.position.view( - 1, # solve_state.batch_size, - graph_success, # solve_state.num_trajopt_seeds, - interpolation_steps, - self._dof, - ) - .transpose(0, 1) - .contiguous() - ) - trajopt_seed_traj = torch.zeros( - (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof), - device=self.tensor_args.device, - dtype=self.tensor_args.dtype, - ) - trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed - trajopt_seed_success = ik_result.success.clone() - trajopt_seed_success[ik_result.success] = graph_result.success - - trajopt_seed_success = trajopt_seed_success.view( - solve_state.batch_size, solve_state.num_trajopt_seeds - ) - trajopt_newton_iters = self.graph_trajopt_iters - else: - _, idx = torch.topk( - graph_result.path_length[graph_result.success], k=1, largest=False - ) - result.interpolated_plan = result.interpolated_plan[idx].squeeze(0) - result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt) - result.optimized_plan = result.interpolated_plan[ - : graph_result.path_buffer_last_tstep[idx.item()] - ] - idx = idx.view(-1) + self._batch_col - result.position_error = ik_result.position_error[ik_result.success][ - graph_result.success - ][idx] - result.rotation_error = ik_result.rotation_error[ik_result.success][ - graph_result.success - ][idx] - result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[ - idx.item() : idx.item() + 1 - ] - result.success = result.success.view(-1)[0:1] - result.success[:] = True - return result - else: - result.success[:] = False - result.status = "Graph Fail" - if not graph_result.valid_query: - result.valid_query = False - if self.store_debug_in_result: - result.debug_info["graph_debug"] = graph_result.debug_info - return result - if plan_config.need_graph_success: - return result - - # do trajopt:: - - if plan_config.enable_opt: - with profiler.record_function("motion_gen/setup_trajopt_seeds"): - self._trajopt_goal_config[:, :ik_success] = goal_config - retract_config = None - if plan_config.use_start_state_as_retract: - retract_config = start_state.position.clone() - goal = Goal( - goal_pose=goal_pose, - current_state=start_state, - links_goal_pose=link_poses, - retract_state=retract_config, - ) - - if ( - trajopt_seed_traj is None - or graph_success < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds - ): - goal_config = self._trajopt_goal_config[0] # batch index == 0 - - goal_state = JointState.from_position( - goal_config, - ) - seed_link_poses = None - if link_poses is not None: - seed_link_poses = {} - - for k in link_poses.keys(): - seed_link_poses[k] = link_poses[k].repeat_seeds( - solve_state.num_trajopt_seeds - ) - seed_goal = Goal( - goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds), - current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds), - goal_state=goal_state, - links_goal_pose=seed_link_poses, - ) - if trajopt_seed_traj is not None: - trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() - # batch, num_seeds, h, dof - if ( - trajopt_seed_success.shape[1] - < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds - ): - trajopt_seed_success_new = torch.zeros( - (1, solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds), - device=self.tensor_args.device, - dtype=torch.bool, - ) - trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = ( - trajopt_seed_success - ) - trajopt_seed_success = trajopt_seed_success_new - # create seeds here: - trajopt_seed_traj = self.trajopt_solver.get_seed_set( - seed_goal, - trajopt_seed_traj, # batch, num_seeds, h, dof - num_seeds=self.noisy_trajopt_seeds, - batch_mode=solve_state.batch_mode, - seed_success=trajopt_seed_success, - ) - trajopt_seed_traj = trajopt_seed_traj.view( - solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds, - solve_state.batch_size, - self.trajopt_solver.action_horizon, - self._dof, - ).contiguous() - if plan_config.enable_finetune_trajopt: - og_value = self.trajopt_solver.interpolation_type - self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA - with profiler.record_function("motion_gen/trajopt"): - log_info("MG: running TO") - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - trajopt_seed_traj, - num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds, - newton_iters=trajopt_newton_iters, - return_all_solutions=plan_config.parallel_finetune - and plan_config.enable_finetune_trajopt, - ) - if plan_config.enable_finetune_trajopt: - self.trajopt_solver.interpolation_type = og_value - if self.store_debug_in_result: - result.debug_info["trajopt_result"] = traj_result - # run finetune - if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0: - with profiler.record_function("motion_gen/finetune_trajopt"): - seed_traj = traj_result.raw_action.clone() # solution.position.clone() - seed_traj = seed_traj.contiguous() - og_solve_time = traj_result.solve_time - seed_override = 1 - opt_dt = traj_result.optimized_dt - - if plan_config.parallel_finetune: - opt_dt = torch.min(opt_dt[traj_result.success]) - seed_override = solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds - - finetune_time = 0 - for k in range(plan_config.finetune_attempts): - newton_iters = None - - scaled_dt = torch.clamp( - opt_dt - * plan_config.finetune_dt_scale - * (plan_config.finetune_dt_decay ** (k)), - self.trajopt_solver.interpolation_dt, - ) - if self.optimize_dt: - self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item()) - - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - seed_traj, - trajopt_instance=self.finetune_trajopt_solver, - num_seeds_override=seed_override, - newton_iters=newton_iters, - ) - finetune_time += traj_result.solve_time - if torch.count_nonzero(traj_result.success) > 0: - break - seed_traj = traj_result.optimized_seeds.detach().clone() - newton_iters = 4 - - traj_result.solve_time = finetune_time - - result.finetune_time = traj_result.solve_time - - traj_result.solve_time = og_solve_time - if self.store_debug_in_result: - result.debug_info["finetune_trajopt_result"] = traj_result - elif plan_config.enable_finetune_trajopt: - traj_result.success = traj_result.success[0:1] - # if torch.count_nonzero(result.success) == 0: - result.status = "Opt Fail" - result.solve_time += traj_result.solve_time + result.finetune_time - result.trajopt_time = traj_result.solve_time - result.trajopt_attempts = 1 - result.success = traj_result.success - - if plan_config.enable_finetune_trajopt and torch.count_nonzero(result.success) == 0: - result.status = "Finetune Opt Fail" - - result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory( - 0, traj_result.path_buffer_last_tstep[0] - ) - result.interpolation_dt = self.trajopt_solver.interpolation_dt - result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep - result.position_error = traj_result.position_error - result.rotation_error = traj_result.rotation_error - result.optimized_dt = traj_result.optimized_dt - result.optimized_plan = traj_result.solution - result.goalset_index = traj_result.goalset_index - return result - - def _plan_js_from_solve_state( - self, - solve_state: ReacherSolveState, - start_state: JointState, - goal_state: JointState, - plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), - ) -> MotionGenResult: - trajopt_seed_traj = None - trajopt_seed_success = None - trajopt_newton_iters = self.js_trajopt_solver.newton_iters - - graph_success = 0 - if len(start_state.shape) == 1: - log_error("Joint state should be not a vector (dof) should be (bxdof)") - - result = MotionGenResult(cspace_error=torch.zeros((1), device=self.tensor_args.device)) - if self.store_debug_in_result: - result.debug_info = {} - # do graph search: - if plan_config.enable_graph and False: - start_config = torch.zeros( - (solve_state.num_graph_seeds, self.js_trajopt_solver.dof), - device=self.tensor_args.device, - dtype=self.tensor_args.dtype, - ) - goal_config = start_config.clone() - start_config[:] = start_state.position - goal_config[:] = goal_state.position - interpolation_steps = None - if plan_config.enable_opt: - interpolation_steps = self.js_trajopt_solver.action_horizon - log_info("MG: running GP") - graph_result = self.graph_search(start_config, goal_config, interpolation_steps) - trajopt_seed_success = graph_result.success - - graph_success = torch.count_nonzero(graph_result.success).item() - result.graph_time = graph_result.solve_time - result.solve_time += graph_result.solve_time - if graph_success > 0: - result.graph_plan = graph_result.interpolated_plan - result.interpolated_plan = graph_result.interpolated_plan - - result.used_graph = True - if plan_config.enable_opt: - trajopt_seed = ( - result.graph_plan.position.view( - 1, # solve_state.batch_size, - graph_success, # solve_state.num_trajopt_seeds, - interpolation_steps, - self._dof, - ) - .transpose(0, 1) - .contiguous() - ) - trajopt_seed_traj = torch.zeros( - (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof), - device=self.tensor_args.device, - dtype=self.tensor_args.dtype, - ) - trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed - trajopt_seed_success = graph_result.success - - trajopt_seed_success = trajopt_seed_success.view( - 1, solve_state.num_trajopt_seeds - ) - trajopt_newton_iters = self.graph_trajopt_iters - else: - _, idx = torch.topk( - graph_result.path_length[graph_result.success], k=1, largest=False - ) - result.interpolated_plan = result.interpolated_plan[idx].squeeze(0) - result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt) - result.optimized_plan = result.interpolated_plan[ - : graph_result.path_buffer_last_tstep[idx.item()] - ] - idx = idx.view(-1) + self._batch_col - result.cspace_error = torch.zeros((1), device=self.tensor_args.device) - - result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[ - idx.item() : idx.item() + 1 - ] - result.success = torch.as_tensor([True], device=self.tensor_args.device) - return result - else: - result.success = torch.as_tensor([False], device=self.tensor_args.device) - result.status = "Graph Fail" - if not graph_result.valid_query: - result.valid_query = False - if self.store_debug_in_result: - result.debug_info["graph_debug"] = graph_result.debug_info - return result - if plan_config.need_graph_success: - return result - - # do trajopt: - if plan_config.enable_opt: - with profiler.record_function("motion_gen/setup_trajopt_seeds"): - - goal = Goal( - current_state=start_state, - goal_state=goal_state, - ) - - if trajopt_seed_traj is None or graph_success < solve_state.num_trajopt_seeds * 1: - seed_goal = Goal( - current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds), - goal_state=goal_state.repeat_seeds(solve_state.num_trajopt_seeds), - ) - if trajopt_seed_traj is not None: - trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() - # batch, num_seeds, h, dof - if trajopt_seed_success.shape[1] < self.js_trajopt_solver.num_seeds: - trajopt_seed_success_new = torch.zeros( - (1, solve_state.num_trajopt_seeds), - device=self.tensor_args.device, - dtype=torch.bool, - ) - trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = ( - trajopt_seed_success - ) - trajopt_seed_success = trajopt_seed_success_new - # create seeds here: - trajopt_seed_traj = self.js_trajopt_solver.get_seed_set( - seed_goal, - trajopt_seed_traj, # batch, num_seeds, h, dof - num_seeds=1, - batch_mode=False, - seed_success=trajopt_seed_success, - ) - trajopt_seed_traj = ( - trajopt_seed_traj.view( - self.js_trajopt_solver.num_seeds * 1, - 1, - self.trajopt_solver.action_horizon, - self._dof, - ) - .contiguous() - .clone() - ) - if plan_config.enable_finetune_trajopt: - og_value = self.trajopt_solver.interpolation_type - self.js_trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA - with profiler.record_function("motion_gen/trajopt"): - log_info("MG: running TO") - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - trajopt_seed_traj, - num_seeds_override=solve_state.num_trajopt_seeds, - newton_iters=trajopt_newton_iters + 2, - return_all_solutions=plan_config.enable_finetune_trajopt, - trajopt_instance=self.js_trajopt_solver, - ) - if plan_config.enable_finetune_trajopt: - self.trajopt_solver.interpolation_type = og_value - # self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate) - if self.store_debug_in_result: - result.debug_info["trajopt_result"] = traj_result - if torch.count_nonzero(traj_result.success) == 0: - result.status = "TrajOpt Fail" - # run finetune - if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0: - with profiler.record_function("motion_gen/finetune_trajopt"): - seed_traj = traj_result.raw_action.clone() # solution.position.clone() - seed_traj = seed_traj.contiguous() - og_solve_time = traj_result.solve_time - - scaled_dt = torch.clamp( - torch.max(traj_result.optimized_dt[traj_result.success]), - self.trajopt_solver.interpolation_dt, - ) - og_dt = self.js_trajopt_solver.solver_dt.clone() - self.js_trajopt_solver.update_solver_dt(scaled_dt.item()) - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - seed_traj, - trajopt_instance=self.js_trajopt_solver, - num_seeds_override=solve_state.num_trajopt_seeds, - newton_iters=trajopt_newton_iters + 4, - ) - self.js_trajopt_solver.update_solver_dt(og_dt) - - result.finetune_time = traj_result.solve_time - - traj_result.solve_time = og_solve_time - if self.store_debug_in_result: - result.debug_info["finetune_trajopt_result"] = traj_result - if torch.count_nonzero(traj_result.success) == 0: - result.status = "Finetune Fail" - elif plan_config.enable_finetune_trajopt: - traj_result.success = traj_result.success[0:1] - result.solve_time += traj_result.solve_time + result.finetune_time - result.trajopt_time = traj_result.solve_time - result.trajopt_attempts = 1 - result.success = traj_result.success - - result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory( - 0, traj_result.path_buffer_last_tstep[0] - ) - - result.interpolation_dt = self.trajopt_solver.interpolation_dt - result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep - result.cspace_error = traj_result.cspace_error - result.optimized_dt = traj_result.optimized_dt - result.optimized_plan = traj_result.solution - result.goalset_index = traj_result.goalset_index - - return result - - def _plan_from_solve_state_batch( - self, - solve_state: ReacherSolveState, - start_state: JointState, - goal_pose: Pose, - plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), - ) -> MotionGenResult: - self._trajopt_goal_config[:] = self.get_retract_config().view(1, 1, self._dof) - trajopt_seed_traj = None - trajopt_seed_success = None - trajopt_newton_iters = None - graph_success = 0 - - # plan ik: - ik_result = self._solve_ik_from_solve_state( - goal_pose, - solve_state, - start_state, - plan_config.use_nn_ik_seed, - plan_config.partial_ik_opt, - ) - - if not plan_config.enable_graph and plan_config.partial_ik_opt: - ik_result.success[:] = True - - # check for success: - result = MotionGenResult( - ik_result.success, - position_error=ik_result.position_error, - rotation_error=ik_result.rotation_error, - ik_time=ik_result.solve_time, - solve_time=ik_result.solve_time, - debug_info={}, - # goalset_index=ik_result.goalset_index, - ) - - ik_success = torch.count_nonzero(ik_result.success) - if ik_success == 0: - result.status = "IK Fail" - result.success = result.success[:, 0] - return result - - # do graph search: - ik_out_seeds = solve_state.num_trajopt_seeds - if plan_config.enable_graph: - ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success) - - # if not plan_config.enable_opt and plan_config.enable_graph: - # self.graph_planner.interpolation_steps = self.interpolation_steps - # self.graph_planner.interpolation_type = self.interpolation_type - # elif plan_config.enable_graph: - # self.graph_planner.interpolation_steps = self.trajopt_solver.traj_tsteps - # self.graph_planner.interpolation_type = InterpolateType.LINEAR - goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof) - - # get shortest path - if plan_config.enable_graph: - interpolation_steps = None - if plan_config.enable_opt: - interpolation_steps = self.trajopt_solver.action_horizon - - start_graph_state = start_state.repeat_seeds(ik_out_seeds) - start_config = start_graph_state.position[ik_result.success.view(-1)].view( - -1, self.ik_solver.dof - ) - graph_result = self.graph_search(start_config, goal_config, interpolation_steps) - graph_success = torch.count_nonzero(graph_result.success).item() - - result.graph_time = graph_result.solve_time - result.solve_time += graph_result.solve_time - if graph_success > 0: - # path = graph_result.interpolated_plan - result.graph_plan = graph_result.interpolated_plan - result.interpolated_plan = graph_result.interpolated_plan - result.used_graph = True - - if plan_config.enable_opt: - trajopt_seed = result.graph_plan.position.view( - graph_success, # solve_state.num_trajopt_seeds, - interpolation_steps, - self._dof, - ).contiguous() - trajopt_seed_traj = torch.zeros( - (1, trajopt_seed.shape[0], self.trajopt_solver.action_horizon, self._dof), - device=self.tensor_args.device, - dtype=self.tensor_args.dtype, - ) - trajopt_seed_traj[0, :, :interpolation_steps, :] = trajopt_seed - trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() - trajopt_seed_success = ik_result.success.clone() - trajopt_seed_success[ik_result.success] = graph_result.success - - trajopt_seed_success = trajopt_seed_success.view( - solve_state.num_trajopt_seeds, solve_state.batch_size - ) - trajopt_newton_iters = self.graph_trajopt_iters - - else: - ik_success = ik_result.success.view(-1).clone() - - # only some might be successful: - - g_dim = torch.nonzero(ik_success).view(-1)[graph_result.success] - - self._batch_graph_search_buffer.copy_at_index( - graph_result.interpolated_plan, g_dim - ) - - # result.graph_plan = JointState.from_position( - # self._batch_graph_search_buffer, - # joint_names=graph_result.interpolated_plan.joint_names, - # ) - result.interpolated_plan = self._batch_graph_search_buffer - g_dim = g_dim.cpu().squeeze().tolist() - for x, x_val in enumerate(g_dim): - self._batch_path_buffer_last_tstep[x_val] = ( - graph_result.path_buffer_last_tstep[x] - ) - result.path_buffer_last_tstep = self._batch_path_buffer_last_tstep - result.optimized_plan = result.interpolated_plan - result.optimized_dt = torch.as_tensor( - [ - self.interpolation_dt - for i in range(result.interpolated_plan.position.shape[0]) - ], - device=self.tensor_args.device, - dtype=self.tensor_args.dtype, - ) - result.success = result.success.view(-1).clone() - result.success[ik_success][graph_result.success] = True - return result - - else: - result.success[:] = False - result.success = result.success[:, 0] - result.status = "Graph Fail" - if not graph_result.valid_query: - result.valid_query = False - if self.store_debug_in_result: - result.debug_info = {"graph_debug": graph_result.debug_info} - return result - - if plan_config.enable_opt: - # get goal configs based on ik success: - self._trajopt_goal_config[ik_result.success] = goal_config - - goal_config = self._trajopt_goal_config # batch index == 0 - - goal = Goal( - goal_pose=goal_pose, - current_state=start_state, - ) - # generate seeds: - if trajopt_seed_traj is None or ( - plan_config.enable_graph and graph_success < solve_state.batch_size - ): - seed_goal = Goal( - goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds), - current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds), - goal_state=JointState.from_position(goal_config.view(-1, self._dof)), - ) - if trajopt_seed_traj is not None: - trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() - - # create seeds here: - trajopt_seed_traj = self.trajopt_solver.get_seed_set( - seed_goal, - trajopt_seed_traj, # batch, num_seeds, h, dof - num_seeds=1, - batch_mode=solve_state.batch_mode, - seed_success=trajopt_seed_success, - ) - trajopt_seed_traj = trajopt_seed_traj.view( - solve_state.num_trajopt_seeds, - solve_state.batch_size, - self.trajopt_solver.action_horizon, - self._dof, - ).contiguous() - if plan_config.enable_finetune_trajopt: - og_value = self.trajopt_solver.interpolation_type - self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA - - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - trajopt_seed_traj, - newton_iters=trajopt_newton_iters, - return_all_solutions=True, - ) - - # output of traj result will have 1 solution per batch - - # run finetune opt on 1 solution per batch: - if plan_config.enable_finetune_trajopt: - self.trajopt_solver.interpolation_type = og_value - # self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate) - if self.store_debug_in_result: - result.debug_info["trajopt_result"] = traj_result - - # run finetune - if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0: - with profiler.record_function("motion_gen/finetune_trajopt"): - seed_traj = traj_result.raw_action.clone() # solution.position.clone() - seed_traj = seed_traj.contiguous() - og_solve_time = traj_result.solve_time - - scaled_dt = torch.clamp( - torch.max(traj_result.optimized_dt[traj_result.success]) - * self.finetune_dt_scale, - self.trajopt_solver.interpolation_dt, - ) - self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item()) - - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - seed_traj, - trajopt_instance=self.finetune_trajopt_solver, - num_seeds_override=solve_state.num_trajopt_seeds, - ) - - result.finetune_time = traj_result.solve_time - - traj_result.solve_time = og_solve_time - if self.store_debug_in_result: - result.debug_info["finetune_trajopt_result"] = traj_result - elif plan_config.enable_finetune_trajopt: - traj_result.success = traj_result.success[:, 0] - - result.success = traj_result.success - - result.interpolated_plan = traj_result.interpolated_solution - result.solve_time += traj_result.solve_time - result.trajopt_time = traj_result.solve_time - result.position_error = traj_result.position_error - result.rotation_error = traj_result.rotation_error - result.cspace_error = traj_result.cspace_error - result.goalset_index = traj_result.goalset_index - result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep - result.optimized_plan = traj_result.solution - result.optimized_dt = traj_result.optimized_dt - if torch.count_nonzero(traj_result.success) == 0: - result.status = "Opt Fail" - result.success[:] = False - if self.store_debug_in_result: - result.debug_info = {"trajopt_result": traj_result} - return result - def compute_kinematics(self, state: JointState) -> KinematicModelState: + """Compute kinematics for a given joint state. + + Args: + state: Joint state of the robot. Only :attr:`JointState.position` is used. + + Returns: + KinematicModelState: Kinematic state of the robot. + """ out = self.rollout_fn.compute_kinematics(state) return out @property - def kinematics(self): + def kinematics(self) -> CudaRobotModel: + """Returns the shared kinematics model of the robot.""" return self.rollout_fn.kinematics @property - def dof(self): + def dof(self) -> int: + """Returns the number of controlled degrees of freedom of the robot.""" return self.kinematics.get_dof() + @property + def collision_cache(self) -> Dict[str, int]: + """Returns the collision cache created by the world collision checker.""" + return self.world_coll_checker.cache + def check_constraints(self, state: JointState) -> RolloutMetrics: + """Compute IK constraints for a given joint state. + + Args: + state: Joint state of the robot. + + Returns: + RolloutMetrics: Metrics for the joint state. + """ metrics = self.ik_solver.check_constraints(state) return metrics def update_world(self, world: WorldConfig): + """Update the world representation for collision checking. + + 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:`MotionGen.collision_cache` as + created during initialization of :class:`MotionGenConfig`. Updating the world also + invalidates the cached roadmaps in the graph planner. See :ref:`world_collision` for more + details. + + Args: + world: New world configuration for collision checking. + """ self.world_coll_checker.load_collision_model(world, fix_cache_reference=self.use_cuda_graph) self.graph_planner.reset_graph() - return True def clear_world_cache(self): + """Remove all collision objects from collision cache.""" + self.world_coll_checker.clear_cache() def reset(self, reset_seed=True): + """Reset the motion generation module. + + Args: + reset_seed: Reset the random seed generator. Resetting this can be time consuming, if + deterministic results are not required, set this to False. + """ self.graph_planner.reset_buffer() if reset_seed: self.reset_seed() def reset_seed(self): + """Reset the random seed generators in all sub-modules of motion generation.""" self.rollout_fn.reset_seed() self.ik_solver.reset_seed() self.graph_planner.reset_seed() self.trajopt_solver.reset_seed() - def get_retract_config(self): + def get_retract_config(self) -> T_DOF: + """Returns the retract/home configuration of the robot.""" return self.rollout_fn.dynamics_model.retract_config def warmup( @@ -2285,6 +1759,23 @@ class MotionGen(MotionGenConfig): warmup_joint_index: int = 0, warmup_joint_delta: float = 0.1, ): + """Warmup planning methods for motion generation. + + Args: + enable_graph: Enable graph search for warmup. + batch: Number of problem queries for batch mode. If None, single query is run. + warmup_js_trajopt: Warmup joint space planning in addition to Cartesian planning. + batch_env_mode: Enable batch world environment mode for warmup. Only used when batch is + not None. + parallel_finetune: Run finetuning trajectory optimization in parallel for warmup. Leave + this to True for most cases. + n_goalset: Number of goal poses to use for warmup. If -1, single goal pose is used. Set + this to the largest number of goals you plan to use with + :meth:`MotionGen.plan_goalset`. After warmup, you can use smaller number of goals + and the method will internally pad the extra goals with the first goal. + warmup_joint_index: Index of the joint to perturb for warmup. + warmup_joint_delta: Delta to perturb the joint for warmup. + """ log_info("Warmup") if warmup_js_trajopt: start_state = JointState.from_position( @@ -2391,7 +1882,8 @@ class MotionGen(MotionGenConfig): MotionGenPlanConfig( max_attempts=10, enable_finetune_trajopt=True, - enable_graph=enable_graph, + enable_graph=False, + enable_graph_attempt=None, ), ) else: @@ -2402,6 +1894,7 @@ class MotionGen(MotionGenConfig): max_attempts=10, enable_finetune_trajopt=True, enable_graph=enable_graph, + enable_graph_attempt=None if not enable_graph else 20, ), ) else: @@ -2420,8 +1913,7 @@ class MotionGen(MotionGenConfig): MotionGenPlanConfig( max_attempts=10, enable_finetune_trajopt=True, - enable_graph=enable_graph, - enable_graph_attempt=20, + enable_graph=False, ), ) else: @@ -2432,12 +1924,11 @@ class MotionGen(MotionGenConfig): max_attempts=10, enable_finetune_trajopt=True, enable_graph=enable_graph, - enable_graph_attempt=20, + enable_graph_attempt=None if not enable_graph else 20, ), ) log_info("Warmup complete") - return True def plan_single_js( self, @@ -2445,7 +1936,23 @@ class MotionGen(MotionGenConfig): goal_state: JointState, plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), ) -> MotionGenResult: - # NOTE: currently runs only one seed + """Plan a single motion to reach a goal joint state from a start joint state. + + This method uses trajectory optimization to find a collision-free path between the start + and goal joint states. If trajectory optimization fails, it uses a graph planner to find + a collision-free path to the goal joint state. The graph plan is then used as a seed for + trajectory optimization. + + Args: + start_state: Start joint state of the robot. + goal_state: Goal joint state of the robot. + plan_config: Planning parameters for motion generation. + + Returns: + MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success` + attribute to see if the query was successful. + """ + time_dict = { "solve_time": 0, "ik_time": 0, @@ -2495,84 +2002,11 @@ class MotionGen(MotionGenConfig): result.solve_time = result.trajopt_time + result.graph_time + result.finetune_time result.total_time = result.solve_time result.attempts = n - - return result - - def plan( - self, - start_state: JointState, - goal_pose: Pose, - enable_graph: bool = True, - enable_opt: bool = True, - use_nn_ik_seed: bool = False, - need_graph_success: bool = False, - max_attempts: int = 100, - timeout: float = 10.0, - enable_graph_attempt: Optional[int] = None, - disable_graph_attempt: Optional[int] = None, - trajopt_attempts: int = 1, - ik_fail_return: Optional[int] = None, - partial_ik_opt: bool = False, - num_ik_seeds: Optional[int] = None, - num_graph_seeds: Optional[int] = None, - num_trajopt_seeds: Optional[int] = None, - ): - plan_config = MotionGenPlanConfig( - enable_graph, - enable_opt, - use_nn_ik_seed, - need_graph_success, - max_attempts, - timeout, - enable_graph_attempt, - disable_graph_attempt, - ik_fail_return, - partial_ik_opt, - num_ik_seeds, - num_graph_seeds, - num_trajopt_seeds, - ) - result = self.plan_single(start_state, goal_pose, plan_config) - return result - - def batch_plan( - self, - start_state: JointState, - goal_pose: Pose, - enable_graph: bool = True, - enable_opt: bool = True, - use_nn_ik_seed: bool = False, - need_graph_success: bool = False, - max_attempts: int = 1, - timeout: float = 10.0, - enable_graph_attempt: Optional[int] = None, - disable_graph_attempt: Optional[int] = None, - success_ratio: float = 1.0, - ik_fail_return: Optional[int] = None, - fail_on_invalid_query: bool = False, - partial_ik_opt: bool = False, - num_ik_seeds: Optional[int] = None, - num_graph_seeds: Optional[int] = None, - num_trajopt_seeds: Optional[int] = None, - ): - plan_config = MotionGenPlanConfig( - enable_graph, - enable_opt, - use_nn_ik_seed, - need_graph_success, - max_attempts, - timeout, - enable_graph_attempt, - disable_graph_attempt, - ik_fail_return, - partial_ik_opt, - num_ik_seeds, - num_graph_seeds, - num_trajopt_seeds, - success_ratio=success_ratio, - fail_on_invalid_query=fail_on_invalid_query, - ) - result = self.plan_batch(start_state, goal_pose, plan_config) + if plan_config.time_dilation_factor is not None and torch.count_nonzero(result.success) > 0: + result.retime_trajectory( + plan_config.time_dilation_factor, + interpolation_kind=self.js_trajopt_solver.interpolation_type, + ) return result def solve_trajopt( @@ -2582,6 +2016,17 @@ class MotionGen(MotionGenConfig): return_all_solutions: bool = False, num_seeds: Optional[int] = None, ): + """Solve trajectory optimization for a given goal. + + Args: + goal: Goal for trajectory optimization. + act_seed: Seed for trajectory optimization. + return_all_solutions: Return results for all seeds in trajectory optimization. + num_seeds: Override number of seeds to use for trajectory optimization. + + Returns: + TrajOptResult: Result of trajectory optimization. + """ result = self.trajopt_solver.solve( goal, act_seed, return_all_solutions=return_all_solutions, num_seeds=num_seeds ) @@ -2591,6 +2036,17 @@ class MotionGen(MotionGenConfig): 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 @@ -2601,6 +2057,19 @@ class MotionGen(MotionGenConfig): start_state: Optional[JointState] = None, goal_pose: Optional[Pose] = None, ) -> bool: + """Update the pose cost metric for :ref:`tut_constrained_planning`. + + Only supports for the main end-effector. Does not support for multiple links that are + specified with `link_poses` in planning methods. + + Args: + metric: Type and parameters for pose constraint to add. + start_state: Start joint state for the constraint. + goal_pose: Goal pose for the constraint. + + Returns: + bool: True if the constraint can be added, False otherwise. + """ # check if constraint is valid: if metric.hold_partial_pose and metric.offset_tstep_fraction < 0.0: start_pose = self.compute_kinematics(start_state).ee_pose.clone() @@ -2648,10 +2117,11 @@ class MotionGen(MotionGenConfig): for rollout in rollouts if isinstance(rollout, ArmReacher) ] - torch.cuda.synchronize() + torch.cuda.synchronize(device=self.tensor_args.device) return True def get_all_rollout_instances(self) -> List[RolloutBase]: + """Get all rollout instances used across components in motion generation.""" if self._rollout_list is None: self._rollout_list = ( self.ik_solver.get_all_rollout_instances() @@ -2663,6 +2133,7 @@ class MotionGen(MotionGenConfig): return self._rollout_list def get_all_solver_rollout_instances(self) -> List[RolloutBase]: + """Get all rollout instances in solvers (IK, TrajOpt).""" if self._solver_rollout_list is None: self._solver_rollout_list = ( self.ik_solver.solver.get_all_rollout_instances() @@ -2673,6 +2144,7 @@ class MotionGen(MotionGenConfig): return self._solver_rollout_list def get_all_pose_solver_rollout_instances(self) -> List[RolloutBase]: + """Get all rollout instances in solvers (IK, TrajOpt) that support Cartesian cost terms.""" if self._pose_solver_rollout_list is None: self._pose_solver_rollout_list = ( self.ik_solver.solver.get_all_rollout_instances() @@ -2682,6 +2154,14 @@ class MotionGen(MotionGenConfig): return self._pose_solver_rollout_list def get_all_kinematics_instances(self) -> List[CudaRobotModel]: + """Get all kinematics instances used across components in motion generation. + + This is deprecated. Use :meth:`MotionGen.kinematics` instead as MotionGen now uses a shared + kinematics instance across all components. + + Returns: + List[CudaRobotModel]: Single kinematics instance, returned as a list for compatibility. + """ if self._kin_list is None: self._kin_list = [ i.dynamics_model.robot_model for i in self.get_all_rollout_instances() @@ -2699,13 +2179,36 @@ class MotionGen(MotionGenConfig): world_objects_pose_offset: Optional[Pose] = None, remove_obstacles_from_world_config: bool = False, ) -> None: - """Attach an object from world to robot's link. + """Attach an object or objects from world to a robot's link. + + This method assumes that the objects exist in the world configuration. If attaching + objects that are not in world, use :meth:`MotionGen.attach_external_objects_to_robot`. Args: - surface_sphere_radius: _description_. Defaults to None. - sphere_tensor: _description_. Defaults to None. - link_name: _description_. Defaults to "attached_object". + joint_state: Joint state of the robot. + object_names: Names of objects in the world to attach to the robot. + surface_sphere_radius: Radius (in meters) to use for points sampled on surface of the + object. A smaller radius will allow for generating motions very close to obstacles. + link_name: Name of the link (frame) to attach the objects to. The assumption is that + this link does not have any geometry and all spheres of this link represent + attached objects. + sphere_fit_type: Sphere fit algorithm to use. See :ref:`attach_object_note` for more + details. The default method :attr:`SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE` + voxelizes the volume of the objects and adds spheres representing the voxels, then + samples points on the surface of the object, adds :attr:`surface_sphere_radius` to + these points. This should be used for most cases. + voxelize_method: Method to use for voxelization, passed to + :py:func:`trimesh.voxel.creation.voxelize`. + world_objects_pose_offset: Offset to apply to the object poses before attaching to the + robot. This is useful when attaching an object that's in contact with the world. + The offset is applied in the world frame before attaching to the robot. + remove_obstacles_from_world_config: Remove the obstacles from the world cache after + attaching to the robot to reduce memory usage. Note that when an object is attached + to the robot, it's disabled in the world collision checker. This flag when enabled, + also removes the object from world cache. For most cases, this should be set to + False. """ + log_info("MG: Attach objects to robot") kin_state = self.compute_kinematics(joint_state) ee_pose = kin_state.ee_pose # w_T_ee @@ -2765,12 +2268,26 @@ class MotionGen(MotionGenConfig): voxelize_method: str = "ray", world_objects_pose_offset: Optional[Pose] = None, ) -> None: - """Attach an object from world to robot's link. + """Attach external objects (not in world model) to a robot's link. Args: - surface_sphere_radius: _description_. Defaults to None. - sphere_tensor: _description_. Defaults to None. - link_name: _description_. Defaults to "attached_object". + joint_state: Joint state of the robot. + external_objects: List of external objects to attach to the robot. + surface_sphere_radius: Radius (in meters) to use for points sampled on surface of the + object. A smaller radius will allow for generating motions very close to obstacles. + link_name: Name of the link (frame) to attach the objects to. The assumption is that + this link does not have any geometry and all spheres of this link represent + attached objects. + sphere_fit_type: Sphere fit algorithm to use. See :ref:`attach_object_note` for more + details. The default method :attr:`SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE` + voxelizes the volume of the objects and adds spheres representing the voxels, then + samples points on the surface of the object, adds :attr:`surface_sphere_radius` to + these points. This should be used for most cases. + voxelize_method: Method to use for voxelization, passed to + :py:func:`trimesh.voxel.creation.voxelize`. + world_objects_pose_offset: Offset to apply to the object poses before attaching to the + robot. This is useful when attaching an object that's in contact with the world. + The offset is applied in the world frame before attaching to the robot. """ log_info("MG: Attach objects to robot") if len(external_objects) == 0: @@ -2815,9 +2332,28 @@ class MotionGen(MotionGenConfig): self.attach_spheres_to_robot(sphere_tensor=sphere_tensor, link_name=link_name) def add_camera_frame(self, camera_observation: CameraObservation, obstacle_name: str): + """Add camera frame to the world collision checker. + + Only supported by :py:class:`~curobo.geom.sdf.world_blox.WorldBloxCollision`. + + Args: + camera_observation: Camera observation to add to the world collision checker. + obstacle_name: Name of the obstacle/layer to add the camera frame to. + """ self.world_coll_checker.add_camera_frame(camera_observation, obstacle_name) def process_camera_frames(self, obstacle_name: Optional[str] = None, process_aux: bool = False): + """Process camera frames for collision checking. + + Only supported by :py:class:`~curobo.geom.sdf.world_blox.WorldBloxCollision`. + + Args: + obstacle_name: Name of the obstacle/layer to process the camera frames for. If None, + processes camera frames on all supported obstacles. + process_aux: Process auxillary information such as mesh integration in nvblox. This is + not required for collision checking and is only needed for debugging. Default is + False to reduce computation time. + """ self.world_coll_checker.process_camera_frames(obstacle_name, process_aux=process_aux) def attach_bounding_box_from_blox_to_robot( @@ -2833,18 +2369,11 @@ class MotionGen(MotionGenConfig): ): """Attach the voxels in a blox layer to robot's link. - NOTE: This is not currently implemented. + .. note:: + This is not currently implemented. - Args: - joint_state (JointState): _description_ - bounding_box (Cuboid): _description_ - blox_layer_name (Optional[str], optional): _description_. Defaults to None. - surface_sphere_radius (float, optional): _description_. Defaults to 0.001. - link_name (str, optional): _description_. Defaults to "attached_object". - sphere_fit_type (SphereFitType, optional): _description_. Defaults to SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE. - voxelize_method (str, optional): _description_. Defaults to "ray". - world_objects_pose_offset (Optional[Pose], optional): _description_. Defaults to None. """ + log_error("Not implemented") kin_state = self.compute_kinematics(joint_state) ee_pose = kin_state.ee_pose # w_T_ee if world_objects_pose_offset is not None: @@ -2895,54 +2424,27 @@ class MotionGen(MotionGenConfig): ): """Attach an object to robot's link. The object to be attached is not in the world model. - Args: - joint_state (JointState): _description_ - obstacle (Obstacle): _description_ - surface_sphere_radius (float, optional): _description_. Defaults to 0.001. - link_name (str, optional): _description_. Defaults to "attached_object". - sphere_fit_type (SphereFitType, optional): _description_. Defaults to SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE. - voxelize_method (str, optional): _description_. Defaults to "ray". - world_objects_pose_offset (Optional[Pose], optional): _description_. Defaults to None. + Deprecated. Use :meth:`MotionGen.attach_external_objects_to_robot` instead. + """ - log_info("MG: Attach objects to robot") - kin_state = self.compute_kinematics(joint_state) - ee_pose = kin_state.ee_pose # w_T_ee - if world_objects_pose_offset is not None: - # add offset from ee: - ee_pose = world_objects_pose_offset.inverse().multiply(ee_pose) - # new ee_pose: - # w_T_ee = offset_T_w * w_T_ee - # ee_T_w - ee_pose = ee_pose.inverse() # ee_T_w to multiply all objects later - max_spheres = self.robot_cfg.kinematics.kinematics_config.get_number_of_spheres(link_name) - n_spheres = max_spheres - sphere_tensor = torch.zeros((max_spheres, 4)) - sphere_tensor[:, 3] = -10.0 - sph_list = [] - if n_spheres == 0: - log_error("MG: No spheres found") - sph = obstacle.get_bounding_spheres( - n_spheres, - surface_sphere_radius, - pre_transform_pose=ee_pose, - tensor_args=self.tensor_args, - fit_type=sphere_fit_type, + log_warn("Deprecated. Use attach_external_objects_to_robot instead") + return self.attach_external_objects_to_robot( + joint_state=joint_state, + external_objects=[obstacle], + surface_sphere_radius=surface_sphere_radius, + link_name=link_name, + sphere_fit_type=sphere_fit_type, voxelize_method=voxelize_method, + world_objects_pose_offset=world_objects_pose_offset, ) - sph_list += [s.position + [s.radius] for s in sph] - - log_info("MG: Computed spheres for attach objects to robot") - - spheres = self.tensor_args.to_device(torch.as_tensor(sph_list)) - - if spheres.shape[0] > max_spheres: - spheres = spheres[: spheres.shape[0]] - sphere_tensor[: spheres.shape[0], :] = spheres.contiguous() - - self.attach_spheres_to_robot(sphere_tensor=sphere_tensor, link_name=link_name) def detach_object_from_robot(self, link_name: str = "attached_object") -> None: - return self.detach_spheres_from_robot(link_name) + """Detach object from robot's link. + + Args: + link_name: Name of the link. + """ + self.detach_spheres_from_robot(link_name) def attach_spheres_to_robot( self, @@ -2953,100 +2455,47 @@ class MotionGen(MotionGenConfig): """Attach spheres to robot's link. Args: - sphere_radius: _description_. Defaults to None. - sphere_tensor: _description_. Defaults to None. - link_name: _description_. Defaults to "attached_object". + sphere_radius: Radius of the spheres. Set to None if :attr:`sphere_tensor` is provided. + sphere_tensor: Sphere x, y, z, r tensor. + link_name: Name of the link to attach the spheres to. Note that this link should + already have pre-allocated spheres. """ self.robot_cfg.kinematics.kinematics_config.attach_object( sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name ) def detach_spheres_from_robot(self, link_name: str = "attached_object") -> None: + """Detach spheres from a robot's link. + + Args: + link_name: Name of the link. + """ self.robot_cfg.kinematics.kinematics_config.detach_object(link_name) 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 add_running_pose_constraint( - self, - lock_x: bool = False, - lock_y: bool = False, - lock_z: bool = False, - lock_rx: bool = False, - lock_ry: bool = False, - lock_rz: bool = False, - ): - raise NotImplementedError() - - def remove_running_pose_constraint(self): - raise NotImplementedError() - - def run_finetune_trajopt( - self, - start_state: JointState, - goal_pose: Pose, - traj_solution: JointState, - traj_dt: Union[float, torch.Tensor, None], - max_attempts: int = 1, - ): - # NOTE: Currently only works for single environment. Need to rewrite for all modes - # finetunes solution - if traj_dt is not None: - self.finetune_trajopt_solver.update_solver_dt(traj_dt.item()) - - # call trajopt with seed: - - # NOTE: currently runs only one seed - time_dict = { - "solve_time": 0, - "ik_time": 0, - "graph_time": 0, - "trajopt_time": 0, - "trajopt_attempts": 0, - } - result = None - # goal_state = JointState.from_position( - # traj_solution.position[..., -2:-1, :].clone(), joint_names=start_state.joint_names - # ) - goal = Goal( - goal_pose=goal_pose, - # goal_state=goal_state, - current_state=start_state, - ) - - for n in range(max_attempts): - traj_result = self.finetune_trajopt_solver.solve_single(goal, traj_solution) - time_dict["trajopt_time"] += traj_result.solve_time - time_dict["trajopt_attempts"] = n - if result is None: - result = MotionGenResult(success=traj_result.success) - - if traj_result.success.item(): - break - - if self.store_debug_in_result: - result.debug_info = {"trajopt_result": traj_result} - - result.position_error = traj_result.position_error - result.rotation_error = traj_result.rotation_error - result.cspace_error = traj_result.cspace_error - result.optimized_dt = traj_result.optimized_dt - result.interpolated_plan = traj_result.interpolated_solution - result.optimized_plan = traj_result.solution - result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep - result.trajopt_time = time_dict["trajopt_time"] - return result - @property def world_model(self) -> WorldConfig: + """Get the world model used for collision checking.""" return self.world_coll_checker.world_model @property def world_collision(self) -> WorldCollision: + """Get the shared instance of world collision checker.""" return self.world_coll_checker @property def project_pose_to_goal_frame(self) -> bool: + """Check if the pose cost metric is projected to goal frame.""" return self.trajopt_solver.rollout_fn.goal_cost.project_distance def update_interpolation_type( @@ -3055,6 +2504,13 @@ class MotionGen(MotionGenConfig): update_graph: bool = True, update_trajopt: bool = True, ): + """Update interpolation type for motion generation. + + Args: + interpolation_type: Type of interpolation to use. + update_graph: Update graph planner with the new interpolation type. + update_trajopt: Update trajectory optimization solvers with the new interpolation type. + """ if update_graph: self.graph_planner.interpolation_type = interpolation_type if update_trajopt: @@ -3065,6 +2521,17 @@ class MotionGen(MotionGenConfig): def update_locked_joints( self, lock_joints: Dict[str, float], robot_config_dict: Union[str, Dict[Any]] ): + """Update locked joints in the robot configuration. + + Use this function to update the joint values of currently locked joints between + planning calls. This function can also be used to change which joints are locked, however + this is only supported when the number of locked joints is the same as the original + robot configuration as the kinematics tensors are pre-allocated. + + Args: + lock_joints: Dictionary of joint names and values to lock. + robot_config_dict: Robot configuration dictionary or path to robot configuration file. + """ if isinstance(robot_config_dict, str): robot_config_dict = load_yaml(join_path(get_robot_configs_path(), robot_config_dict))[ "robot_cfg" @@ -3073,11 +2540,1358 @@ class MotionGen(MotionGenConfig): robot_config_dict = robot_config_dict["robot_cfg"] robot_config_dict["kinematics"]["lock_joints"] = lock_joints robot_cfg = RobotConfig.from_dict(robot_config_dict, self.tensor_args) - - # make sure the new robot config and current have the same joint limits: - new_joint_limits = robot_cfg.kinematics.get_joint_limits() - current_joint_limits = self.robot_cfg.kinematics.get_joint_limits() - # if new_joint_limits != current_joint_limits: - # log_error("Joint limits are different, reinstance motion gen") - self.kinematics.update_kinematics_config(robot_cfg.kinematics.kinematics_config) + + @profiler.record_function("motion_gen/ik") + def _solve_ik_from_solve_state( + self, + goal_pose: Pose, + solve_state: ReacherSolveState, + start_state: JointState, + use_nn_seed: bool, + partial_ik_opt: bool, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve inverse kinematics from solve state, used by motion generation planning call. + + Args: + goal_pose: Goal Pose for the end-effector. + solve_state: Solve state for motion generation. + start_state: Start joint configuration of the robot. + use_nn_seed: Use seed from a neural network. Not implemented. + partial_ik_opt: Only run 50 iterations of inverse kinematics. + link_poses: Goal Poses of any other link in the robot that was specified in + :meth:`curobo.types.robot.RobotConfig.kinematics.link_names`. + + Returns: + IKResult: Result of inverse kinematics. + """ + newton_iters = None + if partial_ik_opt: + newton_iters = self.partial_ik_iters + ik_result = self.ik_solver.solve_any( + solve_state.solve_type, + goal_pose, + start_state.position.view(-1, self._dof), + start_state.position.view(-1, 1, self._dof), + solve_state.num_trajopt_seeds, + solve_state.num_ik_seeds, + use_nn_seed, + newton_iters, + link_poses, + ) + return ik_result + + @profiler.record_function("motion_gen/trajopt_solve") + def _solve_trajopt_from_solve_state( + self, + goal: Goal, + solve_state: ReacherSolveState, + act_seed: Optional[JointState] = None, + use_nn_seed: bool = False, + return_all_solutions: bool = False, + seed_success: Optional[torch.Tensor] = None, + newton_iters: Optional[int] = None, + trajopt_instance: Optional[TrajOptSolver] = None, + num_seeds_override: Optional[int] = None, + ) -> TrajOptResult: + """Solve trajectory optimization from solve state, used by motion generation planning call. + + Args: + goal: Goal containing Pose/Joint targets, current robot state and any other information. + solve_state: Solve state for motion generation. + act_seed: Seed to run trajectory optimization. + use_nn_seed: Use neural network seed for solving trajectory optimization. This is not + implemented. + return_all_solutions: Return all solutions found by trajectory optimization. + seed_success: Success tensor for seeds. + newton_iters: Override Newton iterations to run for solving trajectory optimization. + trajopt_instance: Instance of TrajOptSolver to use for solving trajectory optimization. + num_seeds_override: Override number of seeds to use for solving trajectory optimization. + + Returns: + TrajOptResult: Result of trajectory optimization. + """ + if trajopt_instance is None: + trajopt_instance = self.trajopt_solver + if num_seeds_override is None: + num_seeds_override = solve_state.num_trajopt_seeds + traj_result = trajopt_instance.solve_any( + solve_state.solve_type, + goal, + act_seed, + use_nn_seed, + return_all_solutions, + num_seeds_override, + seed_success, + newton_iters=newton_iters, + ) + return traj_result + + def _get_solve_state( + self, + solve_type: ReacherSolveType, + plan_config: MotionGenPlanConfig, + goal_pose: Pose, + start_state: JointState, + ) -> ReacherSolveState: + """Generate solve state for motion generation based on planning type and configuration. + + MotionGen creates a :class:`ReacherSolveState` for every planning call to keep track of + planning parameters such as number of seeds, batch size, solve type, etc. This solve state + is then compared with existing solve state to determine if solvers (IK, TrajOpt) need to be + re-initialized. Note that changing solve state is not supported when + :attr:`MotionGen.use_cuda_graph` is enabled. + + Args: + solve_type: Type of reacher problem to solve. + plan_config: Planning configuration for motion generation. + goal_pose: Goal Pose for the end-effector. + start_state: Start joint configuration of the robot. + + Raises: + ValueError: If the solve type is not supported. + + Returns: + ReacherSolveState: Solve state for motion generation. + """ + + num_ik_seeds = ( + self.ik_seeds if plan_config.num_ik_seeds is None else plan_config.num_ik_seeds + ) + num_trajopt_seeds = ( + self.trajopt_seeds + if plan_config.num_trajopt_seeds is None + else plan_config.num_trajopt_seeds + ) + + num_graph_seeds = ( + self.trajopt_seeds if plan_config.num_graph_seeds is None else num_trajopt_seeds + ) + solve_state = None + if solve_type == ReacherSolveType.SINGLE: + solve_state = ReacherSolveState( + solve_type, + num_ik_seeds=num_ik_seeds, + num_trajopt_seeds=num_trajopt_seeds, + num_graph_seeds=num_graph_seeds, + batch_size=1, + n_envs=1, + n_goalset=1, + ) + elif solve_type == ReacherSolveType.GOALSET: + solve_state = ReacherSolveState( + solve_type, + num_ik_seeds=num_ik_seeds, + num_trajopt_seeds=num_trajopt_seeds, + num_graph_seeds=num_graph_seeds, + batch_size=1, + n_envs=1, + n_goalset=goal_pose.n_goalset, + ) + elif solve_type == ReacherSolveType.BATCH: + solve_state = ReacherSolveState( + solve_type, + num_ik_seeds=num_ik_seeds, + num_trajopt_seeds=num_trajopt_seeds, + num_graph_seeds=num_graph_seeds, + batch_size=goal_pose.batch, + n_envs=1, + n_goalset=1, + ) + elif solve_type == ReacherSolveType.BATCH_GOALSET: + solve_state = ReacherSolveState( + solve_type, + num_ik_seeds=num_ik_seeds, + num_trajopt_seeds=num_trajopt_seeds, + num_graph_seeds=num_graph_seeds, + batch_size=goal_pose.batch, + n_envs=1, + n_goalset=goal_pose.n_goalset, + ) + elif solve_type == ReacherSolveType.BATCH_ENV: + solve_state = ReacherSolveState( + solve_type, + num_ik_seeds=num_ik_seeds, + num_trajopt_seeds=num_trajopt_seeds, + num_graph_seeds=num_graph_seeds, + batch_size=goal_pose.batch, + n_envs=goal_pose.batch, + n_goalset=1, + ) + elif solve_type == ReacherSolveType.BATCH_ENV_GOALSET: + solve_state = ReacherSolveState( + solve_type, + num_ik_seeds=num_ik_seeds, + num_trajopt_seeds=num_trajopt_seeds, + num_graph_seeds=num_graph_seeds, + batch_size=goal_pose.batch, + n_envs=goal_pose.batch, + n_goalset=goal_pose.n_goalset, + ) + else: + raise ValueError("Unsupported Solve type") + return solve_state + + def _plan_attempts( + self, + solve_state: ReacherSolveState, + start_state: JointState, + goal_pose: Pose, + plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: List[Pose] = None, + ): + """Call many planning attempts for a given reacher solve state. + + Args: + solve_state: Reacher solve state for planning. + start_state: Start joint state for planning. + goal_pose: Goal pose to reach for end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for other links in the robot. + + Returns: + MotionGenResult: Result of planning. + """ + start_time = time.time() + if plan_config.pose_cost_metric is not None: + valid_query = self.update_pose_cost_metric( + plan_config.pose_cost_metric, start_state, goal_pose + ) + if not valid_query: + result = MotionGenResult( + success=torch.as_tensor([False], device=self.tensor_args.device), + valid_query=valid_query, + status="Invalid Hold Partial Pose", + ) + return result + self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size) + 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") + force_graph = plan_config.enable_graph + partial_ik = plan_config.partial_ik_opt + ik_fail_count = 0 + time_dict = { + "solve_time": 0, + "ik_time": 0, + "graph_time": 0, + "trajopt_time": 0, + "trajopt_attempts": 0, + } + best_status = 0 + if plan_config.finetune_dt_scale is None: + plan_config.finetune_dt_scale = self.finetune_dt_scale + for n in range(plan_config.max_attempts): + log_info("MG Iter: " + str(n)) + result = self._plan_from_solve_state( + solve_state, + start_state, + goal_pose, + plan_config, + link_poses, + ) + time_dict["solve_time"] += result.solve_time + time_dict["ik_time"] += result.ik_time + time_dict["graph_time"] += result.graph_time + time_dict["trajopt_time"] += result.trajopt_time + time_dict["trajopt_attempts"] += result.trajopt_attempts + if ( + result.status == MotionGenStatus.IK_FAIL and plan_config.ik_fail_return is not None + ): # IF IK fails the first time, we exist assuming the goal is not reachable + ik_fail_count += 1 + best_status = max(best_status, 1) + + if ik_fail_count > plan_config.ik_fail_return: + break + if result.success[0].item(): + break + + if result.status == MotionGenStatus.FINETUNE_TRAJOPT_FAIL: + plan_config.finetune_dt_scale *= ( + plan_config.finetune_dt_decay**plan_config.finetune_attempts + ) + plan_config.finetune_dt_scale = min(plan_config.finetune_dt_scale, 1.25) + if plan_config.enable_graph_attempt is not None and ( + n >= plan_config.enable_graph_attempt - 1 + and result.status + in [MotionGenStatus.TRAJOPT_FAIL, MotionGenStatus.FINETUNE_TRAJOPT_FAIL] + and not plan_config.enable_graph + ): + plan_config.enable_graph = True + plan_config.partial_ik_opt = False + if plan_config.disable_graph_attempt is not None and ( + n >= plan_config.disable_graph_attempt - 1 + and result.status + in [ + MotionGenStatus.TRAJOPT_FAIL, + MotionGenStatus.GRAPH_FAIL, + MotionGenStatus.FINETUNE_TRAJOPT_FAIL, + ] + and not force_graph + ): + plan_config.enable_graph = False + plan_config.partial_ik_opt = partial_ik + if result.status in [MotionGenStatus.TRAJOPT_FAIL]: + best_status = 3 + elif result.status in [MotionGenStatus.GRAPH_FAIL]: + best_status = 2 + if time.time() - start_time > plan_config.timeout: + break + if not result.valid_query: + result.status = MotionGenStatus.INVALID_QUERY + break + if n == 10: + self.reset_seed() + log_warn("Couldn't find solution with 10 attempts, resetting seeds") + + result.solve_time = time_dict["solve_time"] + result.ik_time = time_dict["ik_time"] + result.graph_time = time_dict["graph_time"] + result.trajopt_time = time_dict["trajopt_time"] + result.trajopt_attempts = time_dict["trajopt_attempts"] + result.attempts = n + 1 + torch.cuda.synchronize(device=self.tensor_args.device) + if plan_config.pose_cost_metric is not None: + self.update_pose_cost_metric(PoseCostMetric.reset_metric()) + if plan_config.time_dilation_factor is not None and torch.count_nonzero(result.success) > 0: + result.retime_trajectory( + plan_config.time_dilation_factor, + interpolation_kind=self.finetune_trajopt_solver.interpolation_type, + ) + + result.total_time = time.time() - start_time + return result + + def _plan_batch_attempts( + self, + solve_state: ReacherSolveState, + start_state: JointState, + goal_pose: Pose, + plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Optional[Dict[str, Pose]] = None, + ): + """Plan batch attempts for a given reacher solve state. + + Args: + solve_state: Reacher solve state for planning. + start_state: Start joint state for planning. + goal_pose: Goal pose to reach for end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for other links in the robot. + + Returns: + MotionGenResult: Result of batched planning. + """ + start_time = time.time() + goal_pose = goal_pose.clone() + if plan_config.pose_cost_metric is not None: + valid_query = self.update_pose_cost_metric( + plan_config.pose_cost_metric, start_state, goal_pose + ) + if not valid_query: + result = MotionGenResult( + success=torch.as_tensor( + [False for _ in solve_state.batch_size], + device=self.motion_gen.tensor_args.device, + ), + valid_query=valid_query, + status="Invalid Hold Partial Pose", + ) + return result + + 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") + if plan_config.enable_graph: + log_error("Graph Search / Geometric Planner not supported in batch_env mode") + + if plan_config.enable_graph or ( + plan_config.enable_graph_attempt is not None + and plan_config.max_attempts >= plan_config.enable_graph_attempt + ): + log_warn("Batch mode enable graph is only supported with num_graph_seeds==1") + plan_config.num_trajopt_seeds = 1 + plan_config.num_graph_seeds = 1 + solve_state.num_trajopt_seeds = 1 + solve_state.num_graph_seeds = 1 + self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size) + + ik_fail_count = 0 + force_graph = plan_config.enable_graph + partial_ik = plan_config.partial_ik_opt + + time_dict = { + "solve_time": 0, + "ik_time": 0, + "graph_time": 0, + "trajopt_time": 0, + "trajopt_attempts": 0, + } + best_result = None + + for n in range(plan_config.max_attempts): + result = self._plan_from_solve_state_batch( + solve_state, + start_state, + goal_pose, + plan_config, + link_poses=link_poses, + ) + + time_dict["solve_time"] += result.solve_time + time_dict["ik_time"] += result.ik_time + + time_dict["graph_time"] += result.graph_time + time_dict["trajopt_time"] += result.trajopt_time + time_dict["trajopt_attempts"] += result.trajopt_attempts + + # if not all have succeeded, store the successful ones and re attempt: + # TODO: update stored based on error + if best_result is None: + best_result = result.clone() + else: + # get success idx: + idx = torch.nonzero(result.success).reshape(-1) + if len(idx) > 0: + best_result.copy_idx(idx, result) + + if ( + result.status == MotionGenStatus.IK_FAIL and plan_config.ik_fail_return is not None + ): # IF IK fails the first time, we exit assuming the goal is not reachable + ik_fail_count += 1 + if ik_fail_count > plan_config.ik_fail_return: + break + + if ( + torch.count_nonzero(best_result.success) + >= goal_pose.batch * plan_config.success_ratio + ): # we want 90% targets to succeed + best_result.status = None + break + if plan_config.enable_graph_attempt is not None and ( + n >= plan_config.enable_graph_attempt - 1 + and result.status != MotionGenStatus.IK_FAIL + and not plan_config.enable_graph + ): + plan_config.enable_graph = True + plan_config.partial_ik_opt = False + + if plan_config.disable_graph_attempt is not None and ( + n >= plan_config.disable_graph_attempt - 1 + and result.status in [MotionGenStatus.TRAJOPT_FAIL, MotionGenStatus.GRAPH_FAIL] + and not force_graph + ): + plan_config.enable_graph = False + plan_config.partial_ik_opt = partial_ik + + if plan_config.fail_on_invalid_query: + if not result.valid_query: + best_result.valid_query = False + best_result.status = "Invalid Problem" + break + if time.time() - start_time > plan_config.timeout: + break + best_result.solve_time = time_dict["solve_time"] + best_result.ik_time = time_dict["ik_time"] + best_result.graph_time = time_dict["graph_time"] + best_result.trajopt_time = time_dict["trajopt_time"] + best_result.trajopt_attempts = time_dict["trajopt_attempts"] + best_result.attempts = n + 1 + torch.cuda.synchronize(device=self.tensor_args.device) + if plan_config.pose_cost_metric is not None: + self.update_pose_cost_metric(PoseCostMetric.reset_metric()) + + if plan_config.time_dilation_factor is not None and torch.count_nonzero(result.success) > 0: + result.retime_trajectory( + plan_config.time_dilation_factor, + interpolation_kind=self.finetune_trajopt_solver.interpolation_type, + ) + best_result.total_time = time.time() - start_time + return best_result + + def _plan_from_solve_state( + self, + solve_state: ReacherSolveState, + start_state: JointState, + goal_pose: Pose, + plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Optional[Dict[str, Pose]] = None, + ) -> MotionGenResult: + """Plan from a given reacher solve state. + + Args: + solve_state: Reacher solve state for planning. + start_state: Start joint state for planning. + goal_pose: Goal pose to reach for end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for other links in the robot. + + Returns: + MotionGenResult: Result of planning. + """ + trajopt_seed_traj = None + trajopt_seed_success = None + trajopt_newton_iters = None + graph_success = 0 + + if len(start_state.shape) == 1: + log_error("Joint state should be not a vector (dof) should be (bxdof)") + + if goal_pose.shape[0] != 1: + log_error( + "Goal position should be of shape [1, n_goalset, -1], current shape: " + + str(goal_pose.shape) + ) + # plan ik: + + ik_result = self._solve_ik_from_solve_state( + goal_pose, + solve_state, + start_state, + plan_config.use_nn_ik_seed, + plan_config.partial_ik_opt, + link_poses, + ) + + if not plan_config.enable_graph and plan_config.partial_ik_opt: + ik_result.success[:] = True + + # check for success: + result = MotionGenResult( + ik_result.success.view(-1)[0:1], # This is not true for batch mode + ik_time=ik_result.solve_time, + solve_time=ik_result.solve_time, + ) + + if self.store_debug_in_result: + result.debug_info = {"ik_result": ik_result} + ik_success = torch.count_nonzero(ik_result.success) + if ik_success == 0: + result.status = MotionGenStatus.IK_FAIL + return result + + # do graph search: + with profiler.record_function("motion_gen/post_ik"): + ik_out_seeds = solve_state.num_trajopt_seeds + if plan_config.enable_graph: + ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success) + + goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof)[ + :ik_success + ] + start_config = tensor_repeat_seeds(start_state.position, ik_out_seeds) + if plan_config.enable_opt: + self._trajopt_goal_config[:] = ik_result.solution + + # do graph search: + if plan_config.enable_graph: + interpolation_steps = None + if plan_config.enable_opt: + interpolation_steps = self.trajopt_solver.action_horizon + log_info("MG: running GP") + graph_result = self.graph_search(start_config, goal_config, interpolation_steps) + trajopt_seed_success = graph_result.success + + graph_success = torch.count_nonzero(graph_result.success).item() + result.graph_time = graph_result.solve_time + result.solve_time += graph_result.solve_time + if graph_success > 0: + log_info("MG: GP Success") + result.graph_plan = graph_result.interpolated_plan + result.interpolated_plan = graph_result.interpolated_plan + + result.used_graph = True + if plan_config.enable_opt: + trajopt_seed = ( + result.graph_plan.position.view( + 1, # solve_state.batch_size, + graph_success, # solve_state.num_trajopt_seeds, + interpolation_steps, + self._dof, + ) + .transpose(0, 1) + .contiguous() + ) + trajopt_seed_traj = torch.zeros( + (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof), + device=self.tensor_args.device, + dtype=self.tensor_args.dtype, + ) + trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed + trajopt_seed_success = ik_result.success.clone() + trajopt_seed_success[ik_result.success] = graph_result.success + + trajopt_seed_success = trajopt_seed_success.view( + solve_state.batch_size, solve_state.num_trajopt_seeds + ) + trajopt_newton_iters = self.graph_trajopt_iters + else: + _, idx = torch.topk( + graph_result.path_length[graph_result.success], k=1, largest=False + ) + result.interpolated_plan = result.interpolated_plan[idx].squeeze(0) + result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt) + result.optimized_plan = result.interpolated_plan[ + : graph_result.path_buffer_last_tstep[idx.item()] + ] + idx = idx.view(-1) + self._batch_col + result.position_error = ik_result.position_error[ik_result.success][ + graph_result.success + ][idx] + result.rotation_error = ik_result.rotation_error[ik_result.success][ + graph_result.success + ][idx] + result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[ + idx.item() : idx.item() + 1 + ] + result.success = result.success.view(-1)[0:1] + result.success[:] = True + return result + else: + result.success[:] = False + result.status = MotionGenStatus.GRAPH_FAIL + if not graph_result.valid_query: + result.valid_query = False + if self.store_debug_in_result: + result.debug_info["graph_debug"] = graph_result.debug_info + return result + if plan_config.need_graph_success: + return result + + # do trajopt:: + + if plan_config.enable_opt: + with profiler.record_function("motion_gen/setup_trajopt_seeds"): + self._trajopt_goal_config[:, :ik_success] = goal_config + retract_config = None + if plan_config.use_start_state_as_retract: + retract_config = start_state.position.clone() + goal = Goal( + goal_pose=goal_pose, + current_state=start_state, + links_goal_pose=link_poses, + retract_state=retract_config, + ) + + if ( + trajopt_seed_traj is None + or graph_success < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds + ): + goal_config = self._trajopt_goal_config[0] # batch index == 0 + + goal_state = JointState.from_position( + goal_config, + ) + seed_link_poses = None + if link_poses is not None: + seed_link_poses = {} + + for k in link_poses.keys(): + seed_link_poses[k] = link_poses[k].repeat_seeds( + solve_state.num_trajopt_seeds + ) + if goal_pose.shape[0] != 1: + log_error( + "Batch size of goal pose should be 1, current shape: " + + str(goal_pose.shape) + ) + seed_goal = Goal( + goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds), + current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds), + goal_state=goal_state, + links_goal_pose=seed_link_poses, + ) + if trajopt_seed_traj is not None: + trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() + # batch, num_seeds, h, dof + if ( + trajopt_seed_success.shape[1] + < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds + ): + trajopt_seed_success_new = torch.zeros( + (1, solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds), + device=self.tensor_args.device, + dtype=torch.bool, + ) + trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = ( + trajopt_seed_success + ) + trajopt_seed_success = trajopt_seed_success_new + # create seeds here: + trajopt_seed_traj = self.trajopt_solver.get_seed_set( + seed_goal, + trajopt_seed_traj, # batch, num_seeds, h, dof + num_seeds=self.noisy_trajopt_seeds, + batch_mode=solve_state.batch_mode, + seed_success=trajopt_seed_success, + ) + trajopt_seed_traj = trajopt_seed_traj.view( + solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds, + solve_state.batch_size, + self.trajopt_solver.action_horizon, + self._dof, + ).contiguous() + if plan_config.enable_finetune_trajopt: + og_value = self.trajopt_solver.interpolation_type + self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA + with profiler.record_function("motion_gen/trajopt"): + log_info("MG: running TO") + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + trajopt_seed_traj, + num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds, + newton_iters=trajopt_newton_iters, + return_all_solutions=plan_config.parallel_finetune + and plan_config.enable_finetune_trajopt, + ) + if plan_config.enable_finetune_trajopt: + self.trajopt_solver.interpolation_type = og_value + if self.store_debug_in_result: + result.debug_info["trajopt_result"] = traj_result + # run finetune + if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0: + with profiler.record_function("motion_gen/finetune_trajopt"): + seed_traj = traj_result.raw_action.clone() # solution.position.clone() + seed_traj = seed_traj.contiguous() + og_solve_time = traj_result.solve_time + seed_override = 1 + opt_dt = traj_result.optimized_dt + + if plan_config.parallel_finetune: + opt_dt = torch.min(opt_dt[traj_result.success]) + seed_override = solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds + + finetune_time = 0 + for k in range(plan_config.finetune_attempts): + newton_iters = None + + scaled_dt = torch.clamp( + opt_dt + * plan_config.finetune_dt_scale + * (plan_config.finetune_dt_decay ** (k)), + self.trajopt_solver.minimum_trajectory_dt, + ) + if self.optimize_dt: + self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item()) + + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + seed_traj, + trajopt_instance=self.finetune_trajopt_solver, + num_seeds_override=seed_override, + newton_iters=newton_iters, + ) + finetune_time += traj_result.solve_time + if torch.count_nonzero(traj_result.success) > 0: + break + seed_traj = traj_result.optimized_seeds.detach().clone() + newton_iters = 4 + + traj_result.solve_time = finetune_time + + result.finetune_time = traj_result.solve_time + + traj_result.solve_time = og_solve_time + if self.store_debug_in_result: + result.debug_info["finetune_trajopt_result"] = traj_result + elif plan_config.enable_finetune_trajopt: + traj_result.success = traj_result.success[0:1] + # if torch.count_nonzero(result.success) == 0: + result.status = MotionGenStatus.TRAJOPT_FAIL + result.solve_time += traj_result.solve_time + result.finetune_time + result.trajopt_time = traj_result.solve_time + result.trajopt_attempts = 1 + result.success = traj_result.success + + if plan_config.enable_finetune_trajopt and torch.count_nonzero(result.success) == 0: + result.status = MotionGenStatus.FINETUNE_TRAJOPT_FAIL + + result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory( + 0, traj_result.path_buffer_last_tstep[0] + ) + result.interpolation_dt = self.trajopt_solver.interpolation_dt + result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep + result.position_error = traj_result.position_error + result.rotation_error = traj_result.rotation_error + result.optimized_dt = traj_result.optimized_dt + result.optimized_plan = traj_result.solution + result.goalset_index = traj_result.goalset_index + return result + + def _plan_js_from_solve_state( + self, + solve_state: ReacherSolveState, + start_state: JointState, + goal_state: JointState, + plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + ) -> MotionGenResult: + """Plan from a given reacher solve state for joint state. + + Args: + solve_state: Reacher solve state for planning. + start_state: Start joint state for planning. + goal_state: Goal joint state to reach. + plan_config: Planning parameters for motion generation. + + Returns: + MotionGenResult: Result of planning. + """ + trajopt_seed_traj = None + trajopt_seed_success = None + trajopt_newton_iters = self.js_trajopt_solver.newton_iters + + graph_success = 0 + if len(start_state.shape) == 1: + log_error("Joint state should be not a vector (dof) should be (bxdof)") + + result = MotionGenResult(cspace_error=torch.zeros((1), device=self.tensor_args.device)) + if self.store_debug_in_result: + result.debug_info = {} + # do graph search: + if plan_config.enable_graph: + start_config = torch.zeros( + (solve_state.num_graph_seeds, self.js_trajopt_solver.dof), + device=self.tensor_args.device, + dtype=self.tensor_args.dtype, + ) + goal_config = start_config.clone() + start_config[:] = start_state.position + goal_config[:] = goal_state.position + interpolation_steps = None + if plan_config.enable_opt: + interpolation_steps = self.js_trajopt_solver.action_horizon + log_info("MG: running GP") + graph_result = self.graph_search(start_config, goal_config, interpolation_steps) + trajopt_seed_success = graph_result.success + + graph_success = torch.count_nonzero(graph_result.success).item() + result.graph_time = graph_result.solve_time + result.solve_time += graph_result.solve_time + if graph_success > 0: + result.graph_plan = graph_result.interpolated_plan + result.interpolated_plan = graph_result.interpolated_plan + + result.used_graph = True + if plan_config.enable_opt: + trajopt_seed = ( + result.graph_plan.position.view( + 1, # solve_state.batch_size, + graph_success, # solve_state.num_trajopt_seeds, + interpolation_steps, + self._dof, + ) + .transpose(0, 1) + .contiguous() + ) + trajopt_seed_traj = torch.zeros( + (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof), + device=self.tensor_args.device, + dtype=self.tensor_args.dtype, + ) + trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed + trajopt_seed_success = graph_result.success + + trajopt_seed_success = trajopt_seed_success.view( + 1, solve_state.num_trajopt_seeds + ) + trajopt_newton_iters = self.graph_trajopt_iters + else: + _, idx = torch.topk( + graph_result.path_length[graph_result.success], k=1, largest=False + ) + result.interpolated_plan = result.interpolated_plan[idx].squeeze(0) + result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt) + result.optimized_plan = result.interpolated_plan[ + : graph_result.path_buffer_last_tstep[idx.item()] + ] + idx = idx.view(-1) + self._batch_col + result.cspace_error = torch.zeros((1), device=self.tensor_args.device) + + result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[ + idx.item() : idx.item() + 1 + ] + result.success = torch.as_tensor([True], device=self.tensor_args.device) + return result + else: + result.success = torch.as_tensor([False], device=self.tensor_args.device) + result.status = MotionGenStatus.GRAPH_FAIL + if not graph_result.valid_query: + result.valid_query = False + if self.store_debug_in_result: + result.debug_info["graph_debug"] = graph_result.debug_info + return result + if plan_config.need_graph_success: + return result + + # do trajopt: + if plan_config.enable_opt: + with profiler.record_function("motion_gen/setup_trajopt_seeds"): + + goal = Goal( + current_state=start_state, + goal_state=goal_state, + ) + + if trajopt_seed_traj is None or graph_success < solve_state.num_trajopt_seeds * 1: + seed_goal = Goal( + current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds), + goal_state=goal_state.repeat_seeds(solve_state.num_trajopt_seeds), + ) + if trajopt_seed_traj is not None: + trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() + # batch, num_seeds, h, dof + if trajopt_seed_success.shape[1] < self.js_trajopt_solver.num_seeds: + trajopt_seed_success_new = torch.zeros( + (1, solve_state.num_trajopt_seeds), + device=self.tensor_args.device, + dtype=torch.bool, + ) + trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = ( + trajopt_seed_success + ) + trajopt_seed_success = trajopt_seed_success_new + # create seeds here: + trajopt_seed_traj = self.js_trajopt_solver.get_seed_set( + seed_goal, + trajopt_seed_traj, # batch, num_seeds, h, dof + num_seeds=1, + batch_mode=False, + seed_success=trajopt_seed_success, + ) + trajopt_seed_traj = ( + trajopt_seed_traj.view( + self.js_trajopt_solver.num_seeds * 1, + 1, + self.trajopt_solver.action_horizon, + self._dof, + ) + .contiguous() + .clone() + ) + if plan_config.enable_finetune_trajopt: + og_value = self.trajopt_solver.interpolation_type + self.js_trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA + with profiler.record_function("motion_gen/trajopt"): + log_info("MG: running TO") + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + trajopt_seed_traj, + num_seeds_override=solve_state.num_trajopt_seeds, + newton_iters=trajopt_newton_iters + 2, + return_all_solutions=plan_config.enable_finetune_trajopt, + trajopt_instance=self.js_trajopt_solver, + ) + if plan_config.enable_finetune_trajopt: + self.trajopt_solver.interpolation_type = og_value + if self.store_debug_in_result: + result.debug_info["trajopt_result"] = traj_result + if torch.count_nonzero(traj_result.success) == 0: + result.status = MotionGenStatus.TRAJOPT_FAIL + # run finetune + if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0: + with profiler.record_function("motion_gen/finetune_trajopt"): + seed_traj = traj_result.raw_action.clone() # solution.position.clone() + seed_traj = seed_traj.contiguous() + og_solve_time = traj_result.solve_time + + scaled_dt = torch.clamp( + torch.max(traj_result.optimized_dt[traj_result.success]), + self.trajopt_solver.minimum_trajectory_dt, + ) + og_dt = self.js_trajopt_solver.solver_dt.clone() + self.js_trajopt_solver.update_solver_dt(scaled_dt.item()) + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + seed_traj, + trajopt_instance=self.js_trajopt_solver, + num_seeds_override=solve_state.num_trajopt_seeds, + newton_iters=trajopt_newton_iters + 4, + ) + self.js_trajopt_solver.update_solver_dt(og_dt) + + result.finetune_time = traj_result.solve_time + + traj_result.solve_time = og_solve_time + if self.store_debug_in_result: + result.debug_info["finetune_trajopt_result"] = traj_result + if torch.count_nonzero(traj_result.success) == 0: + result.status = MotionGenStatus.FINETUNE_TRAJOPT_FAIL + elif plan_config.enable_finetune_trajopt: + traj_result.success = traj_result.success[0:1] + result.solve_time += traj_result.solve_time + result.finetune_time + result.trajopt_time = traj_result.solve_time + result.trajopt_attempts = 1 + result.success = traj_result.success + + result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory( + 0, traj_result.path_buffer_last_tstep[0] + ) + + result.interpolation_dt = self.trajopt_solver.interpolation_dt + result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep + result.cspace_error = traj_result.cspace_error + result.optimized_dt = traj_result.optimized_dt + result.optimized_plan = traj_result.solution + result.goalset_index = traj_result.goalset_index + + return result + + def _plan_from_solve_state_batch( + self, + solve_state: ReacherSolveState, + start_state: JointState, + goal_pose: Pose, + plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), + link_poses: Optional[Dict[str, Pose]] = None, + ) -> MotionGenResult: + """Plan from a given reacher solve state in batch mode. + + Args: + solve_state: Reacher solve state for planning. + start_state: Start joint state for planning. + goal_pose: Goal poses to reach for end-effector. + plan_config: Planning parameters for motion generation. + link_poses: Goal poses for other links in the robot. + + Returns: + MotionGenResult: Result of planning. + """ + self._trajopt_goal_config[:] = self.get_retract_config().view(1, 1, self._dof) + trajopt_seed_traj = None + trajopt_seed_success = None + trajopt_newton_iters = None + graph_success = 0 + + # plan ik: + ik_result = self._solve_ik_from_solve_state( + goal_pose, + solve_state, + start_state, + plan_config.use_nn_ik_seed, + plan_config.partial_ik_opt, + link_poses, + ) + + if not plan_config.enable_graph and plan_config.partial_ik_opt: + ik_result.success[:] = True + + # check for success: + result = MotionGenResult( + ik_result.success, + position_error=ik_result.position_error, + rotation_error=ik_result.rotation_error, + ik_time=ik_result.solve_time, + solve_time=ik_result.solve_time, + debug_info={}, + # goalset_index=ik_result.goalset_index, + ) + + ik_success = torch.count_nonzero(ik_result.success) + if ik_success == 0: + result.status = MotionGenStatus.IK_FAIL + result.success = result.success[:, 0] + return result + + # do graph search: + ik_out_seeds = solve_state.num_trajopt_seeds + if plan_config.enable_graph: + ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success) + + # if not plan_config.enable_opt and plan_config.enable_graph: + # self.graph_planner.interpolation_steps = self.interpolation_steps + # self.graph_planner.interpolation_type = self.interpolation_type + # elif plan_config.enable_graph: + # self.graph_planner.interpolation_steps = self.trajopt_solver.traj_tsteps + # self.graph_planner.interpolation_type = InterpolateType.LINEAR + goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof) + + # get shortest path + if plan_config.enable_graph: + interpolation_steps = None + if plan_config.enable_opt: + interpolation_steps = self.trajopt_solver.action_horizon + + start_graph_state = start_state.repeat_seeds(ik_out_seeds) + start_config = start_graph_state.position[ik_result.success.view(-1)].view( + -1, self.ik_solver.dof + ) + graph_result = self.graph_search(start_config, goal_config, interpolation_steps) + graph_success = torch.count_nonzero(graph_result.success).item() + + result.graph_time = graph_result.solve_time + result.solve_time += graph_result.solve_time + if graph_success > 0: + # path = graph_result.interpolated_plan + result.graph_plan = graph_result.interpolated_plan + result.interpolated_plan = graph_result.interpolated_plan + result.used_graph = True + + if plan_config.enable_opt: + trajopt_seed = result.graph_plan.position.view( + graph_success, # solve_state.num_trajopt_seeds, + interpolation_steps, + self._dof, + ).contiguous() + trajopt_seed_traj = torch.zeros( + (1, trajopt_seed.shape[0], self.trajopt_solver.action_horizon, self._dof), + device=self.tensor_args.device, + dtype=self.tensor_args.dtype, + ) + trajopt_seed_traj[0, :, :interpolation_steps, :] = trajopt_seed + trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() + trajopt_seed_success = ik_result.success.clone() + trajopt_seed_success[ik_result.success] = graph_result.success + + trajopt_seed_success = trajopt_seed_success.view( + solve_state.num_trajopt_seeds, solve_state.batch_size + ) + trajopt_newton_iters = self.graph_trajopt_iters + + else: + ik_success = ik_result.success.view(-1).clone() + + # only some might be successful: + + g_dim = torch.nonzero(ik_success).view(-1)[graph_result.success] + + self._batch_graph_search_buffer.copy_at_index( + graph_result.interpolated_plan, g_dim + ) + + # result.graph_plan = JointState.from_position( + # self._batch_graph_search_buffer, + # joint_names=graph_result.interpolated_plan.joint_names, + # ) + result.interpolated_plan = self._batch_graph_search_buffer + g_dim = g_dim.cpu().squeeze().tolist() + for x, x_val in enumerate(g_dim): + self._batch_path_buffer_last_tstep[x_val] = ( + graph_result.path_buffer_last_tstep[x] + ) + result.path_buffer_last_tstep = self._batch_path_buffer_last_tstep + result.optimized_plan = result.interpolated_plan + result.optimized_dt = torch.as_tensor( + [ + self.interpolation_dt + for i in range(result.interpolated_plan.position.shape[0]) + ], + device=self.tensor_args.device, + dtype=self.tensor_args.dtype, + ) + result.success = result.success.view(-1).clone() + result.success[ik_success][graph_result.success] = True + return result + + else: + result.success[:] = False + result.success = result.success[:, 0] + result.status = MotionGenStatus.GRAPH_FAIL + if not graph_result.valid_query: + result.valid_query = False + if self.store_debug_in_result: + result.debug_info = {"graph_debug": graph_result.debug_info} + return result + + if plan_config.enable_opt: + # get goal configs based on ik success: + self._trajopt_goal_config[ik_result.success] = goal_config + + goal_config = self._trajopt_goal_config # batch index == 0 + + goal = Goal( + goal_pose=goal_pose, + current_state=start_state, + links_goal_pose=link_poses, + ) + # generate seeds: + if trajopt_seed_traj is None or ( + plan_config.enable_graph and graph_success < solve_state.batch_size + ): + seed_link_poses = None + if link_poses is not None: + seed_link_poses = {} + + for k in link_poses.keys(): + seed_link_poses[k] = link_poses[k].repeat_seeds( + solve_state.num_trajopt_seeds + ) + seed_goal = Goal( + goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds), + current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds), + goal_state=JointState.from_position(goal_config.view(-1, self._dof)), + links_goal_pose=seed_link_poses, + ) + if trajopt_seed_traj is not None: + trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous() + + # create seeds here: + trajopt_seed_traj = self.trajopt_solver.get_seed_set( + seed_goal, + trajopt_seed_traj, # batch, num_seeds, h, dof + num_seeds=1, + batch_mode=solve_state.batch_mode, + seed_success=trajopt_seed_success, + ) + trajopt_seed_traj = trajopt_seed_traj.view( + solve_state.num_trajopt_seeds, + solve_state.batch_size, + self.trajopt_solver.action_horizon, + self._dof, + ).contiguous() + if plan_config.enable_finetune_trajopt: + og_value = self.trajopt_solver.interpolation_type + self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA + + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + trajopt_seed_traj, + newton_iters=trajopt_newton_iters, + return_all_solutions=plan_config.enable_finetune_trajopt, + ) + + # output of traj result will have 1 solution per batch + + # run finetune opt on 1 solution per batch: + if plan_config.enable_finetune_trajopt: + self.trajopt_solver.interpolation_type = og_value + if self.store_debug_in_result: + result.debug_info["trajopt_result"] = traj_result + + # run finetune + if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0: + with profiler.record_function("motion_gen/finetune_trajopt"): + seed_traj = traj_result.raw_action.clone() # solution.position.clone() + seed_traj = seed_traj.contiguous() + og_solve_time = traj_result.solve_time + + scaled_dt = torch.clamp( + torch.max(traj_result.optimized_dt[traj_result.success]) + * self.finetune_dt_scale, + self.trajopt_solver.minimum_trajectory_dt, + ) + self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item()) + + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + seed_traj, + trajopt_instance=self.finetune_trajopt_solver, + num_seeds_override=solve_state.num_trajopt_seeds, + ) + + result.finetune_time = traj_result.solve_time + + traj_result.solve_time = og_solve_time + if self.store_debug_in_result: + result.debug_info["finetune_trajopt_result"] = traj_result + elif plan_config.enable_finetune_trajopt and len(traj_result.success.shape) == 2: + traj_result.success = traj_result.success[:, 0] + + result.success = traj_result.success + + result.interpolated_plan = traj_result.interpolated_solution + result.solve_time += traj_result.solve_time + result.trajopt_time = traj_result.solve_time + result.position_error = traj_result.position_error + result.rotation_error = traj_result.rotation_error + result.cspace_error = traj_result.cspace_error + result.goalset_index = traj_result.goalset_index + result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep + result.optimized_plan = traj_result.solution + result.optimized_dt = traj_result.optimized_dt + if torch.count_nonzero(traj_result.success) == 0: + result.status = MotionGenStatus.TRAJOPT_FAIL + result.success[:] = False + if self.store_debug_in_result: + result.debug_info = {"trajopt_result": traj_result} + return result + + def plan( + self, + start_state: JointState, + goal_pose: Pose, + enable_graph: bool = True, + enable_opt: bool = True, + use_nn_ik_seed: bool = False, + need_graph_success: bool = False, + max_attempts: int = 100, + timeout: float = 10.0, + enable_graph_attempt: Optional[int] = None, + disable_graph_attempt: Optional[int] = None, + trajopt_attempts: int = 1, + ik_fail_return: Optional[int] = None, + partial_ik_opt: bool = False, + num_ik_seeds: Optional[int] = None, + num_graph_seeds: Optional[int] = None, + num_trajopt_seeds: Optional[int] = None, + ): + """Deprecated method. Use :meth:`MotionGen.plan_single` or others instead.""" + + log_warn("Deprecated method. Use MotionGen.plan_single or others instead.") + plan_config = MotionGenPlanConfig( + enable_graph, + enable_opt, + use_nn_ik_seed, + need_graph_success, + max_attempts, + timeout, + enable_graph_attempt, + disable_graph_attempt, + ik_fail_return, + partial_ik_opt, + num_ik_seeds, + num_graph_seeds, + num_trajopt_seeds, + ) + result = self.plan_single(start_state, goal_pose, plan_config) + return result + + def batch_plan( + self, + start_state: JointState, + goal_pose: Pose, + enable_graph: bool = True, + enable_opt: bool = True, + use_nn_ik_seed: bool = False, + need_graph_success: bool = False, + max_attempts: int = 1, + timeout: float = 10.0, + enable_graph_attempt: Optional[int] = None, + disable_graph_attempt: Optional[int] = None, + success_ratio: float = 1.0, + ik_fail_return: Optional[int] = None, + fail_on_invalid_query: bool = False, + partial_ik_opt: bool = False, + num_ik_seeds: Optional[int] = None, + num_graph_seeds: Optional[int] = None, + num_trajopt_seeds: Optional[int] = None, + ): + """Deprecated method. Use :meth:`MotionGen.plan_batch` or others instead.""" + + log_warn("Deprecated method. Use MotionGen.plan_batch or others instead.") + + plan_config = MotionGenPlanConfig( + enable_graph, + enable_opt, + use_nn_ik_seed, + need_graph_success, + max_attempts, + timeout, + enable_graph_attempt, + disable_graph_attempt, + ik_fail_return, + partial_ik_opt, + num_ik_seeds, + num_graph_seeds, + num_trajopt_seeds, + success_ratio=success_ratio, + fail_on_invalid_query=fail_on_invalid_query, + ) + result = self.plan_batch(start_state, goal_pose, plan_config) + return result diff --git a/src/curobo/wrap/reacher/mpc.py b/src/curobo/wrap/reacher/mpc.py index 378756d..28919b5 100644 --- a/src/curobo/wrap/reacher/mpc.py +++ b/src/curobo/wrap/reacher/mpc.py @@ -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 + +

+ +

+ + +""" # Standard Library import time @@ -19,6 +43,7 @@ from typing import Dict, Optional, Union import torch # CuRobo +from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel from curobo.geom.sdf.utils import create_collision_checker from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig from curobo.geom.types import WorldConfig @@ -44,22 +69,31 @@ from curobo.wrap.wrap_mpc import WrapConfig, WrapMpc @dataclass class MpcSolverConfig: + """Configuration dataclass for MPC.""" + + #: MPC Solver. solver: WrapMpc + + #: World Collision Checker. world_coll_checker: Optional[WorldCollision] = None + + #: Numeric precision and device to run computations. tensor_args: TensorDeviceType = TensorDeviceType() + + #: Capture full step in MPC as a single CUDA graph. This is not supported currently. use_cuda_graph_full_step: bool = False @staticmethod def load_from_robot_config( robot_cfg: Union[Union[str, dict], RobotConfig], - world_cfg: Union[Union[str, dict], WorldConfig], + world_model: Union[Union[str, dict], WorldConfig], base_cfg: Optional[dict] = None, tensor_args: TensorDeviceType = TensorDeviceType(), compute_metrics: bool = True, use_cuda_graph: Optional[bool] = None, particle_opt_iters: Optional[int] = None, self_collision_check: bool = True, - collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE, + collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH, use_es: Optional[bool] = None, es_learning_rate: Optional[float] = 0.01, use_cuda_graph_metrics: bool = False, @@ -74,9 +108,56 @@ class MpcSolverConfig: use_lbfgs: bool = False, use_mppi: bool = True, ): + """Create an MPC solver configuration from robot and world configuration. + + Args: + robot_cfg: Robot configuration. Can be a path to a YAML file or a dictionary or + an instance of :class:`~curobo.types.robot.RobotConfig`. + world_model: World configuration. Can be a path to a YAML file or a dictionary or + an instance of :class:`~curobo.geom.types.WorldConfig`. + base_cfg: Base configuration for the solver. This file is used to check constraints + and convergence. If None, the default configuration from ``base_cfg.yml`` is used. + tensor_args: Numeric precision and device to run computations. + compute_metrics: Compute metrics on MPC step. + use_cuda_graph: Use CUDA graph for the optimization step. + particle_opt_iters: Number of iterations for the particle optimization. + self_collision_check: Enable self-collision check during MPC optimization. + collision_checker_type: Type of collision checker to use. See :ref:`world_collision`. + use_es: Use Evolution Strategies (ES) solver for MPC. Highly experimental. + es_learning_rate: Learning rate for ES solver. + use_cuda_graph_metrics: Use CUDA graph for computing metrics. + store_rollouts: Store rollouts information for debugging. This will also store the + trajectory taken by the end-effector across the horizon. + use_cuda_graph_full_step: Capture full step in MPC as a single CUDA graph. This is + experimental and might not work reliably. + sync_cuda_time: Synchronize CUDA device with host using + :py:func:`torch.cuda.synchronize` before calculating compute time. + collision_cache: Cache of obstacles to create to load obstacles between planning calls. + An example: ``{"obb": 10, "mesh": 10}``, to create a cache of 10 cuboids and 10 + meshes. + n_collision_envs: Number of collision environments to create for batched planning + across different environments. Only used for :py:meth:`MpcSolver.setup_solve_batch_env` + and :py:meth:`MpcSolver.setup_solve_batch_env_goalset`. + collision_activation_distance: Distance in meters to activate collision cost. A good + value to start with is 0.01 meters. Increase the distance if the robot needs to + stay further away from obstacles. + world_coll_checker: Instance of world collision checker to use for MPC. Leaving this to + None will create a new instance of world collision checker using the provided + :attr:`world_model`. + step_dt: Time step to use between each step in the trajectory. If None, the default + time step from the configuration~(`particle_mpc.yml` or `gradient_mpc.yml`) + is used. This dt should match the control frequency at which you are sending + commands to the robot. This dt should also be greater than than the compute + time for a single step. + use_lbfgs: Use L-BFGS solver for MPC. Highly experimental. + use_mppi: Use MPPI solver for MPC. + + Returns: + MpcSolverConfig: Configuration for the MPC solver. + """ + if use_cuda_graph_full_step: log_error("use_cuda_graph_full_step currently is not supported") - raise ValueError("use_cuda_graph_full_step currently is not supported") task_file = "particle_mpc.yml" config_data = load_yaml(join_path(get_task_configs_path(), task_file)) @@ -108,14 +189,14 @@ class MpcSolverConfig: if isinstance(robot_cfg, dict): robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args) - if isinstance(world_cfg, str): - world_cfg = load_yaml(join_path(get_world_configs_path(), world_cfg)) + if isinstance(world_model, str): + world_model = load_yaml(join_path(get_world_configs_path(), world_model)) - if world_coll_checker is None and world_cfg is not None: - world_cfg = WorldCollisionConfig.load_from_dict( - base_cfg["world_collision_checker_cfg"], world_cfg, tensor_args + if world_coll_checker is None and world_model is not None: + world_model = WorldCollisionConfig.load_from_dict( + base_cfg["world_collision_checker_cfg"], world_model, tensor_args ) - world_coll_checker = create_collision_checker(world_cfg) + world_coll_checker = create_collision_checker(world_model) grad_config_data = None if use_lbfgs: grad_config_data = load_yaml(join_path(get_task_configs_path(), "gradient_mpc.yml")) @@ -134,7 +215,7 @@ class MpcSolverConfig: base_cfg["constraint"], base_cfg["convergence"], base_cfg["world_collision_checker_cfg"], - world_cfg, + world_model, world_coll_checker=world_coll_checker, tensor_args=tensor_args, ) @@ -172,7 +253,7 @@ class MpcSolverConfig: base_cfg["constraint"], base_cfg["convergence"], base_cfg["world_collision_checker_cfg"], - world_cfg, + world_model, world_coll_checker=world_coll_checker, tensor_args=tensor_args, ) @@ -201,13 +282,42 @@ class MpcSolverConfig: class MpcSolver(MpcSolverConfig): - """Model Predictive Control Solver for Arm Reacher task. + """High-level interface for Model Predictive Control (MPC). - Args: - MpcSolverConfig: _description_ + MPC can reach Cartesian poses and joint configurations while avoiding obstacles. The solver + uses Model Predictive Path Integral (MPPI) optimization as the solver. MPC only optimizes + locally so the robot can get stuck near joint limits or behind obstacles. To generate global + trajectories, use :py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`. + + See :ref:`python_mpc_example` for an example. This MPC solver implementation can be used in the + following steps: + + 1. Create a :py:class:`~curobo.rollout.rollout_base.Goal` object with the target pose or joint + configuration. + 2. Create a goal buffer for the problem type using :meth:`setup_solve_single`, + :meth:`setup_solve_goalset`, :meth:`setup_solve_batch`, :meth:`setup_solve_batch_goalset`, + :meth:`setup_solve_batch_env`, or :meth:`setup_solve_batch_env_goalset`. Pass the goal + object from the previous step to this function. This function will update the internal + solve state of MPC and also the goal for MPC. An augmented goal buffer is returned. + 3. Call :meth:`step` with the current joint state to get the next action. + 4. To change the goal, create a :py:class:`~curobo.types.math.Pose` object with new pose or + :py:class:`~curobo.types.state.JointState` object with new joint configuration. Then + copy the target into the augmented goal buffer using + ``goal_buffer.goal_pose.copy_(new_pose)`` or ``goal_buffer.goal_state.copy_(new_state)``. + 5. Call :meth:`update_goal` with the augmented goal buffer to update the goal for MPC. + 6. Call :meth:`step` with the current joint state to get the next action. + + To dynamically change the type of goal reached between pose and joint configuration targets, + create the goal object in step 1 with both targets and then use :meth:`enable_cspace_cost` and + :meth:`enable_pose_cost` to enable or disable reaching joint configuration cost and pose cost. """ def __init__(self, config: MpcSolverConfig) -> None: + """Initializes the MPC solver. + + Args: + config: Configuration parameters for MPC. + """ super().__init__(**vars(config)) self.tensor_args = self.solver.rollout_fn.tensor_args self._goal_buffer = Goal() @@ -222,15 +332,326 @@ class MpcSolver(MpcSolverConfig): self._cu_step_graph = None self._cu_result = None - def _update_batch_size(self, batch_size): - if self.batch_size != batch_size: - self.batch_size = batch_size + def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: + """Creates a goal buffer to solve for a robot to reach target pose or joint configuration. - def update_goal_buffer( + Args: + goal: goal object containing target pose or joint configuration. + num_seeds: Number of seeds to use in the solver. + + Returns: + Goal: Instance of augmented goal buffer. + """ + solve_state = ReacherSolveState( + ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1 + ) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + + self.update_goal(goal_buffer) + + return goal_buffer + + def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: + """Creates a goal buffer to solve for a robot to reach a pose in a set of poses. + + Args: + goal: goal object containing target goalset or joint configuration. + num_seeds: Number of seeds to use in the solver. + + Returns: + Goal: Instance of augmented goal buffer. + """ + solve_state = ReacherSolveState( + ReacherSolveType.GOALSET, + num_mpc_seeds=num_seeds, + batch_size=1, + n_envs=1, + n_goalset=goal.n_goalset, + ) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + self.update_goal(goal_buffer) + return goal_buffer + + def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: + """Creates a goal buffer to solve for a batch of robots to reach targets. + + Args: + goal: goal object containing target poses or joint configurations. + num_seeds: Number of seeds to use in the solver. + + Returns: + Goal: Instance of augmented goal buffer. + """ + solve_state = ReacherSolveState( + ReacherSolveType.BATCH, + num_mpc_seeds=num_seeds, + batch_size=goal.batch, + n_envs=1, + n_goalset=1, + ) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + self.update_goal(goal_buffer) + + return goal_buffer + + def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: + """Creates a goal buffer to solve for a batch of robots to reach a set of poses. + + Args: + goal: goal object containing target goalset or joint configurations. + num_seeds: Number of seeds to use in the solver. + + Returns: + Goal: Instance of augmented goal buffer. + """ + solve_state = ReacherSolveState( + ReacherSolveType.BATCH_GOALSET, + num_mpc_seeds=num_seeds, + batch_size=goal.batch, + n_envs=1, + n_goalset=goal.n_goalset, + ) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + self.update_goal(goal_buffer) + + return goal_buffer + + def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: + """Creates a goal buffer to solve for a batch of robots in different collision worlds. + + Args: + goal: goal object containing target poses or joint configurations. + num_seeds: Number of seeds to use in the solver. + + Returns: + Goal: Instance of augmented goal buffer. + """ + solve_state = ReacherSolveState( + ReacherSolveType.BATCH_ENV, + num_mpc_seeds=num_seeds, + batch_size=goal.batch, + n_envs=1, + n_goalset=1, + ) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + self.update_goal(goal_buffer) + + return goal_buffer + + def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: + """Creates a goal buffer to solve for a batch of robots in different collision worlds. + + Args: + goal: goal object containing target goalset or joint configurations. + num_seeds: Number of seeds to use in the solver. + + Returns: + Goal: Instance of augmented goal buffer. + """ + solve_state = ReacherSolveState( + ReacherSolveType.BATCH_ENV_GOALSET, + num_mpc_seeds=num_seeds, + batch_size=goal.batch, + n_envs=1, + n_goalset=goal.n_goalset, + ) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + self.update_goal(goal_buffer) + + return goal_buffer + + def step( + self, + current_state: JointState, + shift_steps: int = 1, + seed_traj: Optional[JointState] = None, + max_attempts: int = 1, + ): + """Solve for the next action given the current state. + + Args: + current_state: Current joint state of the robot. + shift_steps: Number of steps to shift the trajectory. + seed_traj: Initial trajectory to seed the optimization. If None, the solver + uses the solution from the previous step. + max_attempts: Maximum number of attempts to solve the problem. + + Returns: + WrapResult: Result of the optimization. + """ + converged = True + + for _ in range(max_attempts): + result = self._step_once(current_state.clone(), shift_steps, seed_traj) + if ( + torch.count_nonzero(torch.isnan(result.action.position)) == 0 + and torch.count_nonzero(~result.metrics.feasible) == 0 + ): + converged = True + break + self.reset() + if not converged: + result.action.copy_(current_state) + log_warn("MPC didn't converge") + + return result + + def update_goal(self, goal: Goal): + """Update the goal for MPC. + + Args: + goal: goal object containing target pose or joint configuration. This goal instance + should be created using one of the setup_solve functions. + """ + self.solver.update_params(goal) + + def reset(self): + """Reset the solver.""" + # reset warm start + self.solver.reset() + + def reset_cuda_graph(self): + """Reset captured CUDA graph. This does not work.""" + self.solver.reset_cuda_graph() + + def enable_cspace_cost(self, enable=True): + """Enable or disable reaching joint configuration cost in the solver. + + Args: + enable: Enable or disable reaching joint configuration cost. When False, cspace cost + is disabled. + """ + self.solver.safety_rollout.enable_cspace_cost(enable) + for opt in self.solver.optimizers: + opt.rollout_fn.enable_cspace_cost(enable) + + def enable_pose_cost(self, enable=True): + """Enable or disable reaching pose cost in the solver. + + Args: + enable: Enable or disable reaching pose cost. When False, pose cost is disabled. + """ + self.solver.safety_rollout.enable_pose_cost(enable) + for opt in self.solver.optimizers: + opt.rollout_fn.enable_pose_cost(enable) + + def get_active_js( + self, + in_js: JointState, + ): + """Get controlled joints indexed in MPC order from the input joint state. + + Args: + in_js: Input joint state. + + Returns: + JointState: Joint state with controlled joints. + """ + + opt_jnames = self.rollout_fn.joint_names + opt_js = in_js.get_ordered_joint_state(opt_jnames) + return opt_js + + def update_world(self, world: WorldConfig): + """Update the collision world for the solver. + + This allows for updating the world representation as long as the new world representation + does not have a larger number of obstacles than the :attr:`MpcSolver.collision_cache` as + created during initialization of :class:`MpcSolverConfig`. + + Args: + world: New collision world configuration. See :ref:`world_collision` for more details. + """ + self.world_coll_checker.load_collision_model(world) + + def get_visual_rollouts(self): + """Get rollouts for debugging.""" + return self.solver.optimizers[0].get_rollouts() + + @property + def joint_names(self): + """Get the ordered joint names of the robot.""" + return self.rollout_fn.joint_names + + @property + def collision_cache(self) -> Dict[str, int]: + """Returns the collision cache created by the world collision checker.""" + return self.world_coll_checker.cache + + @property + def kinematics(self) -> CudaRobotModel: + """Get kinematics instance of the robot.""" + return self.solver.safety_rollout.dynamics_model.robot_model + + @property + def world_collision(self) -> WorldCollision: + """Get the world collision checker.""" + return self.world_coll_checker + + @property + def rollout_fn(self) -> ArmReacher: + """Get the rollout function.""" + return self.solver.safety_rollout + + def _step_once( + self, + current_state: JointState, + shift_steps: int = 1, + seed_traj: Optional[JointState] = None, + ) -> WrapResult: + """Solve for the next action given the current state. + + Args: + current_state: Current joint state of the robot. + shift_steps: Number of steps to shift the trajectory. + seed_traj: Initial trajectory to seed the optimization. If None, the solver + uses the solution from the previous step. + + Returns: + WrapResult: Result of the optimization. + """ + # Create cuda graph for whole solve step including computation of metrics + # Including updation of goal buffers + + if self._solve_state is None: + log_error("Need to first setup solve state before calling solve()") + + if self.use_cuda_graph_full_step: + st_time = time.time() + if not self._cu_step_init: + self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj) + self._cu_state_in.copy_(current_state) + if seed_traj is not None: + self._cu_seed.copy_(seed_traj) + self._cu_step_graph.replay() + result = self._cu_result.clone() + torch.cuda.synchronize(device=self.tensor_args.device) + result.solve_time = time.time() - st_time + else: + self._step_goal_buffer.current_state.copy_(current_state) + result = self._solve_from_solve_state( + self._solve_state, + self._step_goal_buffer, + shift_steps, + seed_traj, + ) + + return result + + def _update_solve_state_and_goal_buffer( self, solve_state: ReacherSolveState, goal: Goal, ) -> Goal: + """Update solve state and goal for MPC. + + Args: + solve_state: New solve state. + goal: New goal buffer. + + Returns: + Goal: Updated goal buffer. + """ self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal( goal, self._solve_state, @@ -250,71 +671,64 @@ class MpcSolver(MpcSolverConfig): ) return self._goal_buffer - def step( + def _update_batch_size(self, batch_size: int): + """Update the batch size of the solver. + + Args: + batch_size: Number of problems to solve in parallel. + """ + if self.batch_size != batch_size: + self.batch_size = batch_size + + def _solve_from_solve_state( self, - current_state: JointState, - shift_steps: int = 1, - seed_traj: Optional[JointState] = None, - max_attempts: int = 1, - ): - converged = True - - for _ in range(max_attempts): - result = self.step_once(current_state.clone(), shift_steps, seed_traj) - if ( - torch.count_nonzero(torch.isnan(result.action.position)) == 0 - and torch.max(torch.abs(result.action.position)) < 10.0 - and torch.count_nonzero(~result.metrics.feasible) == 0 - ): - converged = True - break - self.reset() - if not converged: - result.action.copy_(current_state) - log_warn("NOT CONVERGED") - - return result - - def step_once( - self, - current_state: JointState, + solve_state: ReacherSolveState, + goal: Goal, shift_steps: int = 1, seed_traj: Optional[JointState] = None, ) -> WrapResult: - # Create cuda graph for whole solve step including computation of metrics - # Including updation of goal buffers + """Solve for the next action given the current state. - if self._solve_state is None: - log_error("Need to first setup solve state before calling solve()") + Args: + solve_state: solve state object containing information about the current MPC problem. + goal: goal object containing target pose or joint configuration. + shift_steps: Number of steps to shift the trajectory before optimization. + seed_traj: Initial trajectory to seed the optimization. If None, the solver + uses the solution from the previous step. - if self.use_cuda_graph_full_step: - st_time = time.time() - if not self._cu_step_init: - self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj) - self._cu_state_in.copy_(current_state) - if seed_traj is not None: - self._cu_seed.copy_(seed_traj) - self._cu_step_graph.replay() - result = self._cu_result.clone() - torch.cuda.synchronize() - result.solve_time = time.time() - st_time - else: - self._step_goal_buffer.current_state.copy_(current_state) - result = self._solve_from_solve_state( - self._solve_state, - self._step_goal_buffer, - shift_steps, - seed_traj, - ) + Returns: + WrapResult: Result of the optimization. + """ + if solve_state.batch_env: + if solve_state.batch_size > self.world_coll_checker.n_envs: + log_error("Batch Env is less that goal batch") + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) + + if seed_traj is not None: + self.solver.update_init_seed(seed_traj) + + result = self.solver.solve(goal_buffer, seed_traj, shift_steps) + result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action) return result - def _step( + def _mpc_step( self, current_state: JointState, shift_steps: int = 1, seed_traj: Optional[JointState] = None, ): + """One step function that is used to create a CUDA graph. + + Args: + current_state: Current joint state of the robot. + shift_steps: Number of steps to shift the trajectory. + seed_traj: Initial trajectory to seed the optimization. If None, the solver + uses the solution from the previous step. + + Returns: + WrapResult: Result of the optimization. + """ self._step_goal_buffer.current_state.copy_(current_state) result = self._solve_from_solve_state( self._solve_state, @@ -331,6 +745,14 @@ class MpcSolver(MpcSolverConfig): shift_steps: int = 1, seed_traj: Optional[JointState] = None, ): + """Create a CUDA graph for the full step of MPC. + + Args: + current_state: Current joint state of the robot. + shift_steps: Number of steps to shift the trajectory. + seed_traj: Initial trajectory to seed the optimization. If None, the solver + uses the solution from the previous step. + """ log_info("MpcSolver: Creating Cuda Graph") self._cu_state_in = current_state.clone() if seed_traj is not None: @@ -339,7 +761,7 @@ class MpcSolver(MpcSolverConfig): s.wait_stream(torch.cuda.current_stream(device=self.tensor_args.device)) with torch.cuda.stream(s): for _ in range(3): - self._cu_result = self._step( + self._cu_result = self._mpc_step( self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed ) torch.cuda.current_stream(device=self.tensor_args.device).wait_stream(s) @@ -350,142 +772,3 @@ class MpcSolver(MpcSolverConfig): self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed ) self._cu_step_init = True - - def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: - solve_state = ReacherSolveState( - ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1 - ) - goal_buffer = self.update_goal_buffer(solve_state, goal) - return goal_buffer - - def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: - solve_state = ReacherSolveState( - ReacherSolveType.GOALSET, - num_mpc_seeds=num_seeds, - batch_size=1, - n_envs=1, - n_goalset=goal.n_goalset, - ) - goal_buffer = self.update_goal_buffer(solve_state, goal) - return goal_buffer - - def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: - solve_state = ReacherSolveState( - ReacherSolveType.BATCH, - num_mpc_seeds=num_seeds, - batch_size=goal.batch, - n_envs=1, - n_goalset=1, - ) - goal_buffer = self.update_goal_buffer(solve_state, goal) - return goal_buffer - - def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: - solve_state = ReacherSolveState( - ReacherSolveType.BATCH_GOALSET, - num_mpc_seeds=num_seeds, - batch_size=goal.batch, - n_envs=1, - n_goalset=goal.n_goalset, - ) - goal_buffer = self.update_goal_buffer(solve_state, goal) - return goal_buffer - - def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: - solve_state = ReacherSolveState( - ReacherSolveType.BATCH_ENV, - num_mpc_seeds=num_seeds, - batch_size=goal.batch, - n_envs=1, - n_goalset=1, - ) - goal_buffer = self.update_goal_buffer(solve_state, goal) - return goal_buffer - - def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal: - solve_state = ReacherSolveState( - ReacherSolveType.BATCH_ENV_GOALSET, - num_mpc_seeds=num_seeds, - batch_size=goal.batch, - n_envs=1, - n_goalset=goal.n_goalset, - ) - goal_buffer = self.update_goal_buffer(solve_state, goal) - return goal_buffer - - def _solve_from_solve_state( - self, - solve_state: ReacherSolveState, - goal: Goal, - shift_steps: int = 1, - seed_traj: Optional[JointState] = None, - ) -> WrapResult: - if solve_state.batch_env: - if solve_state.batch_size > self.world_coll_checker.n_envs: - raise ValueError("Batch Env is less that goal batch") - - goal_buffer = self.update_goal_buffer(solve_state, goal) - # NOTE: implement initialization from seed set here: - if seed_traj is not None: - self.solver.update_init_seed(seed_traj) - - result = self.solver.solve(goal_buffer, seed_traj, shift_steps) - result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action) - return result - - def fn(self): - # this will run one step of optimization and get new command - pass - - def update_goal(self, goal: Goal): - self.solver.update_params(goal) - return True - - def reset(self): - # reset warm start - self.solver.reset() - pass - - def reset_cuda_graph(self): - self.solver.reset_cuda_graph() - - @property - def rollout_fn(self): - return self.solver.safety_rollout - - def enable_cspace_cost(self, enable=True): - self.solver.safety_rollout.enable_cspace_cost(enable) - for opt in self.solver.optimizers: - opt.rollout_fn.enable_cspace_cost(enable) - - def enable_pose_cost(self, enable=True): - self.solver.safety_rollout.enable_pose_cost(enable) - for opt in self.solver.optimizers: - opt.rollout_fn.enable_pose_cost(enable) - - def get_active_js( - self, - in_js: JointState, - ): - opt_jnames = self.rollout_fn.joint_names - opt_js = in_js.get_ordered_joint_state(opt_jnames) - return opt_js - - @property - def joint_names(self): - return self.rollout_fn.joint_names - - def update_world(self, world: WorldConfig): - self.world_coll_checker.load_collision_model(world) - return True - - def get_visual_rollouts(self): - return self.solver.optimizers[0].get_rollouts() - - @property - def kinematics(self): - return self.solver.safety_rollout.dynamics_model.robot_model - - @property - def world_collision(self): - return self.world_coll_checker diff --git a/src/curobo/wrap/reacher/trajopt.py b/src/curobo/wrap/reacher/trajopt.py index 8ce1dbf..6bcbd60 100644 --- a/src/curobo/wrap/reacher/trajopt.py +++ b/src/curobo/wrap/reacher/trajopt.py @@ -8,13 +8,29 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # +""" +Trajectory optimization module contains :meth:`TrajOptSolver` class which can optimize for +minimum-jerk trajectories by first running a particle-based solver (MPPI) and then refining with +a gradient-based solver (L-BFGS). The module also provides linear interpolation functions for +generating seeds from start joint configuration to goal joint configurations. To generate +trajectories for reaching Cartesian poses or joint configurations, use the higher-level wrapper +:py:class:`~curobo.wrap.reacher.motion_gen.MotionGen`. +Trajectory Optimization uses joint positions as optimization variables with cost terms for +avoiding world collisions, self-collisions, and joint limits. The joint velocities, accelerations, +and jerks are computed using five point stencil. A squared l2-norm cost term on joint accelerations +and jerks is used to encourage smooth trajectories. A cost term for the terminal state to reach +either a Cartesian pose or joint configuration is also used. Read :ref:`research_page` for +more details. +""" + +from __future__ import annotations # Standard Library import math import time from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Sequence, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union # Third Party import torch @@ -52,6 +68,8 @@ from curobo.wrap.wrap_base import WrapBase, WrapConfig, WrapResult @dataclass class TrajOptSolverConfig: + """Configuration parameters for TrajOptSolver.""" + robot_config: RobotConfig solver: WrapBase rollout_fn: ArmReacher @@ -98,6 +116,7 @@ class TrajOptSolverConfig: interpolation_type: InterpolateType = InterpolateType.LINEAR_CUDA, interpolation_steps: int = 10000, interpolation_dt: float = 0.01, + minimum_trajectory_dt: Optional[float] = None, use_cuda_graph: bool = True, self_collision_check: bool = False, self_collision_opt: bool = True, @@ -130,8 +149,156 @@ class TrajOptSolverConfig: optimize_dt: bool = True, project_pose_to_goal_frame: bool = True, ): - # NOTE: Don't have default optimize_dt, instead read from a configuration file. - # use default values, disable environment collision checking + """Load TrajOptSolver configuration from robot configuration. + + Args: + robot_cfg: Robot configuration to use for motion generation. This can be a path to a + yaml file, a dictionary, or an instance of :class:`RobotConfig`. See + :ref:`available_robot_list` for a list of available robots. You can also create a + a configuration file for your robot using :ref:`tut_robot_configuration`. + world_model: World configuration to use for motion generation. This can be a path to a + yaml file, a dictionary, or an instance of :class:`WorldConfig`. See + :ref:`world_collision` for more details. + tensor_args: Numerical precision and compute device to use for motion generation. + position_threshold: Position threshold between target position and reached position in + meters. 0.005 is a good value for most tasks (5mm). + rotation_threshold: Rotation threshold between target orientation and reached + orientation. The metric is q^T * q, where q is the quaternion difference between + target and reached orientation. The metric is not easy to interpret and a future + release will provide a more intuitive metric. For now, use 0.05 as a good value. + cspace_threshold: Joint space threshold in radians for revolute joints and meters for + linear joints between reached joint configuration and target joint configuration + used to measure success. Default of 0.05 has been found to be a good value for most + cases. + world_coll_checker: Instance of world collision checker to use for motion generation. + Leaving this to None will create a new instance of world collision checker using + the provided attr:`world_model`. + base_cfg_file: Base configuration file containing convergence and constraint criteria + to measure success. + particle_file: Optimizer configuration file to use for particle-based + optimization during trajectory optimization. + gradient_file: Optimizer configuration file to use for gradient-based + optimization during trajectory optimization. + trajopt_tsteps: Number of waypoints to use for trajectory optimization. Default of 32 + is found to be a good number for most cases. + interpolation_type: Interpolation type to use for generating dense waypoints from + optimized trajectory. Default of + :py:attr:`curobo.util.trajectory.InterpolateType.LINEAR_CUDA` is found to be a + good choice for most cases. Other suitable options for real robot execution are + :py:attr:`curobo.util.trajectory.InterpolateType.QUINTIC` and + :py:attr:`curobo.util.trajectory.InterpolateType.CUBIC`. + interpolation_steps: Buffer size to use for storing interpolated trajectory. Default of + 5000 is found to be a good number for most cases. + interpolation_dt: Time step in seconds to use for generating interpolated trajectory + from optimized trajectory. Change this if you want to generate a trajectory with + a fixed timestep between waypoints. + minimum_trajectory_dt: Minimum time step in seconds allowed for trajectory + optimization. + use_cuda_graph: Record compute ops as cuda graphs and replay recorded graphs where + implemented. This can speed up execution by upto 10x. Default of True is + recommended. Enabling this will prevent changing solve type or batch size + after the first call to the solver. + self_collision_check: Enable self collision checks for generated motions. Default of + True is recommended. Set this to False to debug planning failures. Setting this to + False will also set self_collision_opt to False. + self_collision_opt: Enable self collision cost during optimization (IK, TrajOpt). + Default of True is recommended. + grad_trajopt_iters: Number of L-BFGS iterations to run trajectory optimization. + num_seeds: Number of seeds to use for trajectory optimization per problem. + seed_ratio: Ratio of linear and bias seeds to use for trajectory optimization. + Linear seed will generate a linear interpolated trajectory from start state + to IK solutions. Bias seed will add a mid-waypoint through the retract + configuration. Default of 1.0 linear and 0.0 bias is recommended. This can be + changed to 0.5 linear and 0.5 bias, along with changing num_seeds to 2. + trajopt_particle_opt: Enable particle-based optimization during trajectory + optimization. Default of True is recommended as particle-based optimization moves + the interpolated seeds away from bad local minima. + collision_checker_type: Type of collision checker to use for motion generation. Default + of CollisionCheckerType.MESH supports world represented by Cuboids and Meshes. See + :ref:`world_collision` for more details. + traj_evaluator_config: Configuration for trajectory evaluator. Default of None will + create a new instance of TrajEvaluatorConfig. After trajectory optimization across + many seeds, the best trajectory is selected based on this configuration. This + evaluator also checks if the optimized dt is within + :py:attr:`curobo.wrap.reacher.evaluator.TrajEvaluatorConfig.max_dt`. This check is + needed to measure smoothness of the optimized trajectory as bad trajectories can + have very high dt to fit within velocity, acceleration, and jerk limits. + traj_evaluator: Instance of trajectory evaluator to use for trajectory optimization. + Default of None will create a new instance of TrajEvaluator. In case you want to + use a custom evaluator, you can create a child instance of TrajEvaluator and + pass it. + minimize_jerk: Minimize jerk as regularization during trajectory optimizaiton. + use_gradient_descent: Use gradient descent instead of L-BFGS for trajectory + optimization. Default of False is recommended. Set to True for debugging gradients + of new cost terms. + 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 optimization + across different environments. Only used for + :py:meth:`TrajOptSolver.solve_batch_env`and + :py:meth:`TrajOptSolver.solve_batch_env_goalset`. + n_collision_envs: Number of collision environments to create for batched planning + across different environments. Only used for :py:meth:`MotionGen.plan_batch_env` + and :py:meth:`MotionGen.plan_batch_env_goalset`. + use_es: Use Evolution Strategies for optimization. Default of None will use MPPI. + es_learning_rate: Learning rate to use for Evolution Strategies. + use_fixed_samples: Use fixed samples for MPPI. Setting to False will increase compute + time as new samples are generated for each iteration of MPPI. + aux_rollout: Rollout instance to use for auxiliary rollouts. + evaluate_interpolated_trajectory: Evaluate interpolated trajectory after optimization. + Default of True is recommended to ensure the optimized trajectory is not passing + through very thin obstacles. + fixed_iters: Use fixed number of iterations of L-BFGS for trajectory + optimization. Default of None will use the setting from the optimizer + configuration. In most cases, fixed iterations of solvers are run as current + solvers treat constraints as costs and there is no guarantee that the constraints + will be satisfied. Instead of checking constraints between iterations of a solver + and exiting, it's computationally cheaper to run a fixed number of iterations. In + addition, running fixed iterations of solvers is more robust to outlier problems. + store_debug: Store debugging information such as values of optimization + variables in TrajOpt result. Setting this to True will set :attr:`use_cuda_graph` + to False. + sync_cuda_time: Synchronize with host using :py:func:`torch.cuda.synchronize` before + measuring compute time. + 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. + trajopt_dt: Time step in seconds to use for trajectory optimization. A good value to + start with is 0.15 seconds. This value is used to compute velocity, acceleration, + and jerk values for waypoints through finite difference. + trim_steps: Trim waypoints from optimized trajectory. The optimized trajectory will + contain the start state at index 0 and have the last two waypoints be the same + as T-2 as trajectory optimization implicitly optimizes for zero acceleration and + velocity at the last waypoint. An example: ``[1,-2]`` will trim the first waypoint + and last 3 waypoints from the optimized trajectory. + store_debug_in_result: Store debugging information in MotionGenResult. This value is + set to True if either store_ik_debug or store_trajopt_debug is set to True. + smooth_weight: Override smooth weight for trajectory optimization. It's not recommended + to set this value for most cases. + state_finite_difference_mode: Finite difference mode to use for computing velocity, + acceleration, and jerk values. Default of None will use the setting from the + optimizer configuration file. The default finite difference method is a five + point stencil to compute the derivatives as this is accurate and provides + faster convergence compared to backward or central difference methods. + filter_robot_command: Filter generated trajectory to remove finite difference + artifacts. Default of True is recommended. + optimize_dt: Optimize dt during trajectory optimization. Default of True is + recommended to find time-optimal trajectories. Setting this to False will use the + provided :attr:`trajopt_dt` for trajectory optimization. Setting to False is + required when optimizing from a non-static start state. + project_pose_to_goal_frame: Project pose to goal frame when calculating distance + between reached and goal pose. Use this to constrain motion to specific axes + either in the global frame or the goal frame. + + Returns: + TrajOptSolverConfig: Trajectory optimization configuration. + """ + + if minimum_trajectory_dt is None: + minimum_trajectory_dt = interpolation_dt + elif minimum_trajectory_dt < interpolation_dt: + log_error("minimum_trajectory_dt cannot be lower than interpolation_dt") if isinstance(robot_cfg, str): robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"] @@ -291,7 +458,6 @@ class TrajOptSolverConfig: elif use_particle_opt: mppi_cfg = ParallelMPPIConfig(**config_dict) parallel_mppi = ParallelMPPI(mppi_cfg) - config_dict = LBFGSOptConfig.create_data_dict( grad_config_data["lbfgs"], arm_rollout_grad, tensor_args ) @@ -323,7 +489,7 @@ class TrajOptSolverConfig: trajopt = WrapBase(cfg) if traj_evaluator_config is None: traj_evaluator_config = TrajEvaluatorConfig.from_basic( - min_dt=interpolation_dt, + min_dt=minimum_trajectory_dt, dof=robot_cfg.kinematics.dof, tensor_args=tensor_args, ) @@ -358,7 +524,9 @@ class TrajOptSolverConfig: @dataclass -class TrajResult(Sequence): +class TrajOptResult(Sequence): + """Trajectory optimization result.""" + success: T_BValue_bool goal: Goal solution: JointState @@ -379,9 +547,15 @@ class TrajResult(Sequence): goalset_index: Optional[torch.Tensor] = None optimized_seeds: Optional[torch.Tensor] = None - def __getitem__(self, idx): - # position_error = rotation_error = cspace_error = path_buffer_last_tstep = None - # metrics = interpolated_solution = None + def __getitem__(self, idx: int) -> TrajOptResult: + """Get item at index. + + Args: + idx: Index of the item to get. + + Returns: + TrajOptResult: Trajectory optimization result at the given index. + """ d_list = [ self.interpolated_solution, @@ -394,7 +568,7 @@ class TrajResult(Sequence): ] idx_vals = list_idx_if_not_none(d_list, idx) - return TrajResult( + return TrajOptResult( goal=self.goal[idx], solution=self.solution[idx], success=self.success[idx], @@ -411,15 +585,38 @@ class TrajResult(Sequence): optimized_seeds=self.optimized_seeds, ) - def __len__(self): + def __len__(self) -> int: + """Get length of the TrajOptResult.""" return self.success.shape[0] +@dataclass +class TrajResult(TrajOptResult): + """Deprecated: Use TrajOptResult instead of TrajResult""" + + def __post_init__(self): + """post-init function for TrajResult""" + log_warn("Deprecated: Use TrajOptResult instead of TrajResult") + + class TrajOptSolver(TrajOptSolverConfig): + """Trajectory Optimization Solver class for generating minimum-jerk trajectories. + + Trajectory Optimization uses joint positions as optimization variables with cost terms for + avoiding world collisions, self-collisions, and joint limits. The joint velocities, accelerations, + and jerks are computed using five point stencil. A squared l2-norm cost term on joint accelerations + and jerks is used to encourage smooth trajectories. A cost term for the terminal state to reach + either a Cartesian pose or joint configuration is also used. Read :ref:`research_page` for + more details. + """ + def __init__(self, config: TrajOptSolverConfig) -> None: + """Initialize TrajOptSolver with configuration parameters. + + Args: + config: Configuration parameters for TrajOptSolver. + """ super().__init__(**vars(config)) - self.dof = self.rollout_fn.d_action - self.action_horizon = self.rollout_fn.action_horizon self.delta_vec = interpolate_kernel(2, self.action_horizon, self.tensor_args).unsqueeze(0) self.waypoint_delta_vec = interpolate_kernel( @@ -427,13 +624,11 @@ class TrajOptSolver(TrajOptSolverConfig): ).unsqueeze(0) assert self.action_horizon / 2 != 0.0 self.solver.update_nproblems(self.num_seeds) - self._max_joint_vel = ( - self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[1, :].reshape( - 1, 1, self.dof - ) - ) - 0.02 - self._max_joint_acc = self.rollout_fn.state_bounds.acceleration[1, :] - 0.02 - self._max_joint_jerk = self.rollout_fn.state_bounds.jerk[1, :] - 0.02 + self._max_joint_vel = self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[ + 1, : + ].reshape(1, 1, self.dof) + self._max_joint_acc = self.rollout_fn.state_bounds.acceleration[1, :] + self._max_joint_jerk = self.rollout_fn.state_bounds.jerk[1, :] self._num_seeds = -1 self._col = None if self.traj_evaluator is None: @@ -450,6 +645,13 @@ class TrajOptSolver(TrajOptSolverConfig): self._rollout_list = None def get_all_rollout_instances(self) -> List[RolloutBase]: + """Get all rollout instances in the solver. + + Useful to update parameters across all rollout instances. + + Returns: + List[RolloutBase]: List of all rollout instances. + """ if self._rollout_list is None: self._rollout_list = [ self.rollout_fn, @@ -458,32 +660,67 @@ class TrajOptSolver(TrajOptSolverConfig): return self._rollout_list def get_all_kinematics_instances(self) -> List[CudaRobotModel]: + """Get all kinematics instances used across components in motion generation. + + This is deprecated. Use :meth:`TrajOptSolver.kinematics` instead as TrajOptSolver now uses + a shared kinematics instance across all components. + + Returns: + List[CudaRobotModel]: Single kinematics instance, returned as a list for compatibility. + """ + log_warn( + "Deprecated: Use TrajOptSolver.kinematics instead as TrajOptSolver now uses a " + + "shared kinematics instance across all components." + ) if self._kin_list is None: self._kin_list = [ i.dynamics_model.robot_model for i in self.get_all_rollout_instances() ] return self._kin_list - def attach_object_to_robot( + def attach_spheres_to_robot( self, sphere_radius: float, sphere_tensor: Optional[torch.Tensor] = None, link_name: str = "attached_object", ) -> None: - for k in self.get_all_kinematics_instances(): - k.attach_object( - sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name - ) + """Attach spheres to robot for collision checking. - def detach_object_from_robot(self, link_name: str = "attached_object") -> None: - for k in self.get_all_kinematics_instances(): - k.detach_object(link_name) + To fit spheres to an obstacle, see + :py:meth:`~curobo.geom.types.Obstacle.get_bounding_spheres` - def update_goal_buffer( + Args: + sphere_radius: Radius of the spheres. Set to None if :attr:`sphere_tensor` is provided. + sphere_tensor: Sphere x, y, z, r tensor. + link_name: Name of the link to attach the spheres to. Note that this link should + already have pre-allocated spheres. + """ + self.kinematics.attach_object( + sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name + ) + + def detach_spheres_from_robot(self, link_name: str = "attached_object") -> None: + """Detach spheres from robot. + + Args: + link_name: Name of the link to detach the spheres from. + """ + self.kinematics.detach_object(link_name) + + def _update_solve_state_and_goal_buffer( self, solve_state: ReacherSolveState, goal: Goal, ): + """Update goal buffer and solve state of current trajectory optimization problem. + + Args: + solve_state: New solve state. + goal: New goal buffer. + + Returns: + Goal: Updated goal buffer with augmentations for new solve state. + """ self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal( goal, self._solve_state, @@ -514,7 +751,28 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: Optional[int] = None, seed_success: Optional[torch.Tensor] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem with any solve type. + + Args: + solve_type: Type of solve to perform. + goal: Goal to reach. + seed_traj: Seed trajectory to start optimization from. This should be + of shape [num_seeds, batch_size, action_horizon, dof]. If None, linearly + interpolated seeds from current state to goal state are used. If goal.goal_state + is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. This cannot be changed after the + first call to solve as CUDA graph re-creation is currently not supported. + seed_success: Success of seeds. This is used to filter out successful seeds from + :attr:`seed_traj`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number + of iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if solve_type == ReacherSolveType.SINGLE: return self.solve_single( goal, @@ -574,7 +832,7 @@ class TrajOptSolver(TrajOptSolverConfig): newton_iters=newton_iters, ) - def solve_from_solve_state( + def _solve_from_solve_state( self, solve_state: ReacherSolveState, goal: Goal, @@ -585,13 +843,39 @@ class TrajOptSolver(TrajOptSolverConfig): seed_success: Optional[torch.Tensor] = None, newton_iters: Optional[int] = None, ): + """Solve trajectory optimization problem with a given solve state. + + Args: + solve_state: Solve state for the optimization problem. + goal: Goal object containing target pose or joint configuration. + seed_traj: Seed trajectory to start optimization from. This should be of + shape [num_seeds, batch_size, action_horizon, dof]. If None, linearly + interpolated seeds from current state to goal state are used. If goal.goal_state + is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. This cannot be changed after the + first call to solve as CUDA graph re-creation is currently not supported. + seed_success: Success of seeds. This is used to filter out successful seeds from + :attr:`seed_traj`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ 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") + log_error("Batch Env is less that goal batch") if newton_iters is not None: self.solver.newton_optimizer.outer_iters = newton_iters self.solver.newton_optimizer.fixed_iters = True - goal_buffer = self.update_goal_buffer(solve_state, goal) + log_info("TrajOpt: solving with Pose batch:" + str(goal.batch)) + goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal) # if self.evaluate_interpolated_trajectory: self.interpolate_rollout.update_params(goal_buffer) # get seeds: @@ -628,7 +912,31 @@ class TrajOptSolver(TrajOptSolverConfig): return_all_solutions: bool = False, num_seeds: Optional[int] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem for a single goal. + + This will use multiple seeds internally and return the best solution. + + Args: + goal: Goal to reach. + seed_traj: Seed trajectory to start optimization from. This should be of shape + [:attr:`num_seeds`, 1, :attr:`TrajOptSolver.action_horizon`, + :attr:`TrajOptSolver.dof`]. If None, linearly interpolated seeds from current state + to goal state are used. If goal.goal_state is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. If None, the number of seeds + is set to :attr:`TrajOptSolver.num_seeds`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if num_seeds is None: num_seeds = self.num_seeds solve_state = ReacherSolveState( @@ -639,7 +947,7 @@ class TrajOptSolver(TrajOptSolverConfig): n_goalset=1, ) - return self.solve_from_solve_state( + return self._solve_from_solve_state( solve_state, goal, seed_traj, @@ -657,7 +965,29 @@ class TrajOptSolver(TrajOptSolverConfig): return_all_solutions: bool = False, num_seeds: Optional[int] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem that uses goalset to represent Pose target. + + Args: + goal: Goal to reach. + seed_traj: Seed trajectory to start optimization from. This should be of shape + [:attr:`num_seeds`, 1, :attr:`TrajOptSolver.action_horizon`, + :attr:`TrajOptSolver.dof`]. If None, linearly interpolated seeds from current state + to goal state are used. If goal.goal_state is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. If None, the number of seeds + is set to :attr:`TrajOptSolver.num_seeds`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if num_seeds is None: num_seeds = self.num_seeds solve_state = ReacherSolveState( @@ -667,7 +997,7 @@ class TrajOptSolver(TrajOptSolverConfig): n_envs=1, n_goalset=goal.n_goalset, ) - return self.solve_from_solve_state( + return self._solve_from_solve_state( solve_state, goal, seed_traj, @@ -686,7 +1016,31 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: Optional[int] = None, seed_success: Optional[torch.Tensor] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem for a batch of goals. + + Args: + goal: Batch of goals to reach, this includes batch of current states. + seed_traj: Seed trajectory to start optimization from. This should be of shape + [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If + None, linearly interpolated seeds from current state to goal state are used. If + goal.goal_state is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. If None, the number of seeds + is set to :attr:`TrajOptSolver.num_seeds`. + seed_success: Success of seeds. This is used to filter out successful seeds from + :attr:`seed_traj`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if num_seeds is None: num_seeds = self.num_seeds solve_state = ReacherSolveState( @@ -696,7 +1050,7 @@ class TrajOptSolver(TrajOptSolverConfig): n_envs=1, n_goalset=1, ) - return self.solve_from_solve_state( + return self._solve_from_solve_state( solve_state, goal, seed_traj, @@ -716,7 +1070,31 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: Optional[int] = None, seed_success: Optional[torch.Tensor] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem for a batch of Poses with goalset. + + Args: + goal: Batch of goals to reach, this includes batch of current states. + seed_traj: Seed trajectory to start optimization from. This should be of shape + [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If + None, linearly interpolated seeds from current state to goal state are used. If + goal.goal_state is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. If None, the number of seeds + is set to :attr:`TrajOptSolver.num_seeds`. + seed_success: Success of seeds. This is used to filter out successful seeds from + :attr:`seed_traj`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if num_seeds is None: num_seeds = self.num_seeds solve_state = ReacherSolveState( @@ -726,7 +1104,7 @@ class TrajOptSolver(TrajOptSolverConfig): n_envs=1, n_goalset=goal.n_goalset, ) - return self.solve_from_solve_state( + return self._solve_from_solve_state( solve_state, goal, seed_traj, @@ -745,7 +1123,31 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: Optional[int] = None, seed_success: Optional[torch.Tensor] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem in a batch of environments. + + Args: + goal: Batch of goals to reach, this includes batch of current states. + seed_traj: Seed trajectory to start optimization from. This should be of shape + [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If + None, linearly interpolated seeds from current state to goal state are used. If + goal.goal_state is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. If None, the number of seeds + is set to :attr:`TrajOptSolver.num_seeds`. + seed_success: Success of seeds. This is used to filter out successful seeds from + :attr:`seed_traj`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if num_seeds is None: num_seeds = self.num_seeds solve_state = ReacherSolveState( @@ -755,7 +1157,7 @@ class TrajOptSolver(TrajOptSolverConfig): n_envs=goal.batch, n_goalset=1, ) - return self.solve_from_solve_state( + return self._solve_from_solve_state( solve_state, goal, seed_traj, @@ -775,7 +1177,31 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: Optional[int] = None, seed_success: Optional[torch.Tensor] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: + ) -> TrajOptResult: + """Solve trajectory optimization problem in a batch of environments with goalset. + + Args: + goal: Batch of goals to reach, this includes batch of current states. + seed_traj: Seed trajectory to start optimization from. This should be of shape + [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If + None, linearly interpolated seeds from current state to goal state are used. If + goal.goal_state is empty, random seeds are generated. + use_nn_seed: Use neural network seed for optimization. This is not implemented. + return_all_solutions: Return solutions for all seeds. + num_seeds: Number of seeds to use for optimization. If None, the number of seeds + is set to :attr:`TrajOptSolver.num_seeds`. + seed_success: Success of seeds. This is used to filter out successful seeds from + :attr:`seed_traj`. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This + is the outer iterations, where each outer iteration will run 25 inner iterations + of LBFGS optimization captured in a CUDA-Graph. Total number of optimization + iterations is 25 * outer_iters. The number of inner iterations can be changed + with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`. + + Returns: + TrajOptResult: Result of the trajectory optimization. + """ if num_seeds is None: num_seeds = self.num_seeds solve_state = ReacherSolveState( @@ -785,7 +1211,7 @@ class TrajOptSolver(TrajOptSolverConfig): n_envs=goal.batch, n_goalset=goal.n_goalset, ) - return self.solve_from_solve_state( + return self._solve_from_solve_state( solve_state, goal, seed_traj, @@ -803,21 +1229,11 @@ class TrajOptSolver(TrajOptSolverConfig): return_all_solutions: bool = False, num_seeds: Optional[int] = None, newton_iters: Optional[int] = None, - ) -> TrajResult: - """Only for single goal - - Args: - goal (Goal): _description_ - seed_traj (Optional[JointState], optional): _description_. Defaults to None. - use_nn_seed (bool, optional): _description_. Defaults to False. - - Raises: - NotImplementedError: _description_ - - Returns: - TrajResult: _description_ - """ - log_warn("TrajOpt.solve() is deprecated, use TrajOpt.solve_single or others instead") + ) -> TrajOptResult: + """Deprecated: Use :meth:`TrajOptSolver.solve_single` or others instead.""" + log_warn( + "TrajOptSolver.solve is deprecated, use TrajOptSolver.solve_single or others instead" + ) if goal.goal_pose.batch == 1 and goal.goal_pose.n_goalset == 1: return self.solve_single( goal, @@ -858,12 +1274,29 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: int, batch_mode: bool = False, ): + """Get result from the optimization problem. + + Args: + result: _description_ + return_all_solutions: _description_ + goal: _description_ + seed_traj: _description_ + num_seeds: _description_ + batch_mode: _description_ + + Raises: + log_error: _description_ + ValueError: _description_ + + Returns: + _description_ + """ st_time = time.time() if self.trim_steps is not None: result.action = result.action.trim_trajectory(self.trim_steps[0], self.trim_steps[1]) interpolated_trajs, last_tstep, opt_dt = self.get_interpolated_trajectory(result.action) if self.sync_cuda_time: - torch.cuda.synchronize() + torch.cuda.synchronize(device=self.tensor_args.device) interpolation_time = time.time() - st_time if self.evaluate_interpolated_trajectory: with profiler.record_function("trajopt/evaluate_interpolated"): @@ -876,7 +1309,6 @@ class TrajOptSolver(TrajOptSolverConfig): result.metrics.rotation_error = metrics.rotation_error result.metrics.cspace_error = metrics.cspace_error result.metrics.goalset_index = metrics.goalset_index - st_time = time.time() if result.metrics.cspace_error is None and result.metrics.position_error is None: raise log_error("convergence check requires either goal_pose or goal_state") @@ -890,6 +1322,7 @@ class TrajOptSolver(TrajOptSolverConfig): self.rotation_threshold, self.cspace_threshold, ) + if False: feasible = torch.all(result.metrics.feasible, dim=-1) @@ -905,7 +1338,7 @@ class TrajOptSolver(TrajOptSolverConfig): success = torch.logical_and(feasible, converge) if return_all_solutions: - traj_result = TrajResult( + traj_result = TrajOptResult( success=success, goal=goal, solution=result.action.scale_by_dt(self.solver_dt_tensor, opt_dt.view(-1, 1, 1)), @@ -976,7 +1409,7 @@ class TrajOptSolver(TrajOptSolverConfig): interpolated_traj = interpolated_trajs[idx] if self.sync_cuda_time: - torch.cuda.synchronize() + torch.cuda.synchronize(device=self.tensor_args.device) if len(best_act_seq.shape) == 3: opt_dt_v = opt_dt.view(-1, 1, 1) else: @@ -990,8 +1423,7 @@ class TrajOptSolver(TrajOptSolverConfig): "interpolation_time": interpolation_time, "select_time": select_time, } - - traj_result = TrajResult( + traj_result = TrajOptResult( success=success, goal=goal, solution=opt_solution, @@ -1022,21 +1454,11 @@ class TrajOptSolver(TrajOptSolverConfig): use_nn_seed: bool = False, return_all_solutions: bool = False, num_seeds: Optional[int] = None, - ) -> TrajResult: - """Only for single goal - - Args: - goal (Goal): _description_ - seed_traj (Optional[JointState], optional): _description_. Defaults to None. - use_nn_seed (bool, optional): _description_. Defaults to False. - - Raises: - NotImplementedError: _description_ - - Returns: - TrajResult: _description_ - """ - log_warn("TrajOpt.solve() is deprecated, use TrajOpt.solve_single or others instead") + ) -> TrajOptResult: + """Deprecated: Use :meth:`TrajOptSolver.solve_batch` or others instead.""" + log_warn( + "TrajOptSolver.batch_solve is deprecated, use TrajOptSolver.solve_batch or others instead" + ) if goal.n_goalset == 1: return self.solve_batch( goal, seed_traj, use_nn_seed, return_all_solutions, num_seeds, seed_success @@ -1046,7 +1468,16 @@ class TrajOptSolver(TrajOptSolverConfig): goal, seed_traj, use_nn_seed, return_all_solutions, num_seeds, seed_success ) - def get_linear_seed(self, start_state, goal_state): + def get_linear_seed(self, start_state, goal_state) -> torch.Tensor: + """Get linearly interpolated seeds from start states to goal states. + + Args: + start_state: start state of the robot. + goal_state: goal state of the robot. + + Returns: + torch.Tensor: Linearly interpolated seeds. + """ start_q = start_state.position.view(-1, 1, self.dof) end_q = goal_state.position.view(-1, 1, self.dof) edges = torch.cat((start_q, end_q), dim=1) @@ -1054,13 +1485,29 @@ class TrajOptSolver(TrajOptSolverConfig): seed = self.delta_vec @ edges return seed - def get_start_seed(self, start_state): + def get_start_seed(self, start_state) -> torch.Tensor: + """Get trajectory seeds with start state repeated. + + Args: + start_state: start state of the robot. + + Returns: + torch.Tensor: Trajectory seeds with start state repeated. + """ start_q = start_state.position.view(-1, 1, self.dof) edges = torch.cat((start_q, start_q), dim=1) seed = self.delta_vec @ edges return seed - def _get_seed_numbers(self, num_seeds): + def _get_seed_numbers(self, num_seeds: int) -> Dict[str, int]: + """Get number of seeds for each seed type. + + Args: + num_seeds: Total number of seeds to generate. + + Returns: + Dict[str, int]: Number of seeds for each seed type. + """ n_seeds = {"linear": 0, "bias": 0, "start": 0, "goal": 0} k = n_seeds.keys t_seeds = 0 @@ -1082,10 +1529,21 @@ class TrajOptSolver(TrajOptSolverConfig): num_seeds: Optional[int] = None, batch_mode: bool = False, ): - # if batch_mode: + """Get seed set for optimization. + + Args: + goal: Goal object containing target pose or joint configuration. + batch: _description_ + h: _description_ + seed_traj: _description_ + dofseed_success: _description_ + n_seedsnum_seeds: _description_ + batch_mode: _description_ + + Returns: + _description_ + """ total_seeds = goal.batch * num_seeds - # else: - # total_seeds = num_seeds if isinstance(seed_traj, JointState): seed_traj = seed_traj.position @@ -1093,7 +1551,6 @@ class TrajOptSolver(TrajOptSolverConfig): if goal.goal_state is not None and self.use_cspace_seed: # get linear seed seed_traj = self.get_seeds(goal.current_state, goal.goal_state, num_seeds=num_seeds) - # .view(batch_size, self.num_seeds, self.action_horizon, self.dof) else: # get start repeat seed: log_info("No goal state found, using current config to seed") @@ -1109,14 +1566,50 @@ class TrajOptSolver(TrajOptSolverConfig): new_seeds = self.get_seeds( goal.current_state, goal.goal_state, num_seeds - seed_traj.shape[0] ) - seed_traj = torch.cat((seed_traj, new_seeds), dim=0) # n_seed, batch, h, dof + seed_traj = torch.cat((seed_traj, new_seeds), dim=0) - seed_traj = seed_traj.view( - total_seeds, self.action_horizon, self.dof - ) # n_seeds,batch, h, dof + if len(seed_traj.shape) == 3: + if ( + seed_traj.shape[0] != total_seeds + or seed_traj.shape[1] != self.action_horizon + or seed_traj.shape[2] != self.dof + ): + log_error( + "Seed traj shape should be [num_seeds * batch, action_horizon, dof]" + + " current shape is " + + str(seed_traj.shape) + ) + elif len(seed_traj.shape) == 4: + if ( + seed_traj.shape[0] * seed_traj.shape[1] != total_seeds + or seed_traj.shape[2] != self.action_horizon + or seed_traj.shape[3] != self.dof + ): + log_error( + "Seed traj shape should be [num_seeds, batch, action_horizon, dof]" + + " or [1, num_seeds * batch, action_horizon, dof]" + + " current shape is " + + str(seed_traj.shape) + ) + else: + log_error("Seed traj shape should have 3 or 4 dimensions: " + str(seed_traj.shape)) + seed_traj = seed_traj.view(total_seeds, self.action_horizon, self.dof) return seed_traj - def get_seeds(self, start_state, goal_state, num_seeds=None): + def get_seeds( + self, start_state: JointState, goal_state: JointState, num_seeds: int = None + ) -> torch.Tensor: + """Get seed trajectories for optimization. + + Args: + start_state: Start state of the robot. + goal_state: Goal state of the robot. + num_seeds: Number of seeds to generate. If None, the number of seeds is set to + :attr:`TrajOptSolver.num_seeds`. + + Returns: + torch.Tensor: Seed trajectories of shape [num_seeds, batch, action_horizon, dof] + """ # repeat seeds: if num_seeds is None: num_seeds = self.num_seeds @@ -1152,24 +1645,45 @@ class TrajOptSolver(TrajOptSolverConfig): 1, n_seeds["goal"], 1, 1 ) seed_set.append(bias_seeds) - all_seeds = torch.cat( - seed_set, dim=1 - ) # .#transpose(0,1).contiguous() # n_seed, batch, h, dof + all_seeds = torch.cat(seed_set, dim=1) return all_seeds - def get_bias_seed(self, start_state, goal_state): + def get_bias_seed(self, start_state, goal_state) -> torch.Tensor: + """Get seed trajectories that pass through the retract configuration at mid-waypoint. + + Args: + start_state: start state of the robot. + goal_state: goal state of the robot. + + Returns: + torch.Tensor: Seed trajectories of shape [num_seeds * batch, action_horizon, dof]. + """ start_q = start_state.position.view(-1, 1, self.dof) end_q = goal_state.position.view(-1, 1, self.dof) bias_q = self.bias_node.view(-1, 1, self.dof).repeat(start_q.shape[0], 1, 1) edges = torch.cat((start_q, bias_q, end_q), dim=1) seed = self.waypoint_delta_vec @ edges - # print(seed) return seed @profiler.record_function("trajopt/interpolation") - def get_interpolated_trajectory(self, traj_state: JointState): + def get_interpolated_trajectory( + self, traj_state: JointState + ) -> Tuple[JointState, torch.Tensor]: + """Get interpolated trajectory from optimized trajectories. + + This function will first find the optimal dt for each trajectory in the batch by scaling + the trajectories to joint limits. Then it will interpolate the trajectory using the optimal + dt. + + Args: + traj_state: Optimized trajectories of shape [num_seeds * batch, action_horizon, dof]. + + Returns: + Tuple[JointState, torch.Tensor, torch.Tensor]: Interpolated trajectory, last time step + for each trajectory in batch, optimal dt for each trajectory in batch. + """ # do interpolation: if ( self._interpolated_traj_buffer is None @@ -1182,24 +1696,34 @@ class TrajOptSolver(TrajOptSolverConfig): self._interpolated_traj_buffer.joint_names = self.rollout_fn.joint_names state, last_tstep, opt_dt = get_batch_interpolated_trajectory( traj_state, + self.solver_dt_tensor, self.interpolation_dt, self._max_joint_vel, self._max_joint_acc, self._max_joint_jerk, - self.solver_dt_tensor, kind=self.interpolation_type, tensor_args=self.tensor_args, out_traj_state=self._interpolated_traj_buffer, min_dt=self.traj_evaluator_config.min_dt, + max_dt=self.traj_evaluator_config.max_dt, optimize_dt=self.optimize_dt, ) - return state, last_tstep, opt_dt def calculate_trajectory_dt( self, trajectory: JointState, + epsilon: float = 1e-6, ) -> torch.Tensor: + """Calculate the optimal dt for a given trajectory by scaling it to joint limits. + + Args: + trajectory: Trajectory to calculate optimal dt for. + epsilon: Small value to improve numerical stability. + + Returns: + torch.Tensor: Optimal trajectory dt. + """ opt_dt = calculate_dt_no_clamp( trajectory.velocity, trajectory.acceleration, @@ -1207,42 +1731,78 @@ class TrajOptSolver(TrajOptSolverConfig): self._max_joint_vel, self._max_joint_acc, self._max_joint_jerk, + epsilon=epsilon, ) return opt_dt def reset_seed(self): + """Reset the seed for random number generators in MPPI and rollout functions.""" self.solver.reset_seed() def reset_cuda_graph(self): + """Clear all recorded CUDA graphs. This does not work.""" self.solver.reset_cuda_graph() self.interpolate_rollout.reset_cuda_graph() self.rollout_fn.reset_cuda_graph() def reset_shape(self): + """Reset the shape of the rollout function and the solver.""" self.solver.reset_shape() self.interpolate_rollout.reset_shape() self.rollout_fn.reset_shape() @property def kinematics(self) -> CudaRobotModel: + """Get the kinematics instance of the robot.""" return self.rollout_fn.dynamics_model.robot_model @property - def retract_config(self): + def retract_config(self) -> torch.Tensor: + """Get the retract/home configuration of the robot. + + Returns: + torch.Tensor: Retract configuration of the robot. + """ return self.rollout_fn.dynamics_model.retract_config.view(1, -1) def fk(self, q: torch.Tensor) -> CudaRobotModelState: + """Compute forward kinematics for the robot. + + Args: + q: Joint configuration of the robot. + + Returns: + CudaRobotModelState: Forward kinematics of the robot. + """ return self.kinematics.get_state(q) @property - def solver_dt(self): + def solver_dt(self) -> torch.Tensor: + """Get the current trajectory dt for the solver. + + Returns: + torch.Tensor: Trajectory dt for the solver. + """ return self.solver.safety_rollout.dynamics_model.traj_dt[0] - # return self.solver.safety_rollout.dynamics_model.dt_traj_params.base_dt @property - def solver_dt_tensor(self): + def solver_dt_tensor(self) -> torch.Tensor: + """Get the current trajectory dt for the solver. + + Returns: + torch.Tensor: Trajectory dt for the solver. + """ return self.solver.safety_rollout.dynamics_model.traj_dt[0] + @property + def minimum_trajectory_dt(self) -> float: + """Get the minimum trajectory dt that is allowed, smaller dt will be clamped to this value. + + Returns: + float: Minimum trajectory dt that is allowed. + """ + return self.traj_evaluator.min_dt + def update_solver_dt( self, dt: Union[float, torch.Tensor], @@ -1250,21 +1810,70 @@ class TrajOptSolver(TrajOptSolverConfig): max_dt: Optional[float] = None, base_ratio: Optional[float] = None, ): + """Update the trajectory dt for the solver. + + This dt is used to calculate the velocity, acceleration and jerk of the trajectory through + five point stencil (finite difference). + + Args: + dt: New trajectory dt. + base_dt: Base dt for the trajectory. This is not supported. + max_dt: Maximum dt for the trajectory. This is not supported. + base_ratio: Ratio in trajectory length to scale from base_dt to max_dt. This is not + supported. + """ all_rollouts = self.get_all_rollout_instances() for rollout in all_rollouts: rollout.update_traj_dt(dt, base_dt, max_dt, base_ratio) - def compute_metrics(self, opt_trajectory: bool, interpolated_trajectory: bool): - self.solver.compute_metrics = opt_trajectory - self.evaluate_interpolated_trajectory = interpolated_trajectory - 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 + def update_pose_cost_metric( self, metric: PoseCostMetric, ): + """Update the pose cost metric for :ref:`tut_constrained_planning`. + + Only supports for the main end-effector. Does not support for multiple links that are + specified with `link_poses` in planning methods. + + Args: + metric: Type and parameters for pose constraint to add. + start_state: Start joint state for the constraint. + goal_pose: Goal pose for the constraint. + + Returns: + bool: True if the constraint can be added, False otherwise. + """ + rollouts = self.get_all_rollout_instances() [ rollout.update_pose_cost_metric(metric) @@ -1274,8 +1883,35 @@ class TrajOptSolver(TrajOptSolverConfig): @property def newton_iters(self): + """Get the number of newton outer iterations during L-BFGS optimization. + + Returns: + int: Number of newton outer iterations during L-BFGS optimization. + """ return self._og_newton_iters + @property + def dof(self) -> int: + """Get the number of degrees of freedom of the robot. + + Returns: + int: Number of degrees of freedom of the robot. + """ + return self.rollout_fn.d_action + + @property + def action_horizon(self) -> int: + """Get the action horizon of the trajectory optimization problem. + + Number of actions in trajectory optimization can be smaller than number of waypoints as + the first waypoint is the current state of the robot and the last two waypoints are + the same as T-2 waypoint to implicitly enforce zero acceleration and zero velocity at T. + + Returns: + int: Action horizon of the trajectory optimization problem. + """ + return self.rollout_fn.action_horizon + @get_torch_jit_decorator() def jit_feasible_success( @@ -1287,6 +1923,7 @@ def jit_feasible_success( rotation_threshold: float, cspace_threshold: float, ): + """JIT function to check if the optimization is successful.""" feasible = torch.all(feasible, dim=-1) converge = feasible if position_error is not None and rotation_error is not None: @@ -1318,6 +1955,7 @@ def jit_trajopt_best_select( col, opt_dt, ): + """JIT function to select the best solution from optimized seeds.""" success[~smooth_label] = False convergence_error = 0 # get the best solution: diff --git a/src/curobo/wrap/wrap_base.py b/src/curobo/wrap/wrap_base.py index 447d7e6..37ed51f 100644 --- a/src/curobo/wrap/wrap_base.py +++ b/src/curobo/wrap/wrap_base.py @@ -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 diff --git a/src/curobo/wrap/wrap_mpc.py b/src/curobo/wrap/wrap_mpc.py index a2f2229..51dbc71 100644 --- a/src/curobo/wrap/wrap_mpc.py +++ b/src/curobo/wrap/wrap_mpc.py @@ -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( diff --git a/tests/interpolation_test.py b/tests/interpolation_test.py index 11fe94d..5cbbf50 100644 --- a/tests/interpolation_test.py +++ b/tests/interpolation_test.py @@ -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, diff --git a/tests/motion_gen_cuda_graph_test.py b/tests/motion_gen_cuda_graph_test.py new file mode 100644 index 0000000..da0ed9e --- /dev/null +++ b/tests/motion_gen_cuda_graph_test.py @@ -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 diff --git a/tests/motion_gen_speed_test.py b/tests/motion_gen_speed_test.py index 9a208ef..773fff1 100644 --- a/tests/motion_gen_speed_test.py +++ b/tests/motion_gen_speed_test.py @@ -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) diff --git a/tests/trajopt_test.py b/tests/trajopt_test.py index 2f7f226..0666faf 100644 --- a/tests/trajopt_test.py +++ b/tests/trajopt_test.py @@ -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