Add re-timing, minimum dt robustness

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 
+- 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, 
+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 
+- 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. 
+default start state in examples be out of bounds.
-- MotionGen uses parallel_finetune by default. To get previous motion gen behavior, pass 
+- 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 
+- 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 
+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. 
+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 new trajopt config to allow for smooth solutions at slow speeds (`velocity_scale<=0.25`).
-add error when `velocity_scale<0.1`.
+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`. 
+- 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 
+- 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. 
+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 
+- On benchmarking Dataset, Planning time is now 42ms on average from 50ms. Higher quality solutions
-are also obtained. See [benchmarks](https://curobo.org/source/getting_started/4_benchmarks.html) for more details.
+are also obtained. See [benchmarks](https://curobo.org/source/getting_started/4_benchmarks.html)
-- Add ``WorldCollisionVoxel``, a new collision checking implementation that uses a voxel grid 
+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 
+- Add batched mode to robot image segmentation, supports single robot multiple camera and batch
-robot batch camera. 
+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 `<axis xyz="0 -1 0"/>` 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 
+- 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 
+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
+- Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and
-- Cleanup docker scripts. Use `build_docker.sh` instead of `build_dev_docker.sh`. Added isaac sim 
+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 `<axis xyz="0 1 0"/>`
  to `<axis xyz="0 -1 0"/>` 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`), 
+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 
+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 
+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()`.
 
 

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"]

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_trajopt_seeds = 4
-    og_collision_activation_distance = 0.03  # 0.03
+    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:
+            mg, robot_cfg, robot_world, world_nvblox = load_curobo(
-                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(
                 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()
+                world_nvblox.clear_cache()
-                mg.world_coll_checker.update_blox_hashes()
+                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
-
+                    torch.cuda.synchronize()
-                    mg.add_camera_frame(cam_obs, "world")
+                    st_int_time = time.time()
-
+                    world_nvblox.add_camera_frame(cam_obs, "world")
-                mg.process_camera_frames("world", False)
+                    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()
+                world_nvblox.update_blox_hashes()
-                torch.cuda.synchronize()
-                # if i > 2:
-                #    mg.world_coll_checker.clear_cache()
-                #    mg.world_coll_checker.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__":

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"]

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
                 )
 
 

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"))

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

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"
+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

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]

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]

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 =

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 
+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 
+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 
+- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
-  :meth:`types.state.JointState`, :meth:`types.camera.CameraObservation`, :meth:`types.math.Pose`.
+  :py:meth:`~types.state.JointState`, :py:meth:`~types.camera.CameraObservation`,
+  :py:meth:`~types.math.Pose`.
 """
 
 

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_link0": 0.1,
-        "panda_link1": 0.05, 
+        "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:

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
+          "radius": 0.011 

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",

src/curobo/content/configs/task/base_cfg.yml

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

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] 
+    weight: [10000.0, 500000.0, 500.0, 500.0]
-    smooth_weight: [0.0,10000.0, 5.0, 0.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 
+  n_iters: 400
-  inner_iters: 25 
+  inner_iters: 25
-  cold_start_n_iters: null 
+  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_convergence: 0.001
-  cost_delta_threshold: 1.0
+  cost_delta_threshold: 0.0001
-  cost_relative_threshold: 0.999 #0.999
+  cost_relative_threshold: 0.99
   epsilon: 0.01
-  history: 27 #15
+  history: 27
   use_cuda_graph: True
   n_problems: 1
   store_debug: False

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
+    run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
-    weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
+    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
+    run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
-    weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
+    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] 
+    weight: [10000.0, 500000.0, 500.0, 500.0] 
-    smooth_weight: [0.0,10000.0, 50.0, 0.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

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
+    weight: 5000.0
-    activation_distance: [0.1]
+    activation_distance: [0.001]
-    null_space_weight: [1.0]
+    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'
-  
+
 

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] 
+    weight: [50000.0, 50000.0, 500.0,500.0]
-    smooth_weight: [0.0,10000.0,5.0, 0.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 
+  inner_iters: 25
-  cold_start_n_iters: null 
+  cold_start_n_iters: null
   min_iters: 25
   line_search_scale: [0.1,0.3,0.7,1.0]
   fixed_iters: True

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'

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 (
+from curobo.cuda_robot_model.types import KinematicsTensorConfig, SelfCollisionKinematicsConfig
-    CudaRobotModelState,
-    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.
+    Currently dof is created only for links that we need to compute kinematics. E.g., for robots
-    E.g., for robots with many serial chains, add all links of the robot to get the correct dof.
+    with many serial chains, add all links of the robot to get the correct dof. This is not an
-    This is not an issue if you are loading collision spheres as that will cover the full geometry
+    issue if you are loading collision spheres as that will cover the full geometry of the robot.
-    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)
 

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

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);
 

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
 
 

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])

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,

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"]
         ):

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:

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.

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()

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
 

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

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,

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_l = p_b[d_id]
-    p_u = p_b[dof + d_id] - eta_p
+    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_l = v_b[d_id]
-    v_u = v_b[dof + d_id] - eta_v
+    v_u = v_b[dof + d_id]
-    a_l = a_b[d_id] + eta_a
+    v_range = v_u - v_l
-    a_u = a_b[dof + d_id] - eta_a
+    eta_v = eta_v * (v_range)
+    v_l += eta_v
+    v_u -= eta_v
 
-    j_l = j_b[d_id] + eta_j
+    a_l = a_b[d_id]
-    j_u = j_b[dof + d_id] - eta_j
+    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_l = p_b[d_id]
-    p_u = p_b[dof + d_id] - eta_p
+    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_l = v_b[d_id]
-    v_u = v_b[dof + d_id] - eta_v
+    v_u = v_b[dof + d_id]
-    a_l = a_b[d_id] + eta_a
+    v_range = v_u - v_l
-    a_u = a_b[dof + d_id] - eta_a
+    eta_v = eta_v * (v_range)
+    v_l += eta_v
+    v_u -= eta_v
 
-    j_l = j_b[d_id] + eta_j
+    a_l = a_b[d_id]
-    j_u = j_b[dof + d_id] - eta_j
+    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

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,
         )

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:

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(

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):
+def setup_curobo_logger(level="info"):
-    logger = logging.getLogger("curobo")
+    """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):
+def log_info(txt: str, logger_name: str = "curobo", *args, **kwargs):
-    logger = logging.getLogger("curobo")
+    """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):
+def log_error(
-    logger = logging.getLogger("curobo")
+    txt: str,
-    logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs)
+    logger_name: str = "curobo",
-    raise
+    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)

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

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_vel: Optional[torch.Tensor] = None,
-    max_acc: torch.Tensor,
+    max_acc: Optional[torch.Tensor] = None,
-    max_jerk: torch.Tensor,
+    max_jerk: Optional[torch.Tensor] = None,
-    raw_dt: float,
     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(
+    if optimize_dt:
-        raw_traj.velocity,
+        traj_steps, steps_max, opt_dt = calculate_tsteps(
-        raw_traj.acceleration,
+            raw_traj.velocity,
-        raw_traj.jerk,
+            raw_traj.acceleration,
-        interpolation_dt,
+            raw_traj.jerk,
-        max_vel,
+            interpolation_dt,
-        max_acc,
+            max_vel,
-        max_jerk,
+            max_acc,
-        raw_dt,
+            max_jerk,
-        min_dt,
+            raw_dt,
-        horizon,
+            min_dt,
-        optimize_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
+    if steps_max <= 0:
-    # To do linear interpolation, we
+        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
+    dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, max_dt)
-    # computed with raw_dt to a new 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)
+    traj_steps, steps_max = calculate_traj_steps(opt_dt, interpolation_dt, horizon)
-    steps_max = torch.max(traj_steps)
     return traj_steps, steps_max, opt_dt

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:

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.cuda_robot_model import CudaRobotModel, CudaRobotModelState
-from curobo.cuda_robot_model.types import 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

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
 

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 
+This module contains :meth:`IKSolver` to solve inverse kinematics, along with
-:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result.
+:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result. A minimal
+example to solve IK is available at :ref:`python_ik_example`.
 
 .. raw:: html
 
@@ -20,6 +21,7 @@ This module contains :meth:`IKSolver` to solve inverse kinematics, along with
     </p>
 This demo is available in :ref:`isaac_ik_reachability`.
 
+
 """
 from __future__ import annotations
 
@@ -190,9 +192,9 @@ class IKSolverConfig:
             grad_iters: Number of iterations for gradient optimization.
             use_particle_opt: Flag to enable particle optimization.
             collision_checker_type: Type of collision checker to use for collision checking.
-            sync_cuda_time: Flag to enable synchronization of cuda device with host before
+            sync_cuda_time: Flag to enable synchronization of cuda device with host using
-                measuring compute time. If you set this to False, then measured time will not
+                :py:func:`torch.cuda.synchronize` before measuring compute time. If you set this to
-                include completion of any launched CUDA kernels.
+                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."""

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
+
+    <p>
+    <video autoplay="True" loop="True" muted="True" preload="auto" width="100%"><source src="../videos/mpc_clip.mp4" type="video/mp4"></video>
+    </p>
+
+
+"""
 
 # Standard Library
 import time
@@ -19,6 +43,7 @@ from typing import Dict, Optional, Union
 import torch
 
 # CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
 from curobo.geom.sdf.utils import create_collision_checker
 from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig
 from curobo.geom.types import WorldConfig
@@ -44,22 +69,31 @@ from curobo.wrap.wrap_mpc import WrapConfig, WrapMpc
 
 @dataclass
 class MpcSolverConfig:
+    """Configuration dataclass for MPC."""
+
+    #: MPC Solver.
     solver: WrapMpc
+
+    #: World Collision Checker.
     world_coll_checker: Optional[WorldCollision] = None
+
+    #: Numeric precision and device to run computations.
     tensor_args: TensorDeviceType = TensorDeviceType()
+
+    #: Capture full step in MPC as a single CUDA graph. This is not supported currently.
     use_cuda_graph_full_step: bool = False
 
     @staticmethod
     def load_from_robot_config(
         robot_cfg: Union[Union[str, dict], RobotConfig],
-        world_cfg: Union[Union[str, dict], WorldConfig],
+        world_model: Union[Union[str, dict], WorldConfig],
         base_cfg: Optional[dict] = None,
         tensor_args: TensorDeviceType = TensorDeviceType(),
         compute_metrics: bool = True,
         use_cuda_graph: Optional[bool] = None,
         particle_opt_iters: Optional[int] = None,
         self_collision_check: bool = True,
-        collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE,
+        collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH,
         use_es: Optional[bool] = None,
         es_learning_rate: Optional[float] = 0.01,
         use_cuda_graph_metrics: bool = False,
@@ -74,9 +108,56 @@ class MpcSolverConfig:
         use_lbfgs: bool = False,
         use_mppi: bool = True,
     ):
+        """Create an MPC solver configuration from robot and world configuration.
+
+        Args:
+            robot_cfg: Robot configuration. Can be a path to a YAML file or a dictionary or
+                an instance of :class:`~curobo.types.robot.RobotConfig`.
+            world_model: World configuration. Can be a path to a YAML file or a dictionary or
+                an instance of :class:`~curobo.geom.types.WorldConfig`.
+            base_cfg: Base configuration for the solver. This file is used to check constraints
+                and convergence. If None, the default configuration from ``base_cfg.yml`` is used.
+            tensor_args: Numeric precision and device to run computations.
+            compute_metrics: Compute metrics on MPC step.
+            use_cuda_graph: Use CUDA graph for the optimization step.
+            particle_opt_iters: Number of iterations for the particle optimization.
+            self_collision_check: Enable self-collision check during MPC optimization.
+            collision_checker_type: Type of collision checker to use. See :ref:`world_collision`.
+            use_es: Use Evolution Strategies (ES) solver for MPC. Highly experimental.
+            es_learning_rate: Learning rate for ES solver.
+            use_cuda_graph_metrics: Use CUDA graph for computing metrics.
+            store_rollouts: Store rollouts information for debugging. This will also store the
+                trajectory taken by the end-effector across the horizon.
+            use_cuda_graph_full_step: Capture full step in MPC as a single CUDA graph. This is
+                experimental and might not work reliably.
+            sync_cuda_time: Synchronize CUDA device with host using
+                :py:func:`torch.cuda.synchronize` before calculating compute time.
+            collision_cache: Cache of obstacles to create to load obstacles between planning calls.
+                An example: ``{"obb": 10, "mesh": 10}``, to create a cache of 10 cuboids and 10
+                meshes.
+            n_collision_envs: Number of collision environments to create for batched planning
+                across different environments. Only used for :py:meth:`MpcSolver.setup_solve_batch_env`
+                and :py:meth:`MpcSolver.setup_solve_batch_env_goalset`.
+            collision_activation_distance: Distance in meters to activate collision cost. A good
+                value to start with is 0.01 meters. Increase the distance if the robot needs to
+                stay further away from obstacles.
+            world_coll_checker: Instance of world collision checker to use for MPC. Leaving this to
+                None will create a new instance of world collision checker using the provided
+                :attr:`world_model`.
+            step_dt: Time step to use between each step in the trajectory. If None, the default
+                time step from the configuration~(`particle_mpc.yml` or `gradient_mpc.yml`)
+                is used. This dt should match the control frequency at which you are sending
+                commands to the robot. This dt should also be greater than than the compute
+                time for a single step.
+            use_lbfgs: Use L-BFGS solver for MPC. Highly experimental.
+            use_mppi: Use MPPI solver for MPC.
+
+        Returns:
+            MpcSolverConfig: Configuration for the MPC solver.
+        """
+
         if use_cuda_graph_full_step:
             log_error("use_cuda_graph_full_step currently is not supported")
-            raise ValueError("use_cuda_graph_full_step currently is not supported")
 
         task_file = "particle_mpc.yml"
         config_data = load_yaml(join_path(get_task_configs_path(), task_file))
@@ -108,14 +189,14 @@ class MpcSolverConfig:
         if isinstance(robot_cfg, dict):
             robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
 
-        if isinstance(world_cfg, str):
+        if isinstance(world_model, str):
-            world_cfg = load_yaml(join_path(get_world_configs_path(), world_cfg))
+            world_model = load_yaml(join_path(get_world_configs_path(), world_model))
 
-        if world_coll_checker is None and world_cfg is not None:
+        if world_coll_checker is None and world_model is not None:
-            world_cfg = WorldCollisionConfig.load_from_dict(
+            world_model = WorldCollisionConfig.load_from_dict(
-                base_cfg["world_collision_checker_cfg"], world_cfg, tensor_args
+                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:
+    MPC can reach Cartesian poses and joint configurations while avoiding obstacles. The solver
-        MpcSolverConfig: _description_
+    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):
+    def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
-        if self.batch_size != batch_size:
+        """Creates a goal buffer to solve for a robot to reach target pose or joint configuration.
-            self.batch_size = batch_size
 
-    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,
+        solve_state: ReacherSolveState,
-        shift_steps: int = 1,
+        goal: Goal,
-        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,
         shift_steps: int = 1,
         seed_traj: Optional[JointState] = None,
     ) -> WrapResult:
-        # Create cuda graph for whole solve step including computation of metrics
+        """Solve for the next action given the current state.
-        # Including updation of goal buffers
 
-        if self._solve_state is None:
+        Args:
-            log_error("Need to first setup solve state before calling solve()")
+            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:
+        Returns:
-            st_time = time.time()
+            WrapResult: Result of the optimization.
-            if not self._cu_step_init:
+        """
-                self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
+        if solve_state.batch_env:
-            self._cu_state_in.copy_(current_state)
+            if solve_state.batch_size > self.world_coll_checker.n_envs:
-            if seed_traj is not None:
+                log_error("Batch Env is less that goal batch")
-                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,
-            )
 
+        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

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
 

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(

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,

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

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)

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
+    error = torch.linalg.norm((goal_state.position - traj[:, -1, :]), dim=-1)
-    assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.005
+    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