diff --git a/CHANGELOG.md b/CHANGELOG.md
index 01d122f..e1350cb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,27 +10,66 @@ its affiliates is strictly prohibited.
-->
# Changelog
+## Version 0.7.2
+
+### New Features
+- Significant improvements for generating slow trajectories. Added re-timing post processing to
+slow down optimized trajectories. Use `MotionGenPlanConfig.time_dilation_factor<1.0` to slow down a
+planned trajectory. This is more robust than setting `velocity_scale<1.0` and also allows for
+changing the speed of trajectories between planning calls
+- `curobo.util.logger` adds `logger_name` as an input, enabling use of logging api with other
+packages.
+
+### Changes in default behavior
+- Move `CudaRobotModelState` from `curobo.cuda_robot_model.types` to
+`curobo.cuda_robot_model.cuda_robot_model`
+- Activation distance for bound cost in now a ratio instead of absolute value to account for very
+small range of joint limits when `velocity_scale<0.1`.
+- `TrajResult` is renamed to `TrajOptResult` to be consistent with other solvers.
+- Order of inputs to `get_batch_interpolated_trajectory` has changed.
+- `MpcSolverConfig.load_from_robot_config` uses `world_model` instead of `world_cfg` to be
+consistent with other wrappers.
+
+### BugFixes & Misc.
+- Fix bug in `MotionGen.plan_batch_env` where graph planner was being set to True. This also fixes
+isaac sim example `batch_motion_gen_reacher.py`.
+- Add `min_dt` as a parameter to `MotionGenConfig` and `TrajOptSolverConfig` to improve readability
+and allow for having smaller `interpolation_dt`.
+- Add `epsilon` to `min_dt` to make sure after time scaling, joint temporal values are not exactly
+at their limits.
+- Remove 0.02 offset for `max_joint_vel` and `max_joint_acc` in `TrajOptSolver`
+- Bound cost now scales the cost by `1/limit_range**2` when `limit_range<1.0` to be robust to small
+joint limits.
+- Added documentation for `curobo.util.logger`, `curobo.wrap.reacher.motion_gen`,
+`curobo.wrap.reacher.mpc`, and `curobo.wrap.reacher.trajopt`.
+- When interpolation buffer is smaller than required, a new buffer is created with a warning
+instead of raising an exception.
+- `torch.cuda.synchronize()` now only synchronizes specified cuda device with
+`torch.cuda.synchronize(device=self.tensor_args.device)`
+- Added python example for MPC.
+
## Version 0.7.1
### New Features
- Add mimic joint parsing and optimization support. Check `ur5e_robotiq_2f_140.yml`.
-- Add `finetune_dt_scale` as a parameter to `MotionGenPlanConfig` to dynamically change the
+- Add `finetune_dt_scale` as a parameter to `MotionGenPlanConfig` to dynamically change the
time-optimal scaling on a per problem instance.
-- `MotionGen.plan_single()` will now try finetuning in a for-loop, with larger and larger dt
+- `MotionGen.plan_single()` will now try finetuning in a for-loop, with larger and larger dt
until convergence. This also warm starts from previous failure.
- Add `high_precision` mode to `MotionGenConfig` to support `<1mm` convergence.
### Changes in default behavior
-- collision_sphere_buffer now supports having offset per link. Also, collision_sphere_buffer only
-applies to world collision while self_collision_buffer applies for self collision. Previously,
+- collision_sphere_buffer now supports having offset per link. Also, collision_sphere_buffer only
+applies to world collision while self_collision_buffer applies for self collision. Previously,
self_collision_buffer was added on top of collision_sphere_buffer.
-- `TrajEvaluatorConfig` cannot be initialized without dof as now per-joint jerk and acceleration
+- `TrajEvaluatorConfig` cannot be initialized without dof as now per-joint jerk and acceleration
limits are used. Use `TrajEvaluatorConfig.from_basic()` to initialize similar to previous behavior.
- `finetune_dt_scale` default value is 0.9 from 0.95.
### BugFixes & Misc.
- Fix bug in `WorldVoxelCollision` where `env_query_idx` was being overwritten.
-- Fix bug in `WorldVoxelCollision` where parent collision types were not getting called in some cases.
+- Fix bug in `WorldVoxelCollision` where parent collision types were not getting called in some
+cases.
- Change voxelization dimensions to include 1 extra voxel per dim.
- Added `seed` parameter to `IKSolverConfig`.
- Added `sampler_seed` parameter `RolloutConfig`.
@@ -41,16 +80,16 @@ limits are used. Use `TrajEvaluatorConfig.from_basic()` to initialize similar to
- Reduced branching in Kinematics kernels and added mimic joint computations.
- Add init_cache to WorldVoxelCollision to create cache for Mesh and Cuboid obstacles.
- `TrajEvaluator` now uses per-joint acceleration and jerk limits.
-- Fixed regression in `batch_motion_gen_reacher.py` example where robot's position was not being
+- Fixed regression in `batch_motion_gen_reacher.py` example where robot's position was not being
set correctly.
- Switched from smooth l2 to l2 for BoundCost as that gives better convergence.
- `requires_grad` is explicitly stored in a varaible before `tensor.detach()` in warp kernel calls
as this can get set to False in some instances.
- Fix dt update in `MotionGen.plan_single_js()` where dt was not reset after finetunestep, causing
joint space planner to fail often.
-- Improve joint space planner success by changing smooth l2 distance cost to l2 distance. Also,
+- Improve joint space planner success by changing smooth l2 distance cost to l2 distance. Also,
added fallback to graph planner when linear path is not possible.
-- Retuned weigths for IKSolver, now 98th percentile accuracy is 10 micrometers wtih 16 seeds
+- Retuned weigths for IKSolver, now 98th percentile accuracy is 10 micrometers wtih 16 seeds
(vs 24 seeds previously).
- Switch float8 precision check from `const` to macro to avoid compile errors in older nvcc, this
fixes docker build issues for isaac sim 2023.1.0.
@@ -58,76 +97,77 @@ fixes docker build issues for isaac sim 2023.1.0.
## Version 0.7.0
### Changes in default behavior
- Increased default collision cache to 50 in RobotWorld.
-- Changed `CSpaceConfig.position_limit_clip` default to 0 as previous default of 0.01 can make
-default start state in examples be out of bounds.
-- MotionGen uses parallel_finetune by default. To get previous motion gen behavior, pass
+- Changed `CSpaceConfig.position_limit_clip` default to 0 as previous default of 0.01 can make
+default start state in examples be out of bounds.
+- MotionGen uses parallel_finetune by default. To get previous motion gen behavior, pass
`warmup(parallel_finetune=False)` and `MotionGenPlanConfig(parallel_finetune=False)`.
- MotionGen loads Mesh Collision checker instead of Primitive by default.
-- UR10e and UR5e now don't have a collision sphere at tool frame for world collision checking. This
+- UR10e and UR5e now don't have a collision sphere at tool frame for world collision checking. This
sphere is only active for self collision avoidance.
-- With torch>=2.0, cuRobo will use `torch.compile` instead of `torch.jit.script` to generate fused
-kernels. This can take several seconds during the first run. To enable this feature, set
+- With torch>=2.0, cuRobo will use `torch.compile` instead of `torch.jit.script` to generate fused
+kernels. This can take several seconds during the first run. To enable this feature, set
environment variable `export CUROBO_TORCH_COMPILE_DISABLE=0`.
### Breaking Changes
- Renamed `copy_if_not_none` to `clone_if_not_none` to be more descriptive. Now `copy_if_not_none`
will try to copy data into reference.
-- Renamed `n_envs` in curobo.opt module to avoid confusion between parallel environments and
+- Renamed `n_envs` in curobo.opt module to avoid confusion between parallel environments and
parallel problems in optimization.
- Added more inputs to pose distance kernels. Check `curobolib/geom.py`.
- Pose cost `run_vec_weight` should now be `[0,0,0,0,0,0]` instead of `[1,1,1,1,1,1]`
- ``max_distance`` is now tensor from ``float`` and is an input to collision kernels.
-- Order of inputs to ``SweptSdfMeshWarpPy`` has changed.
+- Order of inputs to ``SweptSdfMeshWarpPy`` has changed.
### New Features
-- Add function to disable and enable collision for specific links in KinematicsTensorConfig.
+- Add function to disable and enable collision for specific links in KinematicsTensorConfig.
- Add goal index to reacher results to return index of goal reached when goalset planning.
- Add locked joint state update api in MotionGen class.
- Add goalset warmup padding to handle varied number of goals during goalset planning and also when
-calling plan_single after warmup of goalset.
-- Add new trajopt config to allow for smooth solutions at slow speeds (`velocity_scale<=0.25`). Also
-add error when `velocity_scale<0.1`.
+calling plan_single after warmup of goalset.
+- Add new trajopt config to allow for smooth solutions at slow speeds (`velocity_scale<=0.25`).
+Also add error when `velocity_scale<0.1`.
- Add experimental robot image segmentation module to enable robot removal in depth images.
- Add constrained planning mode to motion_gen.
-- Use `torch.compile` to leverage better kernel fusion in place of `torch.jit.script`.
-- Significantly improved collision computation for cuboids and meshes. Mesh collision checker is
-now only 2x slower than cuboid (from 5x slower). Optimization convergence is also improved.
+- Use `torch.compile` to leverage better kernel fusion in place of `torch.jit.script`.
+- Significantly improved collision computation for cuboids and meshes. Mesh collision checker is
+now only 2x slower than cuboid (from 5x slower). Optimization convergence is also improved.
- LBFGS kernels now support ``history <= 31`` from ``history <= 15``.
-- 2x faster LBFGS kernel that allocates upto 68kb of shared memory, preventing use in CUDA devices
+- 2x faster LBFGS kernel that allocates upto 68kb of shared memory, preventing use in CUDA devices
with compute capability ``<7.0``.
-- On benchmarking Dataset, Planning time is now 42ms on average from 50ms. Higher quality solutions
-are also obtained. See [benchmarks](https://curobo.org/source/getting_started/4_benchmarks.html) for more details.
-- Add ``WorldCollisionVoxel``, a new collision checking implementation that uses a voxel grid
+- On benchmarking Dataset, Planning time is now 42ms on average from 50ms. Higher quality solutions
+are also obtained. See [benchmarks](https://curobo.org/source/getting_started/4_benchmarks.html)
+for more details.
+- Add ``WorldCollisionVoxel``, a new collision checking implementation that uses a voxel grid
of signed distances (SDF) to compute collision avoidance metrics. Documentation coming soon, see
``benchmark/curobo_voxel_benchmark.py`` for an example.
-- Add API for ESDF computation from world representations, see
+- Add API for ESDF computation from world representations, see
``WorldCollision.get_esdf_in_bounding_box()``.
-- Add partial support for isaac sim 2023.1.1. Most examples run for UR robots. `Franka Panda` is
+- Add partial support for isaac sim 2023.1.1. Most examples run for UR robots. `Franka Panda` is
unstable.
### BugFixes & Misc.
- refactored wp.index() instances to `[]` to avoid errors in future releases of warp.
- Fix bug in gaussian transformation to ensure values are not -1 or +1.
- Fix bug in ik_solver loading ee_link_name from argument.
-- Fix bug in batch_goalset planning, where pose cost was selected as GOALSET instead of
+- Fix bug in batch_goalset planning, where pose cost was selected as GOALSET instead of
BATCH_GOALSET.
- Added package data to also export `.so` files.
-- Fixed bug in transforming link visual mesh offset when reading from urdf.
+- Fixed bug in transforming link visual mesh offset when reading from urdf.
- Fixed bug in MotionGenPlanConfig.clone() that didn't clone the state of parallel_finetune.
-- Increased weighting from 1.0 to 10.0 for optimized_dt in TrajEvaluator to select shorter
+- Increased weighting from 1.0 to 10.0 for optimized_dt in TrajEvaluator to select shorter
trajectories.
- Improved determinism by setting global seed for random in `graph_nx.py`.
- Added option to clear obstacles in WorldPrimitiveCollision.
- Raise error when reference of tensors change in MotionGen, IKSolver, and TrajOpt when cuda graph
is enabled.
-- plan_single will get converted to plan_goalset when a plan_goalset was used to initialize cuda
+- plan_single will get converted to plan_goalset when a plan_goalset was used to initialize cuda
graph.
- plan_goalset will pad for extra goals when called with less number of goal than initial creation.
- Improved API documentation for Optimizer class.
- Set `use_cuda_graph` to `True` as default from `None` in `MotionGenConfig.load_from_robot_config`
-- Add batched mode to robot image segmentation, supports single robot multiple camera and batch
-robot batch camera.
+- Add batched mode to robot image segmentation, supports single robot multiple camera and batch
+robot batch camera.
- Add `log_warn` import to `arm_reacher.py`
- Remove negative radius check in self collision kernel to allow for self collision checking with
spheres of negative radius.
@@ -135,9 +175,9 @@ spheres of negative radius.
- Added UR5e robot with robotiq gripper (2f-140) with improved sphere model.
- Fix bug in aarch64.dockerfile where curobo was cloned to wrong path.
- Fix bug in aarch64.dockerfile where python was used instead of python3.
-- Remove unused variables in kernels.
+- Remove unused variables in kernels.
- Added ``pybind11`` as a dependency as some pytorch dockers for Jetson do not have this installed.
-- Fix incorrect dimensions in ``MotionGenResult.success`` in ``MotionGen.plan_batch()`` when
+- Fix incorrect dimensions in ``MotionGenResult.success`` in ``MotionGen.plan_batch()`` when
trajectory optimization fails.
- Added unit tests for collision checking functions.
- Fix bug in linear interpolation which was not reading the new ``optimized_dt`` to interpolate
@@ -147,52 +187,53 @@ between trajopt and finetune_trajopt.
### Known Bugs (WIP)
-- `Franka Panda` robot loading from urdf in isaac sim 2023.1.1 is unstable.
+- `Franka Panda` robot loading from urdf in isaac sim 2023.1.1 is unstable.
## Version 0.6.2
### New Features
- Added support for actuated axis to be negative (i.e., urdf joints with `` are
now natively supported).
-- Improved gradient calculation to account for terminal state. Trajectory optimization can reach
+- Improved gradient calculation to account for terminal state. Trajectory optimization can reach
within 1mm of accuracy (median across 2600 problems at 0.017mm).
-- Improved estimation of previous positions based on start velocity and acceleration. This enables
+- Improved estimation of previous positions based on start velocity and acceleration. This enables
Trajectory optimization to optimize from non-zero start velocity and accelerations.
- Added graph planner and finetuning step to joint space planning (motion_gen.plan_single_js). This
improves success and motion quality when planning to reach joint space targets.
- Added finetuning across many seeds in motion_gen, improving success rate and motion quality.
-- Add urdf support to usd helper to export optimization steps as animated usd files for debugging
+- Add urdf support to usd helper to export optimization steps as animated usd files for debugging
motion generation. Check `examples/usd_examples.py` for an example.
-- Retuned weights for IK and Trajectory optimization. This (+ other fixes) significantly improves
-pose reaching accuracy, IK accuracy improves by 100x (98th percentile < 10 micrometers) and motion
+- Retuned weights for IK and Trajectory optimization. This (+ other fixes) significantly improves
+pose reaching accuracy, IK accuracy improves by 100x (98th percentile < 10 micrometers) and motion
generation median at 0.017mm (with). IK now solves most problems with 24 seeds (vs 30 seeds prev.).
Run `benchmark/ik_benchmark.py` to get the latest results.
-- Added `external_asset_path` to robot configuration to help in loading urdf and meshes from an
+- Added `external_asset_path` to robot configuration to help in loading urdf and meshes from an
external directory.
### BugFixes & Misc.
- Update nvblox wrappers to work with v0.0.5 without segfaults. Significantly improves stability.
-- Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and 2022.2.1
-- Cleanup docker scripts. Use `build_docker.sh` instead of `build_dev_docker.sh`. Added isaac sim
+- Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and
+2022.2.1
+- Cleanup docker scripts. Use `build_docker.sh` instead of `build_dev_docker.sh`. Added isaac sim
development docker.
-- Fixed bug in backward kinematics kernel, helped improve IK and TO pose reaching accuracy..
+- Fixed bug in backward kinematics kernel, helped improve IK and TO pose reaching accuracy..
- Changed `panda_finger_joint2` from ``
to `` in `franka_panda.urdf` to match real robot urdf as cuRobo now supports
negative axis.
-- Changed benchmarking scripts to use lock joint state of [0.025,0.025] for mpinets dataset.
+- Changed benchmarking scripts to use lock joint state of [0.025,0.025] for mpinets dataset.
- Added scaling of mesh to Mesh.get_trimesh_mesh() to help in debugging mesh world.
- Improved stability and accuracy of MPPI for MPC.
-- Added NaN checking in STOMP covariance computation to account for cases when cholesky decomp
+- Added NaN checking in STOMP covariance computation to account for cases when cholesky decomp
fails.
- Added ground truth collision check validation in `benchmarks/curobo_nvblox_benchmark.py`.
### Performance Regressions
- cuRobo now generates significantly shorter paths then previous version. E.g., cuRobo obtains
-2.2 seconds 98th percentile motion time on the 2600 problems (`benchmark/curobo_benchmark.py`),
-where previously it was at 3 seconds (1.36x quicker motions). This was obtained by retuning the
-weights and slight reformulations of trajectory optimization. These changes have led to a slight
+2.2 seconds 98th percentile motion time on the 2600 problems (`benchmark/curobo_benchmark.py`),
+where previously it was at 3 seconds (1.36x quicker motions). This was obtained by retuning the
+weights and slight reformulations of trajectory optimization. These changes have led to a slight
degrade in planning time, 20ms slower on 4090 and 40ms on ORIN MAXN. We will address this slow down
-in a later release. One way to avoid this regression is to set `finetune_dt_scale=1.05` in
+in a later release. One way to avoid this regression is to set `finetune_dt_scale=1.05` in
`MotionGenConfig.load_from_robot_config()`.
diff --git a/benchmark/curobo_benchmark.py b/benchmark/curobo_benchmark.py
index 5c582b1..8499969 100644
--- a/benchmark/curobo_benchmark.py
+++ b/benchmark/curobo_benchmark.py
@@ -209,7 +209,6 @@ def load_curobo(
collision_activation_distance=collision_activation_distance,
trajopt_dt=0.25,
finetune_dt_scale=finetune_dt_scale,
- maximum_trajectory_dt=0.15,
high_precision=args.high_precision,
)
mg = MotionGen(motion_gen_config)
@@ -296,7 +295,6 @@ def benchmark_mb(
enable_graph_attempt=1,
disable_graph_attempt=10,
enable_finetune_trajopt=not args.no_finetune,
- partial_ik_opt=False,
enable_graph=graph_mode or force_graph,
timeout=60,
enable_opt=not graph_mode,
@@ -572,6 +570,7 @@ def benchmark_mb(
if not args.kpi:
try:
+ # Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]
@@ -604,6 +603,7 @@ def benchmark_mb(
g_m = CuroboGroupMetrics.from_list(all_files)
try:
+ # Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]
diff --git a/benchmark/curobo_nvblox_benchmark.py b/benchmark/curobo_nvblox_benchmark.py
index 6217ef5..b80264e 100644
--- a/benchmark/curobo_nvblox_benchmark.py
+++ b/benchmark/curobo_nvblox_benchmark.py
@@ -11,6 +11,7 @@
# Standard Library
import argparse
+import time
from copy import deepcopy
from typing import Optional
@@ -24,7 +25,8 @@ from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
from tqdm import tqdm
# CuRobo
-from curobo.geom.sdf.world import CollisionCheckerType, WorldConfig
+from curobo.geom.sdf.world import CollisionCheckerType, WorldCollisionConfig, WorldConfig
+from curobo.geom.sdf.world_blox import WorldBloxCollision
from curobo.geom.types import Cuboid as curobo_Cuboid
from curobo.geom.types import Mesh
from curobo.types.base import TensorDeviceType
@@ -35,6 +37,7 @@ from curobo.types.state import JointState
from curobo.util.logger import setup_curobo_logger
from curobo.util.metrics import CuroboGroupMetrics, CuroboMetrics
from curobo.util_file import (
+ file_exists,
get_assets_path,
get_robot_configs_path,
get_world_configs_path,
@@ -130,8 +133,10 @@ def load_curobo(
):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = 0.0
+ robot_cfg["kinematics"]["collision_link_names"].remove("attached_object")
+ robot_cfg["kinematics"]["ee_link"] = "panda_hand"
- ik_seeds = 30
+ ik_seeds = 32
if graph_mode:
trajopt_seeds = 4
if trajopt_seeds >= 14:
@@ -147,11 +152,30 @@ def load_curobo(
"world": {
"pose": [0, 0, 0, 1, 0, 0, 0],
"integrator_type": "tsdf",
- "voxel_size": 0.02,
+ "voxel_size": 0.01,
}
}
}
)
+ world_nvblox_config = WorldCollisionConfig.load_from_dict(
+ {"cache": None, "checker_type": "BLOX"},
+ world_cfg,
+ TensorDeviceType(),
+ )
+ world_nvblox = WorldBloxCollision(world_nvblox_config)
+ world_cfg = WorldConfig.from_dict(
+ {
+ "voxel": {
+ "base": {
+ "dims": [2.4, 2.4, 2.4],
+ "pose": [0, 0, 0, 1, 0, 0, 0],
+ "voxel_size": 0.005,
+ "feature_dtype": torch.float8_e4m3fn,
+ },
+ }
+ }
+ )
+
interpolation_steps = 2000
if graph_mode:
interpolation_steps = 100
@@ -164,7 +188,7 @@ def load_curobo(
robot_cfg_instance,
world_cfg,
trajopt_tsteps=tsteps,
- collision_checker_type=CollisionCheckerType.BLOX,
+ collision_checker_type=CollisionCheckerType.VOXEL,
use_cuda_graph=cuda_graph,
position_threshold=0.005, # 0.5 cm
rotation_threshold=0.05,
@@ -177,7 +201,6 @@ def load_curobo(
interpolation_steps=interpolation_steps,
collision_activation_distance=collision_activation_distance,
trajopt_dt=0.25,
- finetune_dt_scale=0.9,
maximum_trajectory_dt=0.15,
finetune_trajopt_iters=200,
)
@@ -188,12 +211,15 @@ def load_curobo(
robot_cfg_instance,
"collision_table.yml",
collision_activation_distance=0.0,
- collision_checker_type=CollisionCheckerType.PRIMITIVE,
+ collision_checker_type=CollisionCheckerType.MESH,
n_cuboids=50,
+ n_meshes=50,
)
+ mg.clear_world_cache()
robot_world = RobotWorld(config)
+ robot_world.clear_world_cache()
- return mg, robot_cfg, robot_world
+ return mg, robot_cfg, robot_world, world_nvblox
def benchmark_mb(
@@ -208,7 +234,7 @@ def benchmark_mb(
# load dataset:
graph_mode = args.graph
interpolation_dt = 0.02
- file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:2]
+ file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][2:]
enable_debug = save_log or plot_cost
all_files = []
@@ -216,8 +242,9 @@ def benchmark_mb(
if override_tsteps is not None:
og_tsteps = override_tsteps
- og_trajopt_seeds = 12
- og_collision_activation_distance = 0.03 # 0.03
+ og_trajopt_seeds = 4
+ og_collision_activation_distance = 0.025
+ count = [0, 0, 0, 0]
if args.graph:
og_trajopt_seeds = 4
for file_path in file_paths:
@@ -228,6 +255,7 @@ def benchmark_mb(
for key, v in tqdm(problems.items()):
scene_problems = problems[key]
m_list = []
+ count[3] += 1
i = -1
ik_fail = 0
trajopt_seeds = og_trajopt_seeds
@@ -236,32 +264,7 @@ def benchmark_mb(
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
- if "cage_panda" in key:
- trajopt_seeds = 8
- else:
- continue
- if "table_under_pick_panda" in key:
- tsteps = 44
- trajopt_seeds = 16
- finetune_dt_scale = 0.98
-
- if "cubby_task_oriented" in key and "merged" not in key:
- trajopt_seeds = 24
- finetune_dt_scale = 0.95
- collision_activation_distance = 0.035
- if "dresser_task_oriented" in key:
- trajopt_seeds = 24
- finetune_dt_scale = 0.95
- collision_activation_distance = 0.035 # 0.035
-
- if "merged_cubby_task_oriented" in key:
- collision_activation_distance = 0.025 # 0.035
- if "tabletop_task_oriented" in key:
- collision_activation_distance = 0.025 # 0.035
- if key in ["dresser_neutral_goal"]:
- trajopt_seeds = 24
- collision_activation_distance = og_collision_activation_distance
- mg, robot_cfg, robot_world = load_curobo(
+ mg, robot_cfg, robot_world, world_nvblox = load_curobo(
1,
enable_debug,
tsteps,
@@ -280,7 +283,6 @@ def benchmark_mb(
max_attempts=10,
enable_graph_attempt=1,
enable_finetune_trajopt=True,
- partial_ik_opt=False,
enable_graph=graph_mode,
timeout=60,
enable_opt=not graph_mode,
@@ -298,14 +300,60 @@ def benchmark_mb(
mg.reset(reset_seed=False)
world = WorldConfig.from_dict(problem["obstacles"])
- mg.world_coll_checker.clear_cache()
- mg.world_coll_checker.update_blox_hashes()
+ world_nvblox.clear_cache()
+ world_nvblox.update_blox_hashes()
+ mg.clear_world_cache()
- save_path = "benchmark/log/nvblox/" + key + "_" + str(i)
+ save_path = "benchmark/log/nvblox_640_new/" + key + "_" + str(i)
+ save_path = (
+ "/home/bala/code/raven_internship/data/render_mpinets_640/reformat/"
+ + key
+ + "_"
+ + str(i)
+ )
+ # save_path = "/home/bala/code/raven_internship/data/render_26k/_0_8/reformat/" + key + "_" + str(i)
+ if not file_exists(save_path):
+ continue
m_dataset = Sun3dDataset(save_path)
tensor_args = mg.tensor_args
+
+ if i == 0:
+ for j in tqdm(range(min(10, len(m_dataset))), leave=False):
+ data = m_dataset[j]
+ cam_obs = CameraObservation(
+ rgb_image=tensor_args.to_device(data["rgba"])
+ .squeeze()
+ .to(dtype=torch.uint8)
+ .permute(1, 2, 0), # data[rgba]: 4 x H x W -> H x W x 4
+ depth_image=tensor_args.to_device(data["depth"]),
+ intrinsics=data["intrinsics"],
+ pose=Pose.from_matrix(data["pose"].to(device=mg.tensor_args.device)),
+ )
+ cam_obs = cam_obs
+ torch.cuda.synchronize()
+ st_int_time = time.time()
+ world_nvblox.add_camera_frame(cam_obs, "world")
+ torch.cuda.synchronize()
+ world_nvblox.process_camera_frames("world", False)
+ torch.cuda.synchronize()
+ world_nvblox.update_blox_hashes()
+
+ # get esdf grid:
+ esdf = world_nvblox.get_esdf_in_bounding_box(
+ curobo_Cuboid(
+ name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2.4, 2.4, 2.4]
+ ),
+ voxel_size=0.005,
+ dtype=torch.float32,
+ )
+ mg.world_coll_checker.update_voxel_data(esdf)
+ world_nvblox.clear_cache()
+ world_nvblox.update_blox_hashes()
+ mg.clear_world_cache()
+
+ int_time = 0
for j in tqdm(range(min(1, len(m_dataset))), leave=False):
data = m_dataset[j]
cam_obs = CameraObservation(
@@ -318,18 +366,26 @@ def benchmark_mb(
pose=Pose.from_matrix(data["pose"].to(device=mg.tensor_args.device)),
)
cam_obs = cam_obs
-
- mg.add_camera_frame(cam_obs, "world")
-
- mg.process_camera_frames("world", False)
+ torch.cuda.synchronize()
+ st_int_time = time.time()
+ world_nvblox.add_camera_frame(cam_obs, "world")
+ torch.cuda.synchronize()
+ int_time += time.time() - st_int_time
+ st_time = time.time()
+ world_nvblox.process_camera_frames("world", False)
torch.cuda.synchronize()
- mg.world_coll_checker.update_blox_hashes()
- torch.cuda.synchronize()
- # if i > 2:
- # mg.world_coll_checker.clear_cache()
- # mg.world_coll_checker.update_blox_hashes()
+ world_nvblox.update_blox_hashes()
- # mg.world_coll_checker.save_layer("world", "test.nvblx")
+ # get esdf grid:
+ esdf = world_nvblox.get_esdf_in_bounding_box(
+ curobo_Cuboid(name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2.4, 2.4, 2.4]),
+ voxel_size=0.005,
+ dtype=torch.float32,
+ )
+ mg.world_coll_checker.update_voxel_data(esdf)
+
+ torch.cuda.synchronize()
+ perception_time = time.time() - st_time + int_time
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single(
@@ -372,6 +428,31 @@ def benchmark_mb(
),
log_scale=False,
)
+ if save_log or write_usd:
+ world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
+
+ # nvblox_obs = world_nvblox.get_mesh_from_blox_layer(
+ # "world",
+ # )
+ # nvblox_obs.vertex_colors = None
+
+ # if nvblox_obs.vertex_colors is not None:
+ # nvblox_obs.vertex_colors = nvblox_obs.vertex_colors.cpu().numpy()
+ # else:
+ # nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
+
+ # nvblox_obs.name = "nvblox_mesh_world"
+ # world.add_obstacle(nvblox_obs)
+
+ coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
+ curobo_Cuboid(name="new_test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2, 2, 2]),
+ voxel_size=0.01,
+ )
+
+ coll_mesh.color = [0.0, 0.8, 0.8, 0.8]
+
+ coll_mesh.name = "nvblox_voxel_world"
+ world.add_obstacle(coll_mesh)
if result.success.item():
q_traj = result.get_interpolated_plan()
problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist()
@@ -425,10 +506,10 @@ def benchmark_mb(
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
- world_coll = WorldConfig.from_dict(problem["obstacles"]).get_obb_world()
+ world_coll = WorldConfig.from_dict(problem["obstacles"]).get_mesh_world()
# check if path is collision free w.r.t. ground truth mesh:
- # robot_world.world_model.clear_cache()
+ robot_world.world_model.clear_cache()
robot_world.update_world(world_coll)
q_int_traj = result.get_interpolated_plan().position.unsqueeze(0)
@@ -448,6 +529,9 @@ def benchmark_mb(
# if not d_mask:
# write_usd = True
# #print(torch.max(d_world).item(), problem_name)
+ if d_int_mask:
+ count[0] += 1
+
current_metrics = CuroboMetrics(
skip=False,
success=True,
@@ -465,37 +549,11 @@ def benchmark_mb(
motion_time=result.motion_time.item(),
solve_time=result.solve_time,
jerk=torch.max(torch.abs(result.optimized_plan.jerk)).item(),
+ perception_time=perception_time,
)
- if save_log or write_usd:
- world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
-
- nvblox_obs = mg.world_coll_checker.get_mesh_from_blox_layer(
- "world",
- )
- # nvblox_obs.vertex_colors = None
-
- if nvblox_obs.vertex_colors is not None:
- nvblox_obs.vertex_colors = nvblox_obs.vertex_colors.cpu().numpy()
- else:
- nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
-
- nvblox_obs.name = "nvblox_mesh_world"
- world.add_obstacle(nvblox_obs)
-
- coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
- curobo_Cuboid(
- name="test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.5, 1.5, 1]
- ),
- voxel_size=0.005,
- )
-
- coll_mesh.color = [0.0, 0.8, 0.8, 0.8]
-
- coll_mesh.name = "nvblox_voxel_world"
- world.add_obstacle(coll_mesh)
# run planner
- if write_usd:
+ if write_usd and not d_mask:
# CuRobo
from curobo.util.usd_helper import UsdHelper
@@ -512,7 +570,7 @@ def benchmark_mb(
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
- # flatten_usd=True,
+ flatten_usd=True,
)
# write_usd = False
# exit()
@@ -537,7 +595,9 @@ def benchmark_mb(
m_list.append(current_metrics)
all_groups.append(current_metrics)
+ count[1] += 1
elif result.valid_query:
+ count[1] += 1
current_metrics = CuroboMetrics()
debug = {
"used_graph": result.used_graph,
@@ -555,6 +615,7 @@ def benchmark_mb(
m_list.append(current_metrics)
all_groups.append(current_metrics)
else:
+ count[2] += 1
# print("invalid: " + problem_name)
debug = {
"used_graph": result.used_graph,
@@ -583,7 +644,7 @@ def benchmark_mb(
write_trajopt=True,
visualize_robot_spheres=True,
grid_space=2,
- # flatten_usd=True,
+ flatten_usd=True,
)
exit()
g_m = CuroboGroupMetrics.from_list(m_list)
@@ -599,6 +660,7 @@ def benchmark_mb(
g_m.cspace_path_length.percent_98,
g_m.motion_time.percent_98,
g_m.perception_interpolated_success,
+ g_m.perception_time.mean,
# g_m.orientation_error.median,
)
print(g_m.attempts)
@@ -633,6 +695,7 @@ def benchmark_mb(
print("ST: ", g_m.solve_time)
print("accuracy: ", g_m.position_error, g_m.orientation_error)
print("Jerk: ", g_m.jerk)
+ print(count)
if __name__ == "__main__":
diff --git a/benchmark/curobo_voxel_benchmark.py b/benchmark/curobo_voxel_benchmark.py
index 08171c3..191d500 100644
--- a/benchmark/curobo_voxel_benchmark.py
+++ b/benchmark/curobo_voxel_benchmark.py
@@ -185,7 +185,7 @@ def load_curobo(
finetune_trajopt_iters=200,
)
mg = MotionGen(motion_gen_config)
- mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=True)
+ mg.warmup(enable_graph=True, warmup_js_trajopt=False)
# create a ground truth collision checker:
world_model = WorldConfig.from_dict(
{
@@ -305,6 +305,7 @@ def benchmark_mb(
voxel_size=0.005,
dtype=torch.float32,
)
+ # esdf.feature_tensor[esdf.feature_tensor < -1.0] = -1000.0
world_voxel_collision = mg.world_coll_checker
world_voxel_collision.update_voxel_data(esdf)
@@ -578,6 +579,7 @@ def benchmark_mb(
print(g_m.attempts)
g_m = CuroboGroupMetrics.from_list(all_groups)
try:
+ # Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]
@@ -611,6 +613,7 @@ def benchmark_mb(
g_m = CuroboGroupMetrics.from_list(all_files)
print("######## FULL SET ############")
try:
+ # Third Party
from tabulate import tabulate
headers = ["Metric", "Value"]
diff --git a/benchmark/generate_nvblox_images.py b/benchmark/generate_nvblox_images.py
index 48681c0..994aae2 100644
--- a/benchmark/generate_nvblox_images.py
+++ b/benchmark/generate_nvblox_images.py
@@ -36,7 +36,7 @@ np.random.seed(0)
def generate_images():
# load dataset:
- file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:]
+ file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:2]
for file_path in file_paths:
problems = file_path()
@@ -57,7 +57,7 @@ def generate_images():
# generate images and write to disk:
MeshDataset(
- None, n_frames=50, image_size=640, save_data_dir=save_path, trimesh_mesh=mesh
+ None, n_frames=1, image_size=640, save_data_dir=save_path, trimesh_mesh=mesh
)
diff --git a/benchmark/ik_benchmark.py b/benchmark/ik_benchmark.py
index 83216b8..6aaf40c 100644
--- a/benchmark/ik_benchmark.py
+++ b/benchmark/ik_benchmark.py
@@ -185,6 +185,7 @@ if __name__ == "__main__":
print("Reported errors are 98th percentile")
df.to_csv(join_path(args.save_path, args.file_name + ".csv"))
try:
+ # Third Party
from tabulate import tabulate
print(tabulate(df, headers="keys", tablefmt="grid"))
diff --git a/docker/user.dockerfile b/docker/user.dockerfile
index 1ad73d4..0b385e1 100644
--- a/docker/user.dockerfile
+++ b/docker/user.dockerfile
@@ -10,12 +10,12 @@
##
# Check architecture and load:
-ARG IMAGE_TAG
+ARG IMAGE_TAG
FROM curobo_docker:${IMAGE_TAG}
# Set variables
ARG USERNAME
ARG USER_ID
-ARG CACHE_DATE=2023-04-11
+ARG CACHE_DATE=2024-04-25
# Set environment variables
@@ -29,7 +29,7 @@ RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
# Set user
-# RUN mkdir /home/$USERNAME && chown $USERNAME:$USERNAME /home/$USERNAME
+# RUN mkdir /home/$USERNAME && chown $USERNAME:$USERNAME /home/$USERNAME
USER $USERNAME
WORKDIR /home/$USERNAME
ENV USER=$USERNAME
diff --git a/docker/x86.dockerfile b/docker/x86.dockerfile
index b853afb..0bf970a 100644
--- a/docker/x86.dockerfile
+++ b/docker/x86.dockerfile
@@ -78,13 +78,13 @@ ENV TORCH_CUDA_ARCH_LIST "7.0+PTX"
ENV LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}"
# Add cache date to avoid using cached layers older than this
-ARG CACHE_DATE=2024-04-11
+ARG CACHE_DATE=2024-04-25
RUN pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"
# if you want to use a different version of curobo, create folder as docker/pkgs and put your
-# version of curobo there. Then uncomment below line and comment the next line that clones from
+# version of curobo there. Then uncomment below line and comment the next line that clones from
# github
# COPY pkgs /pkgs
@@ -124,4 +124,7 @@ RUN cd /pkgs && git clone https://github.com/nvlabs/nvblox_torch.git && \
RUN python -m pip install pyrealsense2 opencv-python transforms3d
# install benchmarks:
-RUN python -m pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"
\ No newline at end of file
+RUN python -m pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"
+
+# update ucx path: https://github.com/openucx/ucc/issues/476
+RUN export LD_LIBRARY_PATH=/opt/hpcx/ucx/lib:$LD_LIBRARY_PATH
\ No newline at end of file
diff --git a/examples/isaac_sim/batch_motion_gen_reacher.py b/examples/isaac_sim/batch_motion_gen_reacher.py
index 320bf19..9785dee 100644
--- a/examples/isaac_sim/batch_motion_gen_reacher.py
+++ b/examples/isaac_sim/batch_motion_gen_reacher.py
@@ -144,17 +144,12 @@ def main():
robot_cfg,
world_cfg_list,
tensor_args,
- trajopt_tsteps=32,
collision_checker_type=CollisionCheckerType.MESH,
use_cuda_graph=True,
- num_trajopt_seeds=12,
- num_graph_seeds=12,
interpolation_dt=0.03,
- collision_cache={"obb": 6, "mesh": 6},
+ collision_cache={"obb": 10, "mesh": 10},
collision_activation_distance=0.025,
- self_collision_check=True,
maximum_trajectory_dt=0.25,
- fixed_iters_trajopt=True,
)
motion_gen = MotionGen(motion_gen_config)
j_names = robot_cfg["kinematics"]["cspace"]["joint_names"]
@@ -162,7 +157,9 @@ def main():
print("warming up...")
motion_gen.warmup(
- enable_graph=False, warmup_js_trajopt=False, batch=n_envs, batch_env_mode=True
+ batch=n_envs,
+ batch_env_mode=True,
+ warmup_js_trajopt=False,
)
add_extensions(simulation_app, args.headless_mode)
@@ -176,7 +173,7 @@ def main():
x_sph[..., 3] = radius
env_query_idx = torch.arange(n_envs, device=tensor_args.device, dtype=torch.int32)
plan_config = MotionGenPlanConfig(
- enable_graph=False, enable_graph_attempt=4, max_attempts=2, enable_finetune_trajopt=True
+ enable_graph=False, max_attempts=2, enable_finetune_trajopt=True
)
prev_goal = None
cmd_plan = [None, None]
diff --git a/examples/motion_gen_example.py b/examples/motion_gen_example.py
index ace6162..c01774a 100644
--- a/examples/motion_gen_example.py
+++ b/examples/motion_gen_example.py
@@ -25,7 +25,7 @@ from curobo.util_file import get_robot_configs_path, get_world_configs_path, joi
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
-def plot_traj(trajectory, dt):
+def plot_traj(trajectory, dt, file_name="test.png"):
# Third Party
import matplotlib.pyplot as plt
@@ -42,7 +42,7 @@ def plot_traj(trajectory, dt):
axs[3].plot(timesteps, qddd[:, i], label=str(i))
plt.legend()
- plt.savefig("test.png")
+ plt.savefig(file_name)
plt.close()
# plt.show()
@@ -91,6 +91,53 @@ def plot_iters_traj_3d(trajectory, d_id=1, dof=7, seed=0):
plt.show()
+def demo_motion_gen_simple():
+ world_config = {
+ "mesh": {
+ "base_scene": {
+ "pose": [10.5, 0.080, 1.6, 0.043, -0.471, 0.284, 0.834],
+ "file_path": "scene/nvblox/srl_ur10_bins.obj",
+ },
+ },
+ "cuboid": {
+ "table": {
+ "dims": [5.0, 5.0, 0.2], # x, y, z
+ "pose": [0.0, 0.0, -0.1, 1, 0, 0, 0.0], # x, y, z, qw, qx, qy, qz
+ },
+ },
+ }
+ motion_gen_config = MotionGenConfig.load_from_robot_config(
+ "ur5e.yml",
+ world_config,
+ interpolation_dt=0.01,
+ )
+ motion_gen = MotionGen(motion_gen_config)
+ motion_gen.warmup()
+
+ retract_cfg = motion_gen.get_retract_config()
+
+ state = motion_gen.rollout_fn.compute_kinematics(
+ JointState.from_position(retract_cfg.view(1, -1))
+ )
+
+ goal_pose = Pose.from_list([-0.4, 0.0, 0.4, 1.0, 0.0, 0.0, 0.0]) # x, y, z, qw, qx, qy, qz
+ start_state = JointState.from_position(
+ torch.zeros(1, 6).cuda(),
+ joint_names=[
+ "shoulder_pan_joint",
+ "shoulder_lift_joint",
+ "elbow_joint",
+ "wrist_1_joint",
+ "wrist_2_joint",
+ "wrist_3_joint",
+ ],
+ )
+
+ result = motion_gen.plan_single(start_state, goal_pose, MotionGenPlanConfig(max_attempts=1))
+ traj = result.get_interpolated_plan() # result.optimized_dt has the dt between timesteps
+ print("Trajectory Generated: ", result.success)
+
+
def demo_motion_gen_mesh():
PLOT = False
tensor_args = TensorDeviceType()
@@ -149,7 +196,7 @@ def demo_motion_gen(js=False):
)
motion_gen = MotionGen(motion_gen_config)
- motion_gen.warmup(parallel_finetune=True)
+ motion_gen.warmup()
# motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js, parallel_finetune=True)
# robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
@@ -169,16 +216,21 @@ def demo_motion_gen(js=False):
result = motion_gen.plan_single_js(
start_state,
goal_state,
- MotionGenPlanConfig(max_attempts=1, parallel_finetune=True),
+ MotionGenPlanConfig(max_attempts=1, time_dilation_factor=0.5),
)
else:
result = motion_gen.plan_single(
start_state,
retract_pose,
MotionGenPlanConfig(
- max_attempts=1, parallel_finetune=True, enable_finetune_trajopt=True
+ max_attempts=1,
+ timeout=5,
+ time_dilation_factor=0.5,
),
)
+ new_result = result.clone()
+ new_result.retime_trajectory(0.5, create_interpolation_buffer=True)
+ print(new_result.optimized_dt, new_result.motion_time, result.motion_time)
print(
"Trajectory Generated: ",
result.success,
@@ -190,6 +242,7 @@ def demo_motion_gen(js=False):
traj = result.get_interpolated_plan()
plot_traj(traj, result.interpolation_dt)
+ plot_traj(new_result.get_interpolated_plan(), new_result.interpolation_dt, "test_slow.png")
# plt.save("test.png")
# plt.close()
# traj = result.debug_info["opt_solution"].position.view(-1,7)
@@ -431,6 +484,7 @@ def demo_motion_gen_batch_env(n_envs: int = 10):
if __name__ == "__main__":
setup_curobo_logger("error")
demo_motion_gen(js=False)
+ # demo_motion_gen_simple()
# demo_motion_gen_batch()
# demo_motion_gen_goalset()
# n = [2, 10]
diff --git a/setup.cfg b/setup.cfg
index 3b8256c..f93dbfe 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -120,7 +120,7 @@ doc =
# NOTE (roflaherty): Flake8 doesn't support pyproject.toml configuration yet.
[flake8]
-max-line-length = 100
+max-line-length = 99
docstring-convention = google
exclude = .git,build,deprecated,dist,venv
ignore =
diff --git a/src/curobo/__init__.py b/src/curobo/__init__.py
index 403c5fa..85e12af 100644
--- a/src/curobo/__init__.py
+++ b/src/curobo/__init__.py
@@ -10,17 +10,18 @@
#
"""
-cuRobo provides accelerated modules for robotics which can be used to build high-performance
+cuRobo provides accelerated modules for robotics which can be used to build high-performance
robotics applications. The library has several modules for numerical optimization, robot kinematics,
-geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for
-performing tasks like collision-free inverse kinematics, model predictive control, and motion
+geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for
+performing tasks like collision-free inverse kinematics, model predictive control, and motion
planning.
High-level APIs:
+- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`.
- Inverse Kinematics: :mod:`curobo.wrap.reacher.ik_solver`.
- Model Predictive Control: :mod:`curobo.wrap.reacher.mpc`.
-- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`.
+- Trajectory Optimization: :mod:`curobo.wrap.reacher.trajopt`.
cuRobo package is split into several modules:
@@ -33,10 +34,11 @@ cuRobo package is split into several modules:
- :mod:`curobo.rollout` contains methods that map actions to costs. This class wraps instances of
:mod:`curobo.cuda_robot_model` and :mod:`curobo.geom` to compute costs given trajectory of actions.
- :mod:`curobo.util` contains utility methods.
-- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of
+- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of
collision-free reacher and batched robot world collision checking.
-- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
- :meth:`types.state.JointState`, :meth:`types.camera.CameraObservation`, :meth:`types.math.Pose`.
+- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
+ :py:meth:`~types.state.JointState`, :py:meth:`~types.camera.CameraObservation`,
+ :py:meth:`~types.math.Pose`.
"""
diff --git a/src/curobo/content/configs/robot/franka.yml b/src/curobo/content/configs/robot/franka.yml
index c555f34..351bbd0 100644
--- a/src/curobo/content/configs/robot/franka.yml
+++ b/src/curobo/content/configs/robot/franka.yml
@@ -28,7 +28,7 @@ robot_cfg:
"panda_finger_joint1": "Y",
"panda_finger_joint2": "Y",
}
-
+
usd_flip_joint_limits: ["panda_finger_joint2"]
urdf_path: "robot/franka_description/franka_panda.urdf"
asset_root_path: "robot/franka_description"
@@ -67,13 +67,13 @@ robot_cfg:
"panda_hand": ["panda_leftfinger", "panda_rightfinger","attached_object"],
"panda_leftfinger": ["panda_rightfinger", "attached_object"],
"panda_rightfinger": ["attached_object"],
-
+
}
-
+
self_collision_buffer:
{
- "panda_link0": 0.1,
- "panda_link1": 0.05,
+ "panda_link0": 0.1,
+ "panda_link1": 0.05,
"panda_link2": 0.0,
"panda_link3": 0.0,
"panda_link4": 0.0,
@@ -101,7 +101,7 @@ robot_cfg:
"panda_rightfinger",
]
lock_joints: {"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}
- extra_links: {"attached_object":{"parent_link_name": "panda_hand" ,
+ extra_links: {"attached_object":{"parent_link_name": "panda_hand" ,
"link_name": "attached_object", "fixed_transform": [0,0,0,1,0,0,0], "joint_type":"FIXED",
"joint_name": "attach_joint" }}
cspace:
diff --git a/src/curobo/content/configs/robot/spheres/franka_mesh.yml b/src/curobo/content/configs/robot/spheres/franka_mesh.yml
index 1dec717..b7ecae5 100644
--- a/src/curobo/content/configs/robot/spheres/franka_mesh.yml
+++ b/src/curobo/content/configs/robot/spheres/franka_mesh.yml
@@ -134,11 +134,11 @@ collision_spheres:
"radius": 0.022
panda_leftfinger:
- "center": [0.0, 0.01, 0.043]
- "radius": 0.011 #0.025 # 0.011
+ "radius": 0.011
- "center": [0.0, 0.02, 0.015]
- "radius": 0.011 #0.025 # 0.011
+ "radius": 0.011
panda_rightfinger:
- "center": [0.0, -0.01, 0.043]
- "radius": 0.011 #0.025 #0.011
+ "radius": 0.011
- "center": [0.0, -0.02, 0.015]
- "radius": 0.011 #0.025 #0.011
\ No newline at end of file
+ "radius": 0.011
\ No newline at end of file
diff --git a/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml b/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml
index e146757..1a377ec 100644
--- a/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml
+++ b/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml
@@ -23,7 +23,7 @@ robot_cfg:
base_link: "base_link"
ee_link: "grasp_frame"
link_names: null
- lock_joints: {'finger_joint': 0.7}
+ lock_joints: {'finger_joint': 0.3}
extra_links: null #{"attached_object":{"parent_link_name": "grasp_frame" ,
#"link_name": "attached_object", "fixed_transform": [0,0,0,1,0,0,0], "joint_type":"FIXED",
diff --git a/src/curobo/content/configs/task/base_cfg.yml b/src/curobo/content/configs/task/base_cfg.yml
index bd1eb42..27b79c2 100644
--- a/src/curobo/content/configs/task/base_cfg.yml
+++ b/src/curobo/content/configs/task/base_cfg.yml
@@ -11,8 +11,8 @@
world_collision_checker_cfg:
- cache: null #{"cube": 41, "capsule": 0, "sphere": 0}
- checker_type: "PRIMITIVE" # ["PRIMITIVE", "BLOX", "MESH"]
+ cache: null
+ checker_type: "PRIMITIVE"
max_distance: 0.1
diff --git a/src/curobo/content/configs/task/finetune_trajopt.yml b/src/curobo/content/configs/task/finetune_trajopt.yml
index 6c1c66d..4ae1f61 100644
--- a/src/curobo/content/configs/task/finetune_trajopt.yml
+++ b/src/curobo/content/configs/task/finetune_trajopt.yml
@@ -19,12 +19,12 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
- base_dt: 0.2
+ base_dt: 0.15
base_ratio: 1.0
- max_dt: 0.2
+ max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
- state_finite_difference_mode: "CENTRAL"
+ state_finite_difference_mode: "CENTRAL"
teleport_mode: False
return_full_act_buffer: True
@@ -35,7 +35,7 @@ cost:
weight: [2000,500000.0,30,60]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
- run_weight: 1.0
+ run_weight: 1.0
use_metric: True
link_pose_cfg:
@@ -47,19 +47,19 @@ cost:
run_weight: 1.0
use_metric: True
-
+
cspace_cfg:
weight: 10000.0
terminal: True
run_weight: 0.00 #1
bound_cfg:
- weight: [10000.0, 500000.0, 500.0, 500.0]
- smooth_weight: [0.0,10000.0, 5.0, 0.0]
+ weight: [10000.0, 500000.0, 500.0, 500.0]
+ smooth_weight: [0.0,10000.0, 5.0, 0.0]
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
- activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk
+ activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -75,21 +75,21 @@ cost:
self_collision_cfg:
weight: 5000.0
classify: False
-
-
+
+
lbfgs:
- n_iters: 300
- inner_iters: 25
- cold_start_n_iters: null
+ n_iters: 400
+ inner_iters: 25
+ cold_start_n_iters: null
min_iters: 25
line_search_scale: [0.1, 0.3, 0.7,1.0] # #
fixed_iters: True
- cost_convergence: 0.01
- cost_delta_threshold: 1.0
- cost_relative_threshold: 0.999 #0.999
+ cost_convergence: 0.001
+ cost_delta_threshold: 0.0001
+ cost_relative_threshold: 0.99
epsilon: 0.01
- history: 27 #15
+ history: 27
use_cuda_graph: True
n_problems: 1
store_debug: False
diff --git a/src/curobo/content/configs/task/finetune_trajopt_slow.yml b/src/curobo/content/configs/task/finetune_trajopt_slow.yml
index 383b903..8eb26c8 100644
--- a/src/curobo/content/configs/task/finetune_trajopt_slow.yml
+++ b/src/curobo/content/configs/task/finetune_trajopt_slow.yml
@@ -19,9 +19,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
- base_dt: 0.2
+ base_dt: 0.15
base_ratio: 1.0
- max_dt: 0.2
+ max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
@@ -31,22 +31,21 @@ model:
cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
- run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position
- weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
+ run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
+ weight: [2000,500000.0,30,60]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
- run_weight: 0.0 #0.05
+ run_weight: 1.0
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
- run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position
- weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
+ run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
+ weight: [2000,500000.0,30,60] #[150.0, 2000.0, 30, 40]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
- run_weight: 0.001
+ run_weight: 1.0
use_metric: True
-
cspace_cfg:
weight: 10000.0
@@ -54,12 +53,12 @@ cost:
run_weight: 0.00 #1
bound_cfg:
- weight: [10000.0, 100000.0, 500.0, 500.0]
- smooth_weight: [0.0,10000.0, 50.0, 0.0]
+ weight: [10000.0, 500000.0, 500.0, 500.0]
+ smooth_weight: [0.0,10000.0, 5.0, 0.0]
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
- activation_distance: [0.1,0.1,0.1,10.0] # for position, velocity, acceleration and jerk
+ activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -75,21 +74,20 @@ cost:
self_collision_cfg:
weight: 5000.0
classify: False
-
-
+
lbfgs:
- n_iters: 400 # 400
+ n_iters: 400
inner_iters: 25
- cold_start_n_iters: 200
+ cold_start_n_iters: null
min_iters: 25
- line_search_scale: [0.01, 0.3, 0.7,1.0] # #
+ line_search_scale: [0.1, 0.3, 0.7,1.0] # #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 1.0
cost_relative_threshold: 0.999 #0.999
epsilon: 0.01
- history: 15 #15
+ history: 27 #15
use_cuda_graph: True
n_problems: 1
store_debug: False
diff --git a/src/curobo/content/configs/task/gradient_ik.yml b/src/curobo/content/configs/task/gradient_ik.yml
index df3d90d..98354df 100644
--- a/src/curobo/content/configs/task/gradient_ik.yml
+++ b/src/curobo/content/configs/task/gradient_ik.yml
@@ -29,32 +29,30 @@ cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [2000,10000,20,40]
- vec_convergence: [0.0, 0.00]
+ vec_convergence: [0.0, 0.00]
terminal: False
use_metric: True
run_weight: 1.0
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [2000,10000,20,40]
- vec_convergence: [0.00, 0.000]
+ vec_convergence: [0.00, 0.000]
terminal: False
use_metric: True
run_weight: 1.0
-
-
cspace_cfg:
weight: 0.000
bound_cfg:
- weight: 500.0
- activation_distance: [0.1]
- null_space_weight: [1.0]
+ weight: 5000.0
+ activation_distance: [0.001]
+ null_space_weight: [0.1]
primitive_collision_cfg:
weight: 5000.0
use_sweep: False
classify: False
activation_distance: 0.01
self_collision_cfg:
- weight: 500.0
+ weight: 5000.0
classify: False
@@ -81,10 +79,10 @@ lbfgs:
use_cuda_line_search_kernel: True
use_cuda_update_best_kernel: True
sync_cuda_time: True
- step_scale: 1.0
+ step_scale: 1.0
use_coo_sparse: True
last_best: 10
debug_info:
visual_traj : null #'ee_pos_seq'
-
+
diff --git a/src/curobo/content/configs/task/gradient_trajopt.yml b/src/curobo/content/configs/task/gradient_trajopt.yml
index 506a25e..da82760 100644
--- a/src/curobo/content/configs/task/gradient_trajopt.yml
+++ b/src/curobo/content/configs/task/gradient_trajopt.yml
@@ -19,12 +19,12 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
- base_dt: 0.2
+ base_dt: 0.15
base_ratio: 1.0
- max_dt: 0.2
+ max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
- state_finite_difference_mode: "CENTRAL"
+ state_finite_difference_mode: "CENTRAL"
teleport_mode: False
return_full_act_buffer: True
@@ -48,17 +48,17 @@ cost:
use_metric: True
cspace_cfg:
- weight: 1000.0
+ weight: 10000.0
terminal: True
run_weight: 0.0
bound_cfg:
- weight: [50000.0, 50000.0, 50000.0,50000.0]
- smooth_weight: [0.0,10000.0,5.0, 0.0]
+ weight: [50000.0, 50000.0, 500.0,500.0]
+ smooth_weight: [0.0,10000.0,5.0, 0.0]
run_weight_velocity: 0.00
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
- activation_distance: [0.1,0.1,0.1,0.1] # for position, velocity, acceleration and jerk
+ activation_distance: [0.05,0.05,0.05,0.05] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
@@ -75,13 +75,13 @@ cost:
self_collision_cfg:
weight: 5000.0
classify: False
-
-
+
+
lbfgs:
n_iters: 100 #175
- inner_iters: 25
- cold_start_n_iters: null
+ inner_iters: 25
+ cold_start_n_iters: null
min_iters: 25
line_search_scale: [0.1,0.3,0.7,1.0]
fixed_iters: True
diff --git a/src/curobo/content/configs/task/particle_trajopt.yml b/src/curobo/content/configs/task/particle_trajopt.yml
index 268a613..0f2e288 100644
--- a/src/curobo/content/configs/task/particle_trajopt.yml
+++ b/src/curobo/content/configs/task/particle_trajopt.yml
@@ -18,9 +18,9 @@ model:
acceleration: 1.0
enable: False
dt_traj_params:
- base_dt: 0.2
+ base_dt: 0.15
base_ratio: 1.0
- max_dt: 0.2
+ max_dt: 0.15
vel_scale: 1.0
control_space: 'POSITION'
teleport_mode: False
@@ -30,22 +30,22 @@ model:
cost:
pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
- run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
+ run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
weight: [250.0, 5000.0, 20, 20]
- vec_convergence: [0.0,0.0,1000.0,1000.0]
+ vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True
run_weight: 0.0
use_metric: True
-
+
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
- run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
+ run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
weight: [0.0, 5000.0, 40, 40]
- vec_convergence: [0.0,0.0,1000.0,1000.0]
+ vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True
- run_weight: 0.00
+ run_weight: 0.00
use_metric: True
-
+
cspace_cfg:
weight: 10000.0
terminal: True
@@ -74,8 +74,8 @@ cost:
self_collision_cfg:
weight: 500.0
classify: False
-
-
+
+
mppi:
@@ -89,7 +89,7 @@ mppi:
alpha : 1
num_particles : 25 # 100
update_cov : True
- cov_type : "DIAG_A" #
+ cov_type : "DIAG_A" #
kappa : 0.001
null_act_frac : 0.0
sample_mode : 'MEAN'
diff --git a/src/curobo/cuda_robot_model/cuda_robot_model.py b/src/curobo/cuda_robot_model/cuda_robot_model.py
index 3135bb5..c7c99e3 100644
--- a/src/curobo/cuda_robot_model/cuda_robot_model.py
+++ b/src/curobo/cuda_robot_model/cuda_robot_model.py
@@ -24,11 +24,7 @@ from curobo.cuda_robot_model.cuda_robot_generator import (
CudaRobotGeneratorConfig,
)
from curobo.cuda_robot_model.kinematics_parser import KinematicsParser
-from curobo.cuda_robot_model.types import (
- CudaRobotModelState,
- KinematicsTensorConfig,
- SelfCollisionKinematicsConfig,
-)
+from curobo.cuda_robot_model.types import KinematicsTensorConfig, SelfCollisionKinematicsConfig
from curobo.curobolib.kinematics import get_cuda_kinematics
from curobo.geom.types import Mesh, Sphere
from curobo.types.base import TensorDeviceType
@@ -148,14 +144,67 @@ class CudaRobotModelConfig:
return self.kinematics_config.n_dof
+@dataclass
+class CudaRobotModelState:
+ """Dataclass that stores kinematics information."""
+
+ #: End-effector position stored as x,y,z in meters [b, 3]. End-effector is defined by
+ #: :attr:`CudaRobotModel.ee_link`.
+ ee_position: torch.Tensor
+
+ #: End-effector orientaiton stored as quaternion qw, qx, qy, qz [b,4]. End-effector is defined
+ #: by :attr:`CudaRobotModel.ee_link`.
+ ee_quaternion: torch.Tensor
+
+ #: Linear Jacobian. Currently not supported.
+ lin_jacobian: Optional[torch.Tensor] = None
+
+ #: Angular Jacobian. Currently not supported.
+ ang_jacobian: Optional[torch.Tensor] = None
+
+ #: Position of links specified by link_names (:attr:`CudaRobotModel.link_names`).
+ links_position: Optional[torch.Tensor] = None
+
+ #: Quaternions of links specified by link names (:attr:`CudaRobotModel.link_names`).
+ links_quaternion: Optional[torch.Tensor] = None
+
+ #: Position of spheres specified by collision spheres (:attr:`CudaRobotModel.robot_spheres`)
+ #: in x, y, z, r format [b,n,4].
+ link_spheres_tensor: Optional[torch.Tensor] = None
+
+ #: Names of links that each index in :attr:`links_position` and :attr:`links_quaternion`
+ #: corresponds to.
+ link_names: Optional[str] = None
+
+ @property
+ def ee_pose(self) -> Pose:
+ """Get end-effector pose as a Pose object."""
+ return Pose(self.ee_position, self.ee_quaternion)
+
+ def get_link_spheres(self) -> torch.Tensor:
+ """Get spheres representing robot geometry as a tensor with [batch,4], [x,y,z,radius]."""
+ return self.link_spheres_tensor
+
+ @property
+ def link_pose(self) -> Dict[str, Pose]:
+ """Get link poses as a dictionary of link name to Pose object."""
+ link_poses = None
+ if self.link_names is not None:
+ link_poses = {}
+ link_pos = self.links_position.contiguous()
+ link_quat = self.links_quaternion.contiguous()
+ for i, v in enumerate(self.link_names):
+ link_poses[v] = Pose(link_pos[..., i, :], link_quat[..., i, :])
+ return link_poses
+
+
class CudaRobotModel(CudaRobotModelConfig):
"""
CUDA Accelerated Robot Model
- NOTE: Currently dof is created only for links that we need to compute kinematics.
- E.g., for robots with many serial chains, add all links of the robot to get the correct dof.
- This is not an issue if you are loading collision spheres as that will cover the full geometry
- of the robot.
+ Currently dof is created only for links that we need to compute kinematics. E.g., for robots
+ with many serial chains, add all links of the robot to get the correct dof. This is not an
+ issue if you are loading collision spheres as that will cover the full geometry of the robot.
"""
def __init__(self, config: CudaRobotModelConfig):
@@ -421,6 +470,10 @@ class CudaRobotModel(CudaRobotModelConfig):
def base_link(self):
return self.kinematics_config.base_link
+ @property
+ def robot_spheres(self):
+ return self.kinematics_config.link_spheres
+
def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig):
self.kinematics_config.copy_(new_kin_config)
diff --git a/src/curobo/cuda_robot_model/types.py b/src/curobo/cuda_robot_model/types.py
index 8c4c20e..a213e5a 100644
--- a/src/curobo/cuda_robot_model/types.py
+++ b/src/curobo/cuda_robot_model/types.py
@@ -496,52 +496,3 @@ class SelfCollisionKinematicsConfig:
collision_matrix: Optional[torch.Tensor] = None
experimental_kernel: bool = True
checks_per_thread: int = 32
-
-
-@dataclass
-class CudaRobotModelState:
- """Dataclass that stores kinematics information."""
-
- #: End-effector position stored as x,y,z in meters [b, 3]. End-effector is defined by
- #: :py:attr:`curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGeneratorConfig.ee_link`.
- ee_position: torch.Tensor
-
- #: End-effector orientaiton stored as quaternion qw, qx, qy, qz [b,4]. End-effector is defined
- # by :py:attr:`CudaRobotModelConfig.ee_link`.
- ee_quaternion: torch.Tensor
-
- #: Linear Jacobian. Currently not supported.
- lin_jacobian: Optional[torch.Tensor] = None
-
- #: Angular Jacobian. Currently not supported.
- ang_jacobian: Optional[torch.Tensor] = None
-
- #: Position of links specified by link_names (:py:attr:`CudaRobotModelConfig.link_names`).
- links_position: Optional[torch.Tensor] = None
-
- #: Quaternions of links specified by link names (:py:attr:`CudaRobotModelConfig.link_names`).
- links_quaternion: Optional[torch.Tensor] = None
-
- #: Position of spheres specified by collision spheres (:py:attr:`CudaRobotModelConfig.collision_spheres`)
- #: in x, y, z, r format [b,n,4].
- link_spheres_tensor: Optional[torch.Tensor] = None
-
- link_names: Optional[str] = None
-
- @property
- def ee_pose(self):
- return Pose(self.ee_position, self.ee_quaternion)
-
- def get_link_spheres(self):
- return self.link_spheres_tensor
-
- @property
- def link_pose(self):
- link_poses = None
- if self.link_names is not None:
- link_poses = {}
- link_pos = self.links_position.contiguous()
- link_quat = self.links_quaternion.contiguous()
- for i, v in enumerate(self.link_names):
- link_poses[v] = Pose(link_pos[..., i, :], link_quat[..., i, :])
- return link_poses
diff --git a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu
index da250af..687344e 100644
--- a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu
+++ b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu
@@ -194,13 +194,13 @@ namespace Curobo
const scalar_t x = transform_mat[4];
const scalar_t y = transform_mat[5];
const scalar_t z = transform_mat[6];
-
- if ((x != 0) || (y != 0) || (z != 0))
+ if(x != 0 || y!= 0 || z!=0)
{
C.x = p_arr.x + w * w * sphere_pos.x + 2 * y * w * sphere_pos.z -
2 * z * w * sphere_pos.y + x * x * sphere_pos.x +
2 * y * x * sphere_pos.y + 2 * z * x * sphere_pos.z -
z * z * sphere_pos.x - y * y * sphere_pos.x;
+
C.y = p_arr.y + 2 * x * y * sphere_pos.x + y * y * sphere_pos.y +
2 * z * y * sphere_pos.z + 2 * w * z * sphere_pos.x -
z * z * sphere_pos.y + w * w * sphere_pos.y - 2 * x * w * sphere_pos.z -
@@ -714,7 +714,7 @@ float4 &sum_pt)
// compute interpolation distance between voxel origin and sphere location:
- get_array_value(grid_features, voxel_idx, interpolated_distance);// + 0.5 * loc_grid_params.w;
+ get_array_value(grid_features, voxel_idx, interpolated_distance);
if(!INTERPOLATION)
{
interpolated_distance += 0.5 * loc_grid_params.w;//max(0.0, (0.3 * loc_grid_params.w) - loc_sphere.w);
@@ -821,7 +821,7 @@ float4 &sum_pt)
d_grad = (d_plus - d_minus) * (1/(2*voxel_size));
}
- else
+ if (!CENTRAL_DIFFERENCE)
{
// x difference:
int x_minus,y_minus, z_minus;
@@ -1252,7 +1252,9 @@ float4 &sum_pt)
const int32_t env_idx,
const int bn_sph_idx,
const int sph_idx,
- dist_scalar_t *out_distance, grad_scalar_t *closest_pt, uint8_t *sparsity_idx,
+ dist_scalar_t *out_distance,
+ grad_scalar_t *closest_pt,
+ uint8_t *sparsity_idx,
const float *weight,
const float *activation_distance,
const float *max_distance,
@@ -1863,7 +1865,9 @@ float4 &sum_pt)
const int32_t env_idx,
const int bn_sph_idx,
const int sph_idx,
- dist_scalar_t *out_distance, grad_scalar_t *closest_pt, uint8_t *sparsity_idx,
+ dist_scalar_t *out_distance,
+ grad_scalar_t *closest_pt,
+ uint8_t *sparsity_idx,
const float *weight,
const float *activation_distance,
const float *max_distance,
@@ -1879,7 +1883,6 @@ float4 &sum_pt)
const float eta = max_distance_local;
float max_dist = -1 * max_distance_local;
max_distance_local = -1 * max_distance_local;
-;
float3 max_grad = make_float3(0.0, 0.0, 0.0);
diff --git a/src/curobo/geom/cv.py b/src/curobo/geom/cv.py
index b2b9da0..8d0e007 100644
--- a/src/curobo/geom/cv.py
+++ b/src/curobo/geom/cv.py
@@ -47,7 +47,9 @@ def project_depth_to_pointcloud(depth_image: torch.Tensor, intrinsics_matrix: to
@get_torch_jit_decorator()
-def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor):
+def get_projection_rays(
+ height: int, width: int, intrinsics_matrix: torch.Tensor, depth_to_meter: float = 0.001
+):
"""Projects numpy depth image to point cloud.
Args:
@@ -82,7 +84,7 @@ def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor
rays = torch.stack([output_x, output_y, input_z], -1).reshape(
intrinsics_matrix.shape[0], width * height, 3
)
- rays = rays * (1.0 / 1000.0)
+ rays = rays * depth_to_meter
return rays
diff --git a/src/curobo/geom/sdf/world.py b/src/curobo/geom/sdf/world.py
index 03d6550..fbe10eb 100644
--- a/src/curobo/geom/sdf/world.py
+++ b/src/curobo/geom/sdf/world.py
@@ -264,7 +264,7 @@ class WorldCollisionConfig:
n_envs: int = 1
checker_type: CollisionCheckerType = CollisionCheckerType.PRIMITIVE
max_distance: Union[torch.Tensor, float] = 0.1
- max_esdf_distance: Union[torch.Tensor, float] = 1000.0
+ max_esdf_distance: Union[torch.Tensor, float] = 100.0
def __post_init__(self):
if self.world_model is not None and isinstance(self.world_model, list):
@@ -398,6 +398,7 @@ class WorldCollision(WorldCollisionConfig):
self._cache_voxelization is None
or self._cache_voxelization.voxel_size != new_grid.voxel_size
or self._cache_voxelization.dims != new_grid.dims
+ or self._cache_voxelization.xyzr_tensor is None
):
self._cache_voxelization = new_grid
self._cache_voxelization.xyzr_tensor = self._cache_voxelization.create_xyzr_tensor(
@@ -458,7 +459,6 @@ class WorldCollision(WorldCollisionConfig):
self.update_cache_voxelization(new_grid)
xyzr = self._cache_voxelization.xyzr_tensor
- voxel_shape = xyzr.shape
xyzr = xyzr.view(-1, 1, 1, 4)
weight = self.tensor_args.to_device([1.0])
diff --git a/src/curobo/geom/sdf/world_blox.py b/src/curobo/geom/sdf/world_blox.py
index 3340406..0cba6d1 100644
--- a/src/curobo/geom/sdf/world_blox.py
+++ b/src/curobo/geom/sdf/world_blox.py
@@ -127,6 +127,7 @@ class WorldBloxCollision(WorldVoxelCollision):
collision_query_buffer: CollisionQueryBuffer,
weight,
activation_distance,
+ return_loss: bool = False,
compute_esdf: bool = False,
):
d = self._blox_mapper.query_sphere_sdf_cost(
@@ -136,8 +137,11 @@ class WorldBloxCollision(WorldVoxelCollision):
collision_query_buffer.blox_collision_buffer.sparsity_index_buffer,
weight,
activation_distance,
+ self.max_esdf_distance,
self._blox_tensor_list[0],
self._blox_tensor_list[1],
+ return_loss,
+ compute_esdf,
)
return d
@@ -191,12 +195,12 @@ class WorldBloxCollision(WorldVoxelCollision):
sum_collisions=sum_collisions,
compute_esdf=compute_esdf,
)
-
d = self._get_blox_sdf(
query_sphere,
collision_query_buffer,
weight=weight,
activation_distance=activation_distance,
+ return_loss=return_loss,
compute_esdf=compute_esdf,
)
@@ -246,6 +250,7 @@ class WorldBloxCollision(WorldVoxelCollision):
collision_query_buffer,
weight=weight,
activation_distance=activation_distance,
+ return_loss=return_loss,
)
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
"mesh" not in self.collision_types or not self.collision_types["mesh"]
@@ -404,7 +409,7 @@ class WorldBloxCollision(WorldVoxelCollision):
cuboid: Cuboid,
layer_name: Optional[str] = None,
):
- raise NotImplementedError
+ log_error("Not implemented")
def get_bounding_spheres(
self,
@@ -418,7 +423,8 @@ class WorldBloxCollision(WorldVoxelCollision):
clear_region: bool = False,
) -> List[Sphere]:
mesh = self.get_mesh_in_bounding_box(bounding_box, obstacle_name, clear_region=clear_region)
-
+ if clear_region:
+ self.clear_bounding_box(bounding_box, obstacle_name)
spheres = mesh.get_bounding_spheres(
n_spheres=n_spheres,
surface_sphere_radius=surface_sphere_radius,
diff --git a/src/curobo/geom/sdf/world_voxel.py b/src/curobo/geom/sdf/world_voxel.py
index 9eb9c62..0cdaad5 100644
--- a/src/curobo/geom/sdf/world_voxel.py
+++ b/src/curobo/geom/sdf/world_voxel.py
@@ -80,6 +80,9 @@ class WorldVoxelCollision(WorldMeshCollision):
if feature_dtype in [torch.float32, torch.float16, torch.bfloat16]:
voxel_features[:] = -1.0 * self.max_esdf_distance
else:
+ if self.max_esdf_distance > 100.0:
+ log_warn("Using fp8 for WorldVoxelCollision will reduce max_esdf_distance to 100")
+ self.max_esdf_distance = 100.0
voxel_features = (voxel_features.to(dtype=torch.float16) - self.max_esdf_distance).to(
dtype=feature_dtype
)
@@ -503,9 +506,10 @@ class WorldVoxelCollision(WorldMeshCollision):
False,
use_batch_env,
False,
- True,
+ False,
False,
)
+
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
"mesh" not in self.collision_types or not self.collision_types["mesh"]
):
diff --git a/src/curobo/geom/sphere_fit.py b/src/curobo/geom/sphere_fit.py
index d8aa65d..2542aa0 100644
--- a/src/curobo/geom/sphere_fit.py
+++ b/src/curobo/geom/sphere_fit.py
@@ -92,6 +92,7 @@ def voxel_fit_surface_mesh(
def get_voxelgrid_from_mesh(mesh: trimesh.Trimesh, n_spheres: int, voxelize_method: str = "ray"):
+ """Get voxel grid from mesh using :py:func:`trimesh.voxel.creation.voxelize`."""
pitch = get_voxel_pitch(mesh, n_spheres)
radius = pitch / 2.0
try:
diff --git a/src/curobo/geom/types.py b/src/curobo/geom/types.py
index 8a15488..fe8b500 100644
--- a/src/curobo/geom/types.py
+++ b/src/curobo/geom/types.py
@@ -58,9 +58,9 @@ class Obstacle:
texture: Optional[str] = None
#: material properties to apply in visualization.
- material: Material = Material()
+ material: Material = field(default_factory=Material)
- tensor_args: TensorDeviceType = TensorDeviceType()
+ tensor_args: TensorDeviceType = field(default_factory=TensorDeviceType)
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True) -> trimesh.Trimesh:
"""Create a trimesh instance from the obstacle representation.
diff --git a/src/curobo/opt/newton/newton_base.py b/src/curobo/opt/newton/newton_base.py
index ea9061d..590d18e 100644
--- a/src/curobo/opt/newton/newton_base.py
+++ b/src/curobo/opt/newton/newton_base.py
@@ -125,7 +125,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
def _shift(self, shift_steps=1):
# TODO: shift best q?:
- self.best_cost[:] = 500000.0
+ self.best_cost[:] = 5000000.0
self.best_iteration[:] = 0
self.current_iteration[:] = 0
return True
@@ -159,8 +159,9 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
with profiler.record_function("newton/reset"):
self.i = -1
self._opt_finished = False
- self.best_cost[:] = 500000.0
+ self.best_cost[:] = 5000000.0
self.best_iteration[:] = 0
+ self.current_iteration[:] = 0
super().reset()
@@ -448,6 +449,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self.cost_delta_threshold,
self.cost_relative_threshold,
)
+ # print(self.best_cost[0], self.best_q[0])
else:
cost = cost.detach()
q = q.detach()
diff --git a/src/curobo/opt/opt_base.py b/src/curobo/opt/opt_base.py
index f105685..1d41bb3 100644
--- a/src/curobo/opt/opt_base.py
+++ b/src/curobo/opt/opt_base.py
@@ -169,7 +169,7 @@ class Optimizer(OptimizerConfig):
st_time = time.time()
out = self._optimize(opt_tensor, shift_steps, n_iters)
if self.sync_cuda_time:
- torch.cuda.synchronize()
+ torch.cuda.synchronize(device=self.tensor_args.device)
self.opt_dt = time.time() - st_time
return out
diff --git a/src/curobo/rollout/arm_base.py b/src/curobo/rollout/arm_base.py
index f98279f..8560a42 100644
--- a/src/curobo/rollout/arm_base.py
+++ b/src/curobo/rollout/arm_base.py
@@ -384,6 +384,7 @@ class ArmBase(RolloutBase, ArmBaseConfig):
# setup constraint terms:
constraint = self.bound_constraint.forward(state.state_seq)
+
constraint_list = [constraint]
if (
self.constraint_cfg.primitive_collision_cfg is not None
@@ -407,7 +408,9 @@ class ArmBase(RolloutBase, ArmBaseConfig):
self_constraint = self.robot_self_collision_constraint.forward(state.robot_spheres)
constraint_list.append(self_constraint)
constraint = cat_sum(constraint_list)
+
feasible = constraint == 0.0
+
if out_metrics is None:
out_metrics = RolloutMetrics()
out_metrics.feasible = feasible
diff --git a/src/curobo/rollout/arm_reacher.py b/src/curobo/rollout/arm_reacher.py
index 77b970a..ac2cebc 100644
--- a/src/curobo/rollout/arm_reacher.py
+++ b/src/curobo/rollout/arm_reacher.py
@@ -263,7 +263,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
goal_cost = self.goal_cost.forward(
ee_pos_batch, ee_quat_batch, self._goal_buffer
)
-
+ # print(self._compute_g_dist, goal_cost.view(-1))
cost_list.append(goal_cost)
with profiler.record_function("cost/link_poses"):
if self._goal_buffer.links_goal_pose is not None and self.cost_cfg.pose_cfg is not None:
@@ -286,6 +286,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
and self.cost_cfg.cspace_cfg is not None
and self.dist_cost.enabled
):
+
joint_cost = self.dist_cost.forward_target_idx(
self._goal_buffer.goal_state.position,
state_batch.position,
diff --git a/src/curobo/rollout/cost/bound_cost.py b/src/curobo/rollout/cost/bound_cost.py
index 3524cf0..0c91f75 100644
--- a/src/curobo/rollout/cost/bound_cost.py
+++ b/src/curobo/rollout/cost/bound_cost.py
@@ -21,6 +21,7 @@ import warp as wp
from curobo.cuda_robot_model.types import JointLimits
from curobo.types.robot import JointState
from curobo.types.tensor import T_DOF
+from curobo.util.logger import log_error
from curobo.util.torch_utils import get_torch_jit_decorator
from curobo.util.warp import init_warp
@@ -53,6 +54,16 @@ class BoundCostConfig(CostConfig):
self.joint_limits = bounds.clone()
if teleport_mode:
self.cost_type = BoundCostType.POSITION
+ if self.cost_type != BoundCostType.POSITION:
+ if torch.max(self.joint_limits.velocity[1] - self.joint_limits.velocity[0]) == 0.0:
+ log_error("Joint velocity limits is zero")
+ if (
+ torch.max(self.joint_limits.acceleration[1] - self.joint_limits.acceleration[0])
+ == 0.0
+ ):
+ log_error("Joint acceleration limits is zero")
+ if torch.max(self.joint_limits.jerk[1] - self.joint_limits.jerk[0]) == 0.0:
+ log_error("Joint jerk limits is zero")
def __post_init__(self):
if isinstance(self.activation_distance, List):
@@ -675,8 +686,14 @@ def forward_bound_pos_warp(
target_p = retract_config[target_id]
p_l = p_b[d_id]
p_u = p_b[dof + d_id]
+
+ p_range = p_u - p_l
+ eta_p = eta_p * (p_range)
p_l += eta_p
p_u -= eta_p
+ if p_range < 1.0:
+ w = (1.0 / p_range) * w
+
# compute cost:
b_addrs = b_id * horizon * dof + h_id * dof + d_id
@@ -690,33 +707,15 @@ def forward_bound_pos_warp(
g_p = 2.0 * n_w * error
# bound cost:
- # if c_p < p_l:
- # delta = p_l - c_p
- # if (delta) > eta_p or eta_p == 0.0:
- # c_total += w * (delta - 0.5 * eta_p)
- # g_p += -w
- # else:
- # c_total += w * (0.5 / eta_p) * delta * delta
- # g_p += -w * (1.0 / eta_p) * delta
- # elif c_p > p_u:
- # delta = c_p - p_u
- # if (delta) > eta_p or eta_p == 0.0:
- # c_total += w * (delta - 0.5 * eta_p)
- # g_p += w
- # else:
- # c_total += w * (0.5 / eta_p) * delta * delta
- # g_p += w * (1.0 / eta_p) * delta
- # bound cost:
+ delta = 0.0
if c_p < p_l:
delta = c_p - p_l
- c_total += w * delta * delta
- g_p += 2.0 * w * delta
elif c_p > p_u:
delta = c_p - p_u
- c_total += w * delta * delta
- g_p += 2.0 * w * delta
+ c_total += w * delta * delta
+ g_p += 2.0 * w * delta
out_cost[b_addrs] = c_total
# compute gradient
@@ -807,16 +806,33 @@ def forward_bound_warp(
c_j = jerk[b_addrs]
- p_l = p_b[d_id] + eta_p
- p_u = p_b[dof + d_id] - eta_p
+ p_l = p_b[d_id]
+ p_u = p_b[dof + d_id]
+ p_range = p_u - p_l
+ eta_p = eta_p * (p_range)
+ p_l += eta_p
+ p_u -= eta_p
- v_l = v_b[d_id] + eta_v
- v_u = v_b[dof + d_id] - eta_v
- a_l = a_b[d_id] + eta_a
- a_u = a_b[dof + d_id] - eta_a
+ v_l = v_b[d_id]
+ v_u = v_b[dof + d_id]
+ v_range = v_u - v_l
+ eta_v = eta_v * (v_range)
+ v_l += eta_v
+ v_u -= eta_v
- j_l = j_b[d_id] + eta_j
- j_u = j_b[dof + d_id] - eta_j
+ a_l = a_b[d_id]
+ a_u = a_b[dof + d_id]
+ a_range = a_u - a_l
+ eta_a = eta_a * (a_range)
+ a_l += eta_a
+ a_u -= eta_a
+
+ j_l = j_b[d_id]
+ j_u = j_b[dof + d_id]
+ j_range = j_u - j_l
+ eta_j = eta_j * (j_range)
+ j_l += eta_j
+ j_u -= eta_j
delta = float(0.0)
if n_w > 0.0:
@@ -826,6 +842,7 @@ def forward_bound_warp(
# bound cost:
delta = 0.0
+
if c_p < p_l:
delta = c_p - p_l
elif c_p > p_u:
@@ -995,19 +1012,47 @@ def forward_bound_smooth_warp(
r_wj = run_weight_jerk[h_id]
c_j = jerk[b_addrs]
- p_l = p_b[d_id] + eta_p
- p_u = p_b[dof + d_id] - eta_p
+ p_l = p_b[d_id]
+ p_u = p_b[dof + d_id]
+ p_range = p_u - p_l
+ eta_p = eta_p * (p_range)
+ p_l += eta_p
+ p_u -= eta_p
- v_l = v_b[d_id] + eta_v
- v_u = v_b[dof + d_id] - eta_v
- a_l = a_b[d_id] + eta_a
- a_u = a_b[dof + d_id] - eta_a
+ v_l = v_b[d_id]
+ v_u = v_b[dof + d_id]
+ v_range = v_u - v_l
+ eta_v = eta_v * (v_range)
+ v_l += eta_v
+ v_u -= eta_v
- j_l = j_b[d_id] + eta_j
- j_u = j_b[dof + d_id] - eta_j
+ a_l = a_b[d_id]
+ a_u = a_b[dof + d_id]
+ a_range = a_u - a_l
+ eta_a = eta_a * (a_range)
+ a_l += eta_a
+ a_u -= eta_a
+
+ j_l = j_b[d_id]
+ j_u = j_b[dof + d_id]
+ j_range = j_u - j_l
+ eta_j = eta_j * (j_range)
+ j_l += eta_j
+ j_u -= eta_j
delta = float(0.0)
+ if p_range < 1.0:
+ w = w * (1.0 / (p_range * p_range))
+
+ if v_range < 1.0:
+ b_wv = b_wv * (1.0 / (v_range * v_range))
+ if a_range < 1.0:
+ b_wa = b_wa * (1.0 / (a_range * a_range))
+ w_a = w_a * (1.0 / (a_range * a_range))
+ if j_range < 1.0:
+ b_wj = b_wj * (1.0 / (j_range * j_range))
+ w_j = w_j * (1.0 / (j_range * j_range))
# position:
if n_w > 0.0:
error = c_p - target_p
diff --git a/src/curobo/rollout/cost/self_collision_cost.py b/src/curobo/rollout/cost/self_collision_cost.py
index c4469f9..76dc949 100644
--- a/src/curobo/rollout/cost/self_collision_cost.py
+++ b/src/curobo/rollout/cost/self_collision_cost.py
@@ -68,7 +68,6 @@ class SelfCollisionCost(CostBase, SelfCollisionCostConfig):
self.self_collision_kin_config.thread_location,
self.self_collision_kin_config.thread_max,
self.self_collision_kin_config.checks_per_thread,
- # False,
self.self_collision_kin_config.experimental_kernel,
self.return_loss,
)
diff --git a/src/curobo/types/math.py b/src/curobo/types/math.py
index 98908e6..97f11dd 100644
--- a/src/curobo/types/math.py
+++ b/src/curobo/types/math.py
@@ -271,7 +271,7 @@ class Pose(Sequence):
if tensor_args is None and device is None:
log_error("Pose.to() requires tensor_args or device")
if tensor_args is not None:
- t_type = vars(tensor_args.as_torch_dict())
+ t_type = tensor_args.as_torch_dict()
else:
t_type = {"device": device}
if self.position is not None:
diff --git a/src/curobo/types/robot.py b/src/curobo/types/robot.py
index 61d1ede..8095430 100644
--- a/src/curobo/types/robot.py
+++ b/src/curobo/types/robot.py
@@ -43,6 +43,8 @@ class RobotConfig:
@staticmethod
def from_dict(data_dict_in, tensor_args=TensorDeviceType()):
+ if "robot_cfg" in data_dict_in:
+ data_dict_in = data_dict_in["robot_cfg"]
data_dict = deepcopy(data_dict_in)
if isinstance(data_dict["kinematics"], dict):
data_dict["kinematics"] = CudaRobotModelConfig.from_config(
diff --git a/src/curobo/util/logger.py b/src/curobo/util/logger.py
index 9b79a45..917de79 100644
--- a/src/curobo/util/logger.py
+++ b/src/curobo/util/logger.py
@@ -8,11 +8,26 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
+"""
+This module provides logging API, wrapping :py:class:`logging.Logger`. These functions are used
+to log messages in the cuRobo package. The functions can also be used in other packages by
+creating a new logger (:py:meth:`setup_logger`) with the desired name.
+"""
# Standard Library
import logging
+import sys
-def setup_curobo_logger(level="info"):
+def setup_logger(level="info", logger_name: str = "curobo"):
+ """Set up logger level.
+
+ Args:
+ level: Log level. Default is "info". Other options are "debug", "warning", "error".
+ logger_name: Name of the logger. Default is "curobo".
+
+ Raises:
+ ValueError: If log level is not one of [info, debug, warning, error].
+ """
FORMAT = "[%(levelname)s] [%(name)s] %(message)s"
if level == "info":
level = logging.INFO
@@ -25,21 +40,64 @@ def setup_curobo_logger(level="info"):
else:
raise ValueError("Log level should be one of [info,debug, warn, error]")
logging.basicConfig(format=FORMAT, level=level)
- logger = logging.getLogger("curobo")
+ logger = logging.getLogger(logger_name)
logger.setLevel(level=level)
-def log_warn(txt: str, *args, **kwargs):
- logger = logging.getLogger("curobo")
+def setup_curobo_logger(level="info"):
+ """Set up logger level for curobo package. Deprecated. Use :py:meth:`setup_logger` instead."""
+ return setup_logger(level, "curobo")
+
+
+def log_warn(txt: str, logger_name: str = "curobo", *args, **kwargs):
+ """Log warning message. Also see :py:meth:`logging.Logger.warning`.
+
+ Args:
+ txt: Warning message.
+ logger_name: Name of the logger. Default is "curobo".
+ """
+ logger = logging.getLogger(logger_name)
logger.warning(txt, *args, **kwargs)
-def log_info(txt: str, *args, **kwargs):
- logger = logging.getLogger("curobo")
+def log_info(txt: str, logger_name: str = "curobo", *args, **kwargs):
+ """Log info message. Also see :py:meth:`logging.Logger.info`.
+
+ Args:
+ txt: Info message.
+ logger_name: Name of the logger. Default is "curobo".
+ """
+ logger = logging.getLogger(logger_name)
logger.info(txt, *args, **kwargs)
-def log_error(txt: str, exc_info=True, stack_info=True, *args, **kwargs):
- logger = logging.getLogger("curobo")
- logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs)
- raise
+def log_error(
+ txt: str,
+ logger_name: str = "curobo",
+ exc_info=True,
+ stack_info=False,
+ stacklevel: int = 2,
+ *args,
+ **kwargs
+):
+ """Log error and raise ValueError.
+
+ Args:
+ txt: Helpful message that conveys the error.
+ logger_name: Name of the logger. Default is "curobo".
+ exc_info: Add exception info to message. See :py:meth:`logging.Logger.error`.
+ stack_info: Add stacktracke to message. See :py:meth:`logging.Logger.error`.
+ stacklevel: See :py:meth:`logging.Logger.error`. Default value of 2 removes this function
+ from the stack trace.
+
+ Raises:
+ ValueError: Error message with exception.
+ """
+ logger = logging.getLogger(logger_name)
+ if sys.version_info.major == 3 and sys.version_info.minor <= 7:
+ logger.error(txt, exc_info=exc_info, stack_info=stack_info, *args, **kwargs)
+ else:
+ logger.error(
+ txt, exc_info=exc_info, stack_info=stack_info, stacklevel=stacklevel, *args, **kwargs
+ )
+ raise ValueError(txt)
diff --git a/src/curobo/util/metrics.py b/src/curobo/util/metrics.py
index 208eaf1..9c6a90c 100644
--- a/src/curobo/util/metrics.py
+++ b/src/curobo/util/metrics.py
@@ -38,6 +38,7 @@ class CuroboMetrics(TrajectoryMetrics):
perception_success: bool = False
perception_interpolated_success: bool = False
jerk: float = np.inf
+ perception_time: float = 0.0
@dataclass
@@ -47,6 +48,7 @@ class CuroboGroupMetrics(TrajectoryGroupMetrics):
perception_success: float = 0.0
perception_interpolated_success: float = 0.0
jerk: float = np.inf
+ perception_time: float = np.inf
@classmethod
def from_list(cls, group: List[CuroboMetrics]):
@@ -60,5 +62,6 @@ class CuroboGroupMetrics(TrajectoryGroupMetrics):
[m.perception_interpolated_success for m in group]
)
data.jerk = Statistic.from_list([m.jerk for m in successes])
+ data.perception_time = Statistic.from_list([m.perception_time for m in successes])
return data
diff --git a/src/curobo/util/trajectory.py b/src/curobo/util/trajectory.py
index effc1d0..cb8f845 100644
--- a/src/curobo/util/trajectory.py
+++ b/src/curobo/util/trajectory.py
@@ -11,7 +11,7 @@
# Standard Library
import math
from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Tuple
# Third Party
import numpy as np
@@ -121,44 +121,58 @@ def get_spline_interpolated_trajectory(raw_traj, des_horizon, degree=5):
def get_batch_interpolated_trajectory(
raw_traj: JointState,
+ raw_dt: torch.Tensor,
interpolation_dt: float,
- max_vel: torch.Tensor,
- max_acc: torch.Tensor,
- max_jerk: torch.Tensor,
- raw_dt: float,
+ max_vel: Optional[torch.Tensor] = None,
+ max_acc: Optional[torch.Tensor] = None,
+ max_jerk: Optional[torch.Tensor] = None,
kind: InterpolateType = InterpolateType.LINEAR_CUDA,
out_traj_state: Optional[JointState] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
max_deviation: float = 0.1,
min_dt: float = 0.02,
+ max_dt: float = 0.15,
optimize_dt: bool = True,
):
# compute dt across trajectory:
+ if len(raw_traj.shape) == 2:
+ raw_traj = raw_traj.unsqueeze(0)
b, horizon, dof = raw_traj.position.shape # horizon
# given the dt required to run trajectory at maximum velocity,
# we find the number of timesteps required:
- traj_steps, steps_max, opt_dt = calculate_tsteps(
- raw_traj.velocity,
- raw_traj.acceleration,
- raw_traj.jerk,
- interpolation_dt,
- max_vel,
- max_acc,
- max_jerk,
- raw_dt,
- min_dt,
- horizon,
- optimize_dt,
- )
+ if optimize_dt:
+ traj_steps, steps_max, opt_dt = calculate_tsteps(
+ raw_traj.velocity,
+ raw_traj.acceleration,
+ raw_traj.jerk,
+ interpolation_dt,
+ max_vel,
+ max_acc,
+ max_jerk,
+ raw_dt,
+ min_dt,
+ max_dt,
+ horizon,
+ optimize_dt,
+ )
+ else:
+ traj_steps, steps_max = calculate_traj_steps(raw_dt, interpolation_dt, horizon)
+ opt_dt = raw_dt
# traj_steps contains the tsteps for each trajectory
- assert steps_max > 0
- # To do linear interpolation, we
+ if steps_max <= 0:
+ log_error("Steps max is less than 0")
+
+ if out_traj_state is not None and out_traj_state.position.shape[1] < steps_max:
+ log_warn(
+ "Interpolation buffer shape is smaller than steps_max,"
+ + " creating new buffer of shape "
+ + str(steps_max)
+ )
+ out_traj_state = None
+
if out_traj_state is None:
out_traj_state = JointState.zeros([b, steps_max, dof], tensor_args)
- if out_traj_state.position.shape[1] < steps_max:
- log_error("Interpolation buffer shape is smaller than steps_max")
-
if kind in [InterpolateType.LINEAR, InterpolateType.CUBIC]:
# plot and save:
out_traj_state = get_cpu_linear_interpolation(
@@ -187,7 +201,7 @@ def get_batch_interpolated_trajectory(
)
else:
raise ValueError("Unknown interpolation type")
- # opt_dt = (raw_dt) / opt_dt
+
return out_traj_state, traj_steps, opt_dt
@@ -448,7 +462,9 @@ def calculate_dt_fixed(
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
raw_dt: torch.Tensor,
- interpolation_dt: float,
+ min_dt: float,
+ max_dt: float,
+ epsilon: float = 1e-6,
):
# compute scaled dt:
max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof
@@ -465,9 +481,9 @@ def calculate_dt_fixed(
dt_score_jerk = raw_dt * torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3)
dt_score = torch.maximum(dt_score_vel, dt_score_acc)
dt_score = torch.maximum(dt_score, dt_score_jerk)
- dt_score = torch.clamp(dt_score, interpolation_dt, raw_dt)
- # NOTE: this dt score is not dt, rather a scaling to convert velocity, acc, jerk that was
- # computed with raw_dt to a new dt
+
+ dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, max_dt)
+
return dt_score
@@ -480,7 +496,8 @@ def calculate_dt(
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
raw_dt: float,
- interpolation_dt: float,
+ min_dt: float,
+ epsilon: float = 1e-6,
):
# compute scaled dt:
max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof
@@ -497,7 +514,8 @@ def calculate_dt(
dt_score_jerk = raw_dt * torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3)
dt_score = torch.maximum(dt_score_vel, dt_score_acc)
dt_score = torch.maximum(dt_score, dt_score_jerk)
- dt_score = torch.clamp(dt_score, interpolation_dt, raw_dt)
+ dt_score = torch.clamp(dt_score * (1.0 + epsilon), min_dt, raw_dt)
+
# NOTE: this dt score is not dt, rather a scaling to convert velocity, acc, jerk that was
# computed with raw_dt to a new dt
return dt_score
@@ -511,6 +529,7 @@ def calculate_dt_no_clamp(
max_vel: torch.Tensor,
max_acc: torch.Tensor,
max_jerk: torch.Tensor,
+ epsilon: float = 1e-6,
):
# compute scaled dt:
max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof
@@ -527,9 +546,19 @@ def calculate_dt_no_clamp(
dt_score_jerk = torch.pow((torch.max(jerk_scale_dt, dim=-1)[0]), 1 / 3)
dt_score = torch.maximum(dt_score_vel, dt_score_acc)
dt_score = torch.maximum(dt_score, dt_score_jerk)
+ dt_score = dt_score * (1.0 + epsilon)
return dt_score
+@get_torch_jit_decorator()
+def calculate_traj_steps(
+ opt_dt: torch.Tensor, interpolation_dt: float, horizon: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ traj_steps = (torch.ceil((horizon - 1) * ((opt_dt) / interpolation_dt))).to(dtype=torch.int32)
+ steps_max = torch.max(traj_steps)
+ return traj_steps, steps_max
+
+
@get_torch_jit_decorator()
def calculate_tsteps(
vel: torch.Tensor,
@@ -541,15 +570,23 @@ def calculate_tsteps(
max_jerk: torch.Tensor,
raw_dt: torch.Tensor,
min_dt: float,
+ max_dt: float,
horizon: int,
optimize_dt: bool = True,
):
# compute scaled dt:
opt_dt = calculate_dt_fixed(
- vel, acc, jerk, max_vel, max_acc, max_jerk, raw_dt, interpolation_dt
+ vel,
+ acc,
+ jerk,
+ max_vel,
+ max_acc,
+ max_jerk,
+ raw_dt,
+ min_dt,
+ max_dt,
)
if not optimize_dt:
opt_dt[:] = raw_dt
- traj_steps = (torch.ceil((horizon - 1) * ((opt_dt) / interpolation_dt))).to(dtype=torch.int32)
- steps_max = torch.max(traj_steps)
+ traj_steps, steps_max = calculate_traj_steps(opt_dt, interpolation_dt, horizon)
return traj_steps, steps_max, opt_dt
diff --git a/src/curobo/wrap/model/robot_segmenter.py b/src/curobo/wrap/model/robot_segmenter.py
index 173ff03..2948765 100644
--- a/src/curobo/wrap/model/robot_segmenter.py
+++ b/src/curobo/wrap/model/robot_segmenter.py
@@ -40,6 +40,7 @@ class RobotSegmenter:
distance_threshold: float = 0.05,
use_cuda_graph: bool = True,
ops_dtype: torch.dtype = torch.float32,
+ depth_to_meter: float = 0.001,
):
self._robot_world = robot_world
self._projection_rays = None
@@ -53,6 +54,7 @@ class RobotSegmenter:
self.tensor_args = robot_world.tensor_args
self.distance_threshold = distance_threshold
self._ops_dtype = ops_dtype
+ self._depth_to_meter = depth_to_meter
@staticmethod
def from_robot_file(
@@ -95,7 +97,10 @@ class RobotSegmenter:
if len(intrinsics.shape) == 2:
intrinsics = intrinsics.unsqueeze(0)
project_rays = get_projection_rays(
- camera_obs.depth_image.shape[-2], camera_obs.depth_image.shape[-1], intrinsics
+ camera_obs.depth_image.shape[-2],
+ camera_obs.depth_image.shape[-1],
+ intrinsics,
+ self._depth_to_meter,
).to(dtype=self._ops_dtype)
if self._projection_rays is None:
diff --git a/src/curobo/wrap/model/robot_world.py b/src/curobo/wrap/model/robot_world.py
index 5900712..5073eb0 100644
--- a/src/curobo/wrap/model/robot_world.py
+++ b/src/curobo/wrap/model/robot_world.py
@@ -18,8 +18,7 @@ from typing import Dict, List, Optional, Tuple, Union
import torch
# CuRobo
-from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
-from curobo.cuda_robot_model.types import CudaRobotModelState
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelState
from curobo.geom.sdf.utils import create_collision_checker
from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig
from curobo.geom.types import WorldConfig
diff --git a/src/curobo/wrap/reacher/evaluator.py b/src/curobo/wrap/reacher/evaluator.py
index 2dfc351..863b6a9 100644
--- a/src/curobo/wrap/reacher/evaluator.py
+++ b/src/curobo/wrap/reacher/evaluator.py
@@ -78,7 +78,7 @@ class TrajEvaluatorConfig:
max_jerk: float = 500.0,
cost_weight: float = 0.01,
min_dt: float = 0.001,
- max_dt: float = 0.1,
+ max_dt: float = 0.15,
tensor_args: TensorDeviceType = TensorDeviceType(),
) -> TrajEvaluatorConfig:
"""Creates TrajEvaluatorConfig object from basic parameters.
@@ -169,6 +169,7 @@ def compute_smoothness(
traj_dt: torch.Tensor,
min_dt: float,
max_dt: float,
+ epsilon: float = 1e-6,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""JIT compatible function to compute smoothness.
@@ -185,6 +186,8 @@ def compute_smoothness(
traj_dt: dt of trajectory tensor of shape (batch, horizon) or (1, horizon)
min_dt: minimum delta time allowed between steps/waypoints in a trajectory.
max_dt: maximum delta time allowed between steps/waypoints in a trajectory.
+ epsilon: relaxes evaluation of velocity, acceleration, and jerk limits to allow for
+ numerical errors.
Returns:
Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and
@@ -211,7 +214,7 @@ def compute_smoothness(
# clamp dt score:
- dt_score = torch.clamp(dt_score, min_dt, max_dt)
+ dt_score = torch.clamp(dt_score * (1 + epsilon), min_dt, max_dt)
scale_dt = (1 / dt_score).view(-1, 1, 1)
abs_acc = torch.abs(acc) * (scale_dt**2)
# mean_acc_val = torch.max(torch.mean(abs_acc, dim=-1), dim=-1)[0]
@@ -376,7 +379,9 @@ class TrajEvaluator(TrajEvaluatorConfig):
label, cost = compute_smoothness_opt_dt(
js.velocity, js.acceleration, js.jerk, max_vel, self.max_acc, self.max_jerk, opt_dt
)
+
label = torch.logical_and(label, opt_dt <= self.max_dt)
+
pl_cost = compute_path_length_cost(js.velocity, cspace_distance_weight)
return label, pl_cost + self.cost_weight * cost
diff --git a/src/curobo/wrap/reacher/ik_solver.py b/src/curobo/wrap/reacher/ik_solver.py
index 0e7120a..23ffd8e 100644
--- a/src/curobo/wrap/reacher/ik_solver.py
+++ b/src/curobo/wrap/reacher/ik_solver.py
@@ -9,9 +9,10 @@
# its affiliates is strictly prohibited.
#
-"""
-This module contains :meth:`IKSolver` to solve inverse kinematics, along with
-:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result.
+"""
+This module contains :meth:`IKSolver` to solve inverse kinematics, along with
+:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result. A minimal
+example to solve IK is available at :ref:`python_ik_example`.
.. raw:: html
@@ -20,6 +21,7 @@ This module contains :meth:`IKSolver` to solve inverse kinematics, along with
This demo is available in :ref:`isaac_ik_reachability`.
+
"""
from __future__ import annotations
@@ -190,9 +192,9 @@ class IKSolverConfig:
grad_iters: Number of iterations for gradient optimization.
use_particle_opt: Flag to enable particle optimization.
collision_checker_type: Type of collision checker to use for collision checking.
- sync_cuda_time: Flag to enable synchronization of cuda device with host before
- measuring compute time. If you set this to False, then measured time will not
- include completion of any launched CUDA kernels.
+ sync_cuda_time: Flag to enable synchronization of cuda device with host using
+ :py:func:`torch.cuda.synchronize` before measuring compute time. If you set this to
+ False, then measured time will not include completion of any launched CUDA kernels.
use_gradient_descent: Flag to enable gradient descent optimization instead of LBFGS.
collision_cache: Number of objects to cache for collision checking.
n_collision_envs: Number of collision environments to use for IK. This is useful when
@@ -1065,7 +1067,6 @@ class IKSolver(IKSolverConfig):
IKResult object with solutions to the IK problems.
"""
success = self._get_success(result.metrics, num_seeds=num_seeds)
-
if result.metrics.null_space_error is not None:
result.metrics.pose_error += result.metrics.null_space_error
if result.metrics.cspace_error is not None:
@@ -1524,6 +1525,36 @@ class IKSolver(IKSolverConfig):
if isinstance(rollout, ArmReacher)
]
+ def get_full_js(self, active_js: JointState) -> JointState:
+ """Get full joint state from controlled joint state, appending locked joints.
+
+ Args:
+ active_js: Controlled joint state
+
+ Returns:
+ JointState: Full joint state.
+ """
+ return self.rollout_fn.get_full_dof_from_solution(active_js)
+
+ def get_active_js(
+ self,
+ in_js: JointState,
+ ):
+ """Get controlled joint state from input joint state.
+
+ This is used to get the joint state for only joints that are optimization variables. This
+ also re-orders the joints to match the order of optimization variables.
+
+ Args:
+ in_js: Input joint state.
+
+ Returns:
+ JointState: Active joint state.
+ """
+ opt_jnames = self.rollout_fn.joint_names
+ opt_js = in_js.get_ordered_joint_state(opt_jnames)
+ return opt_js
+
@property
def joint_names(self) -> List[str]:
"""Get ordered names of all joints used in optimization with IKSolver."""
diff --git a/src/curobo/wrap/reacher/motion_gen.py b/src/curobo/wrap/reacher/motion_gen.py
index 1468a5c..2edf67c 100644
--- a/src/curobo/wrap/reacher/motion_gen.py
+++ b/src/curobo/wrap/reacher/motion_gen.py
@@ -9,20 +9,17 @@
# its affiliates is strictly prohibited.
#
-"""
-This module contains :meth:`MotionGen` class that provides a high-level interface for motion
-generation. Motion Generation can take goals either as joint configurations
+"""
+This module contains :meth:`MotionGen` class that provides a high-level interface for motion
+generation. Motion Generation can take goals either as joint configurations
:meth:`MotionGen.plan_single_js` or as Cartesian poses :meth:`MotionGen.plan_single`. When Cartesian
-pose is given as target, inverse kinematics is first done to generate seeds for trajectory
-optimization. Motion generation fallback to using a graph planner when linear interpolated
+pose is given as target, inverse kinematics is first done to generate seeds for trajectory
+optimization. Motion generation fallback to using a graph planner when linear interpolated
trajectory optimization seeds are not successful. Reaching one Cartesian pose in a goalset is also
supported using :meth:`MotionGen.plan_goalset`. Batched planning in the same world environment (
-:meth:`MotionGen.plan_batch`) and different world environments (:meth:`MotionGen.plan_batch_env`)
+:meth:`MotionGen.plan_batch`) and different world environments (:meth:`MotionGen.plan_batch_env`)
is also provided.
-A motion generation request can be configured using :meth:`MotionGenPlanConfig`. The result
-of motion generation is returned as a :meth:`MotionGenResult`.
-
.. raw:: html
@@ -30,6 +27,11 @@ of motion generation is returned as a :meth:`MotionGenResult`.
+
+A motion generation request can be configured using :meth:`MotionGenPlanConfig`. The result
+of motion generation is returned as a :meth:`MotionGenResult`. A minimal example is availble at
+:ref:`python_motion_gen_example`.
+
"""
from __future__ import annotations
@@ -64,10 +66,10 @@ from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
-from curobo.types.tensor import T_BDOF, T_BValue_bool, T_BValue_float
+from curobo.types.tensor import T_BDOF, T_DOF, T_BValue_float
from curobo.util.logger import log_error, log_info, log_warn
from curobo.util.tensor_util import tensor_repeat_seeds
-from curobo.util.trajectory import InterpolateType
+from curobo.util.trajectory import InterpolateType, get_batch_interpolated_trajectory
from curobo.util.warp import init_warp
from curobo.util_file import (
get_robot_configs_path,
@@ -78,7 +80,7 @@ from curobo.util_file import (
)
from curobo.wrap.reacher.evaluator import TrajEvaluator, TrajEvaluatorConfig
from curobo.wrap.reacher.ik_solver import IKResult, IKSolver, IKSolverConfig
-from curobo.wrap.reacher.trajopt import TrajOptSolver, TrajOptSolverConfig
+from curobo.wrap.reacher.trajopt import TrajOptResult, TrajOptSolver, TrajOptSolverConfig
from curobo.wrap.reacher.types import ReacherSolveState, ReacherSolveType
@@ -159,14 +161,14 @@ class MotionGenConfig:
#: tried again until :attr:`MotionGenPlanConfig.finetune_attempts`.
optimize_dt: bool = True
- @staticmethod
@profiler.record_function("motion_gen_config/load_from_robot_config")
+ @staticmethod
def load_from_robot_config(
robot_cfg: Union[Union[str, Dict], RobotConfig],
world_model: Optional[Union[Union[str, Dict], WorldConfig]] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
num_ik_seeds: int = 32,
- num_graph_seeds: int = 1,
+ num_graph_seeds: int = 4,
num_trajopt_seeds: int = 4,
num_batch_ik_seeds: int = 32,
num_batch_trajopt_seeds: int = 1,
@@ -228,8 +230,9 @@ class MotionGenConfig:
finetune_smooth_weight: Optional[List[float]] = None,
state_finite_difference_mode: Optional[str] = None,
finetune_dt_scale: float = 0.9,
+ minimum_trajectory_dt: Optional[float] = None,
maximum_trajectory_time: Optional[float] = None,
- maximum_trajectory_dt: float = 0.1,
+ maximum_trajectory_dt: Optional[float] = None,
velocity_scale: Optional[Union[List[float], float]] = None,
acceleration_scale: Optional[Union[List[float], float]] = None,
jerk_scale: Optional[Union[List[float], float]] = None,
@@ -239,13 +242,249 @@ class MotionGenConfig:
graph_seed: int = 1531,
high_precision: bool = False,
):
+ """Create a motion generation configuration from robot and world configuration.
+
+ Args:
+ robot_cfg: Robot configuration to use for motion generation. This can be a path to a
+ yaml file, a dictionary, or an instance of :class:`RobotConfig`. See
+ :ref:`available_robot_list` for a list of available robots. You can also create a
+ a configuration file for your robot using :ref:`tut_robot_configuration`.
+ world_model: World configuration to use for motion generation. This can be a path to a
+ yaml file, a dictionary, or an instance of :class:`WorldConfig`. See
+ :ref:`world_collision` for more details.
+ tensor_args: Numerical precision and compute device to use for motion generation.
+ num_ik_seeds: Number of seeds to use for solving inverse kinematics. Default of 32 is
+ found to be a good number for most cases. In sparse environments, a lower number of
+ 16 can also be used.
+ num_graph_seeds: Number of seeds to use for graph planner per problem query. When graph
+ planning is used to generate seeds for trajectory optimization, graph planner will
+ attempt to find collision-free paths from the start state to the many inverse
+ kinematics solutions.
+ num_trajopt_seeds: Number of seeds to use for trajectory optimization per problem
+ query. Default of 4 is found to be a good number for most cases. Increasing this
+ will increase memory usage.
+ num_batch_ik_seeds: Number of seeds to use for inverse kinematics during batched
+ planning. Default of 32 is found to be a good number for most cases.
+ num_batch_trajopt_seeds: Number of seeds to use for trajectory optimization during
+ batched planning. Using more than 1 will disable graph planning for batched
+ planning.
+ num_trajopt_noisy_seeds: Number of augmented trajectories to use per trajectory seed.
+ The augmentation is done by adding random noise to the trajectory. This
+ augmentation has not been found to be useful and it's recommended to keep this to
+ 1. The noisy seeds can also be used in conjunction with the trajopt_seed_ratio to
+ generate seeds that go through a bias point.
+ position_threshold: Position threshold in meters between reached position and target
+ position used to measure success.
+ rotation_threshold: Rotation threshold between reached orientation and target
+ orientation used to measure success. The metric is q^T * q, where q is the
+ quaternion difference between target and reached orientation. The metric is not
+ easy to interpret and a future release will provide a more intuitive metric. For
+ now, use 0.05 as a good value.
+ cspace_threshold: Joint space threshold in radians for revolute joints and meters for
+ linear joints between reached joint configuration and target joint configuration
+ used to measure success. Default of 0.05 has been found to be a good value for most
+ cases.
+ world_coll_checker: Instance of world collision checker to use for motion generation.
+ Leaving this to None will create a new instance of world collision checker using
+ the provided attr:`world_model`.
+ base_cfg_file: Base configuration file containing convergence and constraint criteria
+ to measure success.
+ particle_ik_file: Optimizer configuration file to use for particle-based optimization
+ during inverse kinematics.
+ gradient_ik_file: Optimizer configuration file to use for gradient-based optimization
+ during inverse kinematics.
+ graph_file: Configuration file to use for graph planner.
+ particle_trajopt_file: Optimizer configuration file to use for particle-based
+ optimization during trajectory optimization.
+ gradient_trajopt_file: Optimizer configuration file to use for gradient-based
+ optimization during trajectory optimization.
+ finetune_trajopt_file: Optimizer configuration file to use for finetuning trajectory
+ optimization.
+ trajopt_tsteps: Number of waypoints to use for trajectory optimization. Default of 32
+ is found to be a good number for most cases.
+ interpolation_steps: Buffer size to use for storing interpolated trajectory. Default of
+ 5000 is found to be a good number for most cases.
+ interpolation_dt: Time step in seconds to use for generating interpolated trajectory
+ from optimized trajectory. Change this if you want to generate a trajectory with
+ a fixed timestep between waypoints.
+ interpolation_type: Interpolation type to use for generating dense waypoints from
+ optimized trajectory. Default of
+ :py:attr:`curobo.util.trajectory.InterpolateType.LINEAR_CUDA` is found to be a
+ good choice for most cases. Other suitable options for real robot execution are
+ :py:attr:`curobo.util.trajectory.InterpolateType.QUINTIC` and
+ :py:attr:`curobo.util.trajectory.InterpolateType.CUBIC`.
+ use_cuda_graph: Record compute ops as cuda graphs and replay recorded graphs where
+ implemented. This can speed up execution by upto 10x. Default of True is
+ recommended. Enabling this will prevent changing problem type or batch size
+ after the first call to the solver.
+ self_collision_check: Enable self collision checks for generated motions. Default of
+ True is recommended. Set this to False to debug planning failures. Setting this to
+ False will also set self_collision_opt to False.
+ self_collision_opt: Enable self collision cost during optimization (IK, TrajOpt).
+ Default of True is recommended.
+ grad_trajopt_iters: Number of iterations to run trajectory optimization.
+ trajopt_seed_ratio: Ratio of linear and bias seeds to use for trajectory optimization.
+ Linear seed will generate a linear interpolated trajectory from start state
+ to IK solutions. Bias seed will add a mid-waypoint through the retract
+ configuration. Default of 1.0 linear and 0.0 bias is recommended. This can be
+ changed to 0.5 linear and 0.5 bias, along with changing trajopt_noisy_seeds to 2.
+ ik_opt_iters: Number of iterations to run inverse kinematics.
+ ik_particle_opt: Enable particle-based optimization during inverse kinematics. Default
+ of True is recommended as particle-based optimization moves the random seeds to
+ a regions of local minima.
+ collision_checker_type: Type of collision checker to use for motion generation. Default
+ of CollisionCheckerType.MESH supports world represented by Cuboids and Meshes. See
+ :ref:`world_collision` for more details.
+ sync_cuda_time: Synchronize with host using :py:func:`torch.cuda.synchronize` before
+ measuring compute time.
+ trajopt_particle_opt: Enable particle-based optimization during trajectory
+ optimization. Default of True is recommended as particle-based optimization moves
+ the interpolated seeds away from bad local minima.
+ traj_evaluator_config: Configuration for trajectory evaluator. Default of None will
+ create a new instance of TrajEvaluatorConfig. After trajectory optimization across
+ many seeds, the best trajectory is selected based on this configuration. This
+ evaluator also checks if the optimized dt is within
+ :py:attr:`curobo.wrap.reacher.evaluator.TrajEvaluatorConfig.max_dt`. This check is
+ needed to measure smoothness of the optimized trajectory as bad trajectories can
+ have very high dt to fit within velocity, acceleration, and jerk limits.
+ traj_evaluator: Instance of trajectory evaluator to use for trajectory optimization.
+ Default of None will create a new instance of TrajEvaluator. In case you want to
+ use a custom evaluator, you can create a child instance of TrajEvaluator and
+ pass it.
+ minimize_jerk: Minimize jerk as regularization during trajectory optimizaiton.
+ filter_robot_command: Filter generated trajectory to remove finite difference
+ artifacts. Default of True is recommended.
+ use_gradient_descent: Use gradient descent instead of L-BFGS for trajectory
+ optimization.
+ collision_cache: Cache of obstacles to create to load obstacles between planning calls.
+ An example: ``{"obb": 10, "mesh": 10}``, to create a cache of 10 cuboids and 10
+ meshes.
+ n_collision_envs: Number of collision environments to create for batched planning
+ across different environments. Only used for :py:meth:`MotionGen.plan_batch_env`
+ and :py:meth:`MotionGen.plan_batch_env_goalset`.
+ ee_link_name: End effector link/frame to use for reaching Cartesian poses. Default of
+ None will use the end effector link from the robot configuration. This cannot
+ be changed after creating the robot configuration.
+ use_es_ik: Use evolutionary strategy for as the particle-based optimizer for inverse
+ kinematics. Default of None will use MPPI as the optimization algorithm. ES is not
+ recommended as it's unstable and provided optimization parameters were not tuned.
+ use_es_trajopt: Use evolutionary strategy as the particle-based optimizer for
+ trajectory optimization. Default of None will use MPPI as the optimization
+ algorithm. ES is not recommended as it's unstable and provided optimization
+ parameters were not tuned.
+ es_ik_learning_rate: Learning rate to use for evolutionary strategy in IK.
+ es_trajopt_learning_rate: Learning rate to use for evolutionary strategy in TrajOpt.
+ use_ik_fixed_samples: Use fixed samples of noise during particle-based optimization
+ of IK. Default of None will use the setting from the optimizer configuration file
+ (``particle_ik.yml``).
+ use_trajopt_fixed_samples: Use fixed samples of noise during particle-based
+ optimization of trajectory optimization. Default of None will use the setting from
+ the optimizer configuration file (``particle_trajopt.yml``).
+ evaluate_interpolated_trajectory: Evaluate interpolated trajectory after optimization.
+ Default of True is recommended to ensure the optimized trajectory is not passing
+ through very thin obstacles.
+ partial_ik_iters: Number of iterations of L-BFGS to run inverse kinematics when
+ only partial IK is needed.
+ fixed_iters_trajopt: Use fixed number of iterations of L-BFGS for trajectory
+ optimization. Default of None will use the setting from the optimizer
+ configuration. In most cases, fixed iterations of solvers are run as current
+ solvers treat constraints as costs and there is no guarantee that the constraints
+ will be satisfied. Instead of checking constraints between iterations of a solver
+ and exiting, it's computationally cheaper to run a fixed number of iterations. In
+ addition, running fixed iterations of solvers is more robust to outlier problems.
+ store_ik_debug: Store debugging information such as values of optimization variables
+ at every iteration in IK result. Setting this to True will set
+ :attr:`use_cuda_graph` to False.
+ store_trajopt_debug: Store debugging information such as values of optimization
+ variables in TrajOpt result. Setting this to True will set :attr:`use_cuda_graph`
+ to False.
+ graph_trajopt_iters: Number of iterations to run trajectory optimization when seeded
+ from a graph plan. Default of None will use the same number of iterations as
+ linear seeded trajectory optimization. This can be set to a higher value of 200
+ in case where trajectories obtained are not of requird quality.
+ collision_max_outside_distance: Maximum distance to check for collision outside a
+ obstacle. Increasing this value will slow down collision checks with Meshes as
+ closest point queries will be run up to this distance outside an obstacle.
+ collision_activation_distance: Distance in meters to activate collision cost. A good
+ value to start with is 0.01 meters. Increase the distance if the robot needs to
+ stay further away from obstacles.
+ trajopt_dt: Time step in seconds to use for trajectory optimization. A good value to
+ start with is 0.15 seconds. This value is used to compute velocity, acceleration,
+ and jerk values for waypoints through finite difference.
+ js_trajopt_dt: Time step in seconds to use for trajectory optimization when reaching
+ joint space targets. A good value to start with is 0.15 seconds. This value is used
+ to compute velocity, acceleration, and jerk values for waypoints through finite
+ difference.
+ js_trajopt_tsteps: Number of waypoints to use for trajectory optimization when reaching
+ joint space targets. Default of None will use the same number of waypoints as
+ Cartesian trajectory optimization.
+ trim_steps: Trim waypoints from optimized trajectory. The optimized trajectory will
+ contain the start state at index 0 and have the last two waypoints be the same
+ as T-2 as trajectory optimization implicitly optimizes for zero acceleration and
+ velocity at the last waypoint. An example: ``[1,-2]`` will trim the first waypoint
+ and last 3 waypoints from the optimized trajectory.
+ store_debug_in_result: Store debugging information in MotionGenResult. This value is
+ set to True if either store_ik_debug or store_trajopt_debug is set to True.
+ finetune_trajopt_iters: Number of iterations to run trajectory optimization for
+ finetuning after an initial collision-free trajectory is obtained.
+ smooth_weight: Override smooth weight for trajectory optimization. It's not recommended
+ to set this value for most cases.
+ finetune_smooth_weight: Override smooth weight for finetuning trajectory optimization.
+ It's not recommended to set this value for most cases.
+ state_finite_difference_mode: Finite difference mode to use for computing velocity,
+ acceleration, and jerk values. Default of None will use the setting from the
+ optimizer configuration file. The default finite difference method is a five
+ point stencil to compute the derivatives as this is accurate and provides
+ faster convergence compared to backward or central difference methods.
+ finetune_dt_scale: Scale initial estimated dt by this value to finetune trajectory
+ optimization. This is deprecated and will be removed in future releases. Use
+ :py:attr:`MotionGenPlanConfig.finetune_dt_scale` instead.
+ minimum_trajectory_dt: Minimum time step in seconds allowed for trajectory
+ optimization.
+ maximum_trajectory_time: Maximum time in seconds allowed for trajectory optimization.
+ maximum_trajectory_dt: Maximum time step in seconds allowed for trajectory
+ optimization.
+ velocity_scale: Scale velocity limits by this value. Default of None will not scale
+ the velocity limits. To generate slower trajectories, use
+ :py:attr:`MotionGenPlanConfig.time_dilation_factor` < 1.0 instead. Changing this
+ value is not recommended as it changes the scale of cost terms and they would
+ require retuning.
+ acceleration_scale: Scale acceleration limits by this value. Default of None will not
+ scale the acceleration limits. To generate slower trajectories, use
+ :py:attr:`MotionGenPlanConfig.time_dilation_factor` < 1.0 instead. Changing this
+ value is not recommended as it changes the scale of cost terms and they would
+ require retuning.
+ jerk_scale: Scale jerk limits by this value. Default of None will not scale the jerk
+ limits. To generate slower trajectories, use
+ :py:attr:`MotionGenPlanConfig.time_dilation_factor` < 1.0 instead. Changing this
+ value is not recommended as it changes the scale of cost terms and they would
+ require retuning.
+ optimize_dt: Optimize dt during trajectory optimization. Default of True is
+ recommended to find time-optimal trajectories. Setting this to False will use the
+ provided :attr:`trajopt_dt` for trajectory optimization. Setting to False is
+ required when optimizing from a non-static start state.
+ project_pose_to_goal_frame: Project pose to goal frame when calculating distance
+ between reached and goal pose. Use this to constrain motion to specific axes
+ either in the global frame or the goal frame.
+ ik_seed: Random seed to use for inverse kinematics.
+ graph_seed: Random seed to use for graph planner.
+ high_precision: Use high precision settings for motion generation. This will increase
+ the number of iterations for optimization solvers and reduce the thresholds for
+ position to 1mm and rotation to 0.025. Default of False is recommended for most
+ cases as standard motion generation settings reach within 0.5mm on most problems.
+
+ Returns:
+ MotionGenConfig: Instance of motion generation configuration.
+ """
if position_threshold <= 0.001:
high_precision = True
if high_precision:
finetune_trajopt_iters = (
300 if finetune_trajopt_iters is None else max(300, finetune_trajopt_iters)
)
- grad_trajopt_iters = 200 if grad_trajopt_iters is None else max(200, grad_trajopt_iters)
+ if grad_trajopt_iters is None:
+ grad_trajopt_iters = 200
+ grad_trajopt_iters = max(200, grad_trajopt_iters)
position_threshold = min(position_threshold, 0.001)
rotation_threshold = min(rotation_threshold, 0.025)
cspace_threshold = min(cspace_threshold, 0.01)
@@ -256,28 +495,55 @@ class MotionGenConfig:
if trajopt_tsteps is not None:
js_trajopt_tsteps = trajopt_tsteps
if velocity_scale is not None and isinstance(velocity_scale, float):
+ log_warn(
+ "To slow down trajectories, use MotionGenPlanConfig.time_dilation_factor"
+ + " instead of velocity_scale"
+ )
velocity_scale = [velocity_scale]
if acceleration_scale is not None and isinstance(acceleration_scale, float):
+ log_warn(
+ "To slow down trajectories, use MotionGenPlanConfig.time_dilation_factor"
+ + " instead of acceleration_scale"
+ )
acceleration_scale = [acceleration_scale]
if jerk_scale is not None and isinstance(jerk_scale, float):
jerk_scale = [jerk_scale]
if store_ik_debug or store_trajopt_debug:
store_debug_in_result = True
+
+ if (
+ velocity_scale is not None
+ and min(velocity_scale) < 0.1
+ and finetune_trajopt_file is None
+ and maximum_trajectory_dt is None
+ ):
+ log_error(
+ "velocity scale<0.1 requires a user determined maximum_trajectory_dt as"
+ + " default scaling will likely fail. A good value to start with would be 30"
+ + " seconds"
+ )
+
+ if maximum_trajectory_dt is None:
+ maximum_trajectory_dt = 0.15
+ maximum_trajectory_dt_acc = maximum_trajectory_dt
+ maximum_trajectory_dt_vel = maximum_trajectory_dt
if (
acceleration_scale is not None
and min(acceleration_scale) < 1.0
- and maximum_trajectory_dt <= 0.1
+ and maximum_trajectory_dt <= 0.2
):
- maximum_trajectory_dt = np.sqrt(1.0 / min(acceleration_scale)) * maximum_trajectory_dt
- elif (
+ maximum_trajectory_dt_acc = (
+ np.sqrt(1.0 / min(acceleration_scale)) * maximum_trajectory_dt * 3
+ )
+ if (
velocity_scale is not None
and min(velocity_scale) < 1.0
- and maximum_trajectory_dt <= 0.1
+ and maximum_trajectory_dt <= 0.2
):
- maximum_trajectory_dt = (1.0 / min(velocity_scale)) * maximum_trajectory_dt
-
+ maximum_trajectory_dt_vel = (1.0 / min(velocity_scale)) * maximum_trajectory_dt * 3
+ maximum_trajectory_dt = max(maximum_trajectory_dt_acc, maximum_trajectory_dt_vel)
if maximum_trajectory_dt is not None:
if trajopt_dt is None:
trajopt_dt = maximum_trajectory_dt
@@ -288,16 +554,6 @@ class MotionGenConfig:
if velocity_scale is not None and min(velocity_scale) < 0.5:
fixed_iters_trajopt = True
- if (
- velocity_scale is not None
- and min(velocity_scale) < 0.1
- and finetune_trajopt_file is None
- ):
- log_error(
- "velocity scale<0.1 is not supported with default finetune_trajopt.yml "
- + "provide your own finetune_trajopt_file to override this error"
- )
-
if (
velocity_scale is not None
and min(velocity_scale) <= 0.25
@@ -338,7 +594,10 @@ class MotionGenConfig:
robot_cfg["kinematics"]["cspace"]["velocity_scale"] = velocity_scale
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
-
+ if minimum_trajectory_dt is None:
+ minimum_trajectory_dt = interpolation_dt
+ elif minimum_trajectory_dt < interpolation_dt:
+ log_error("minimum_trajectory_dt cannot be lower than interpolation_dt")
if traj_evaluator_config is None:
if maximum_trajectory_dt is not None:
max_dt = maximum_trajectory_dt
@@ -347,7 +606,7 @@ class MotionGenConfig:
if acceleration_scale is not None:
max_dt = max_dt * (1.0 / np.sqrt(min(acceleration_scale)))
traj_evaluator_config = TrajEvaluatorConfig.from_basic(
- min_dt=interpolation_dt, max_dt=max_dt, dof=robot_cfg.kinematics.dof
+ min_dt=minimum_trajectory_dt, max_dt=max_dt, dof=robot_cfg.kinematics.dof
)
traj_evaluator_config.max_acc = robot_cfg.kinematics.get_joint_limits().acceleration[1]
@@ -362,7 +621,8 @@ class MotionGenConfig:
if n_collision_envs is not None:
base_config_data["world_collision_checker_cfg"]["n_envs"] = n_collision_envs
if collision_max_outside_distance is not None:
- assert collision_max_outside_distance >= 0.0
+ if collision_max_outside_distance < 0.0:
+ log_error("collision_max_outside_distance cannot be negative")
base_config_data["world_collision_checker_cfg"][
"max_distance"
] = collision_max_outside_distance
@@ -394,7 +654,6 @@ class MotionGenConfig:
grad_iters=ik_opt_iters,
use_particle_opt=ik_particle_opt,
sync_cuda_time=sync_cuda_time,
- # use_gradient_descent=use_gradient_descent,
use_es=use_es_ik,
es_learning_rate=es_ik_learning_rate,
use_fixed_samples=use_ik_fixed_samples,
@@ -529,7 +788,7 @@ class MotionGenConfig:
use_particle_opt=False,
traj_evaluator_config=traj_evaluator_config,
traj_evaluator=traj_evaluator,
- evaluate_interpolated_trajectory=True,
+ evaluate_interpolated_trajectory=evaluate_interpolated_trajectory,
fixed_iters=fixed_iters_trajopt,
store_debug=store_trajopt_debug,
collision_activation_distance=collision_activation_distance,
@@ -653,7 +912,7 @@ class MotionGenPlanConfig:
#: Run finetuning trajectory optimization after running 100 iterations of
#: trajectory optimization. This will provide shorter and smoother trajectories. When
- # :attr:`MotionGenConfig.optimize_dt` is True, this flag will also scale the trajectory
+ #: :attr:`MotionGenConfig.optimize_dt` is True, this flag will also scale the trajectory
#: optimization by a new dt. Leave this to True for most cases. If you are not interested in
#: finding time-optimal solutions and only want to use motion generation as a feasibility check,
#: set this to False. Note that when set to False, the resulting trajectory is only guaranteed
@@ -681,11 +940,21 @@ class MotionGenPlanConfig:
#: fails to find a solution. This is only used when :attr:`MotionGenConfig.optimize_dt` is True.
finetune_dt_decay: float = 1.025
+ #: Slow down optimized trajectory by re-timing with a dilation factor. This is useful to
+ #: execute trajectories at a slower speed for debugging. Use this to generate slower
+ #: trajectories instead of reducing :attr:`MotionGenConfig.velocity_scale` or
+ #: :attr:`MotionGenConfig.acceleration_scale` as those parameters will require re-tuning
+ #: of the cost terms while :attr:`MotionGenPlanConfig.time_dilation_factor` will only
+ #: post-process the trajectory.
+ time_dilation_factor: Optional[float] = None
+
def __post_init__(self):
+ """Post initialization checks."""
if not self.enable_opt and not self.enable_graph:
log_error("Graph search and Optimization are both disabled, enable one")
def clone(self) -> MotionGenPlanConfig:
+ """Clone the current planning configuration."""
return MotionGenPlanConfig(
enable_graph=self.enable_graph,
enable_opt=self.enable_opt,
@@ -710,6 +979,7 @@ class MotionGenPlanConfig:
),
finetune_dt_scale=self.finetune_dt_scale,
finetune_attempts=self.finetune_attempts,
+ time_dilation_factor=self.time_dilation_factor,
)
@@ -814,6 +1084,7 @@ class MotionGenResult:
goalset_index: Optional[torch.Tensor] = None
def clone(self):
+ """Clone the current result."""
m = MotionGenResult(
self.success.clone(),
valid_query=self.valid_query,
@@ -842,29 +1113,16 @@ class MotionGenResult:
)
return m
- @staticmethod
- def _check_none_and_copy_idx(
- current_tensor: Union[torch.Tensor, JointState, None],
- source_tensor: Union[torch.Tensor, JointState, None],
- idx: int,
- ):
- if source_tensor is not None:
- if current_tensor is None:
- current_tensor = source_tensor.clone()
- else:
- if isinstance(current_tensor, torch.Tensor) and isinstance(
- source_tensor, torch.Tensor
- ):
- current_tensor[idx] = source_tensor[idx]
- elif isinstance(current_tensor, JointState) and isinstance(
- source_tensor, JointState
- ):
- source_state = source_tensor[idx]
- current_tensor.copy_at_index(source_state, idx)
-
- return current_tensor
-
def copy_idx(self, idx: torch.Tensor, source_result: MotionGenResult):
+ """Copy data from source result to current result at index.
+
+ Args:
+ idx: index to copy data at.
+ source_result: source result to copy data from.
+
+ Returns:
+ MotionGenResult: copied result.
+ """
idx = idx.to(dtype=torch.long)
# copy data from source result:
self.success[idx] = source_result.success[idx]
@@ -906,6 +1164,15 @@ class MotionGenResult:
return self
def get_paths(self) -> List[JointState]:
+ """Get interpolated trajectories from the result. Use for batched queries.
+
+ This will return unsuccessful trajectories as well. Use
+ :meth:`MotionGenResult.get_successful_paths` to get only successful trajectories.
+
+ Returns:
+ List[JointState]: Interpolated trajectories. Check
+ :attr:`MotionGenResult.interpolation_dt` for the time between steps.
+ """
path = [
self.interpolated_plan[x].trim_trajectory(0, self.path_buffer_last_tstep[x])
for x in range(len(self.interpolated_plan))
@@ -913,6 +1180,12 @@ class MotionGenResult:
return path
def get_successful_paths(self) -> List[torch.Tensor]:
+ """Get successful interpolated trajectories from the result. Use for batched queries.
+
+ Returns:
+ List[JointState]: Interpolated trajectories. Check
+ :attr:`MotionGenResult.interpolation_dt` for the time between steps.
+ """
path = []
nz_i = torch.nonzero(self.success.view(-1)).view(-1)
path = self.interpolated_plan[nz_i]
@@ -928,26 +1201,133 @@ class MotionGenResult:
return path_list
def get_interpolated_plan(self) -> JointState:
+ """Get interpolated trajectory from the result.
+
+ Returns:
+ JointState: Interpolated trajectory. Check :attr:`MotionGenResult.interpolation_dt` for
+ the time between steps.
+ """
if self.path_buffer_last_tstep is None:
return self.interpolated_plan
if len(self.path_buffer_last_tstep) > 1:
- raise ValueError("only single result is supported")
+ log_error("only single result is supported")
return self.interpolated_plan.trim_trajectory(0, self.path_buffer_last_tstep[0])
+ def retime_trajectory(
+ self,
+ time_dilation_factor: float,
+ interpolate_trajectory: bool = True,
+ interpolation_dt: Optional[float] = None,
+ interpolation_kind: InterpolateType = InterpolateType.LINEAR_CUDA,
+ create_interpolation_buffer: bool = False,
+ ):
+ """Retime the optimized trajectory by a dilation factor.
+
+ Args:
+ time_dilation_factor: Time factor to slow down the trajectory. Should be less than 1.0.
+ interpolate_trajectory: Interpolate the trajectory after retiming.
+ interpolation_dt: Time between steps in the interpolated trajectory. If None,
+ :attr:`MotionGenResult.interpolation_dt` is used.
+ interpolation_kind: Interpolation type to use.
+ create_interpolation_buffer: Create a new buffer for interpolated trajectory. Set this
+ to True if existing buffer is not large enough to store new trajectory.
+ """
+
+ if time_dilation_factor > 1.0:
+ log_error("time_dilation_factor should be less than 1.0")
+ if time_dilation_factor == 1.0:
+ return
+ if len(self.path_buffer_last_tstep) > 1:
+ log_error("only single result is supported")
+
+ new_dt = self.optimized_dt * (1.0 / time_dilation_factor)
+ if len(self.optimized_plan.shape) == 3:
+ new_dt = new_dt.view(-1, 1, 1)
+ else:
+ new_dt = new_dt.view(-1, 1)
+ self.optimized_plan = self.optimized_plan.scale_by_dt(self.optimized_dt, new_dt)
+ self.optimized_dt = new_dt.view(-1)
+ if interpolate_trajectory:
+ if interpolation_dt is not None:
+ self.interpolation_dt = interpolation_dt
+ self.interpolated_plan, last_tstep, _ = get_batch_interpolated_trajectory(
+ self.optimized_plan,
+ self.optimized_dt,
+ self.interpolation_dt,
+ kind=interpolation_kind,
+ out_traj_state=self.interpolated_plan if not create_interpolation_buffer else None,
+ tensor_args=self.interpolated_plan.tensor_args,
+ optimize_dt=False,
+ )
+ self.path_buffer_last_tstep = [last_tstep[i] for i in range(len(last_tstep))]
+ if len(self.optimized_plan.shape) == 2:
+ self.interpolated_plan = self.interpolated_plan.squeeze(0)
+
@property
- def motion_time(self):
+ def motion_time(self) -> Union[float, torch.Tensor]:
+ """Compute motion time in seconds."""
+
# -2 as last three timesteps have the same value
# 0, 1 also have the same position value.
return self.optimized_dt * (self.optimized_plan.position.shape[-2] - 1 - 2 - 1)
+ @staticmethod
+ def _check_none_and_copy_idx(
+ current_tensor: Union[torch.Tensor, JointState, None],
+ source_tensor: Union[torch.Tensor, JointState, None],
+ idx: int,
+ ) -> Union[torch.Tensor, JointState]:
+ """Helper function to copy data from source tensor to current tensor at index.
+
+ Also supports copying from JointState to JointState.
+
+ Args:
+ current_tensor: tensor to copy data to.
+ source_tensor: tensor to copy data from.
+ idx: index to copy data at.
+
+ Returns:
+ Union[torch.Tensor, JointState]: copied tensor.
+ """
+ if source_tensor is not None:
+ if current_tensor is None:
+ current_tensor = source_tensor.clone()
+ else:
+ if isinstance(current_tensor, torch.Tensor) and isinstance(
+ source_tensor, torch.Tensor
+ ):
+ current_tensor[idx] = source_tensor[idx]
+ elif isinstance(current_tensor, JointState) and isinstance(
+ source_tensor, JointState
+ ):
+ source_state = source_tensor[idx]
+ current_tensor.copy_at_index(source_state, idx)
+
+ return current_tensor
+
class MotionGen(MotionGenConfig):
+ """Motion generation wrapper for generating collision-free trajectories.
+
+ This module provides an interface to generate collision-free trajectories for manipulators. It
+ supports Cartesian Pose for end-effectors and joint space goals. When a Cartesian Pose is used
+ as target, it uses batched inverse kinematics to find multiple joint configuration solutions
+ for the Cartesian Pose and then runs trajectory optimization with a linear interpolation from
+ start configuration to each of the IK solutions. When trajectory optimization fails, it uses
+ a graph planner to find collision-free paths to the IK solutions and then uses these paths as
+ seeds for trajectory optimization. The module also supports batched queries for motion
+ generation. Use this module as the high-level API for generating collision-free trajectories.
+ """
+
def __init__(self, config: MotionGenConfig):
+ """Initializes the motion generation module.
+
+ Args:
+ config: Configuration for motion generation.
+ """
super().__init__(**vars(config))
- self.rollout_fn = (
- self.graph_planner.safety_rollout_fn
- ) # TODO: compute_kinematics fn in rolloutbase
+ self.rollout_fn = self.graph_planner.safety_rollout_fn
self._trajopt_goal_config = None
self._dof = self.rollout_fn.d_action
self._batch_graph_search_buffer = None
@@ -960,6 +1340,12 @@ class MotionGen(MotionGenConfig):
self.update_batch_size(seeds=self.trajopt_seeds)
def update_batch_size(self, seeds=10, batch=1):
+ """Update the batch size for motion generation.
+
+ Args:
+ seeds: Number of seeds for trajectory optimization and graph planner.
+ batch: Number of queries to run in batch mode.
+ """
if (
self._trajopt_goal_config is None
or self._trajopt_goal_config.shape[0] != batch
@@ -994,6 +1380,26 @@ class MotionGen(MotionGenConfig):
use_nn_seed: bool = True,
newton_iters: Optional[int] = None,
) -> IKResult:
+ """Solve inverse kinematics for a given Cartesian Pose or a batch of Poses.
+
+ Use this only if your problem size is same as when using motion generation. If you want
+ to solve IK for a different problem size, create an instance of
+ :class:`curobo.wrap.reacher.ik_solver.IKSolver`.
+
+ Args:
+ goal_pose: Goal Pose for the end-effector.
+ retract_config: Retract configuration for the end-effector.
+ seed_config: Seed configuration for inverse kinematics.
+ return_seeds: Number of solutions to return per problem query.
+ num_seeds: Number of seeds to use for solving inverse kinematics.
+ use_nn_seed: Use neural network seed for solving inverse kinematics. This is not
+ implemented.
+ newton_iters: Number of Newton iterations to run for solving inverse kinematics.
+ This is useful to override the default number of iterations.
+
+ Returns:
+ IKResult: Result of inverse kinematics.
+ """
return self.ik_solver.solve(
goal_pose,
retract_config,
@@ -1004,383 +1410,22 @@ class MotionGen(MotionGenConfig):
newton_iters,
)
- @profiler.record_function("motion_gen/ik")
- def _solve_ik_from_solve_state(
- self,
- goal_pose: Pose,
- solve_state: ReacherSolveState,
- start_state: JointState,
- use_nn_seed: bool,
- partial_ik_opt: bool,
- link_poses: Optional[Dict[str, Pose]] = None,
- ) -> IKResult:
- newton_iters = None
- if partial_ik_opt:
- newton_iters = self.partial_ik_iters
- ik_result = self.ik_solver.solve_any(
- solve_state.solve_type,
- goal_pose,
- start_state.position.view(-1, self._dof),
- start_state.position.view(-1, 1, self._dof),
- solve_state.num_trajopt_seeds,
- solve_state.num_ik_seeds,
- use_nn_seed,
- newton_iters,
- link_poses,
- )
- return ik_result
-
- @profiler.record_function("motion_gen/trajopt_solve")
- def _solve_trajopt_from_solve_state(
- self,
- goal: Goal,
- solve_state: ReacherSolveState,
- act_seed: Optional[JointState] = None,
- use_nn_seed: bool = False,
- return_all_solutions: bool = False,
- seed_success: Optional[torch.Tensor] = None,
- newton_iters: Optional[int] = None,
- trajopt_instance: Optional[TrajOptSolver] = None,
- num_seeds_override: Optional[int] = None,
- ) -> TrajOptResult:
- if trajopt_instance is None:
- trajopt_instance = self.trajopt_solver
- if num_seeds_override is None:
- num_seeds_override = solve_state.num_trajopt_seeds
- traj_result = trajopt_instance.solve_any(
- solve_state.solve_type,
- goal,
- act_seed,
- use_nn_seed,
- return_all_solutions,
- num_seeds_override,
- seed_success,
- newton_iters=newton_iters,
- )
- return traj_result
-
@profiler.record_function("motion_gen/graph_search")
def graph_search(
self, start_config: T_BDOF, goal_config: T_BDOF, interpolation_steps: Optional[int] = None
) -> GraphResult:
+ """Run graph search to find collision-free paths between start and goal configurations.
+
+ Args:
+ start_config: Start joint configurations of the robot.
+ goal_config: Goal joint configurations of the robot.
+ interpolation_steps: Number of interpolation steps to interpolate obtained solutions.
+
+ Returns:
+ GraphResult: Result of graph search.
+ """
return self.graph_planner.find_paths(start_config, goal_config, interpolation_steps)
- def _get_solve_state(
- self,
- solve_type: ReacherSolveType,
- plan_config: MotionGenPlanConfig,
- goal_pose: Pose,
- start_state: JointState,
- ):
- # TODO: for batch seeds
- num_ik_seeds = (
- self.ik_seeds if plan_config.num_ik_seeds is None else plan_config.num_ik_seeds
- )
- num_trajopt_seeds = (
- self.trajopt_seeds
- if plan_config.num_trajopt_seeds is None
- else plan_config.num_trajopt_seeds
- )
-
- num_graph_seeds = (
- self.trajopt_seeds if plan_config.num_graph_seeds is None else num_trajopt_seeds
- )
- solve_state = None
- if solve_type == ReacherSolveType.SINGLE:
- solve_state = ReacherSolveState(
- solve_type,
- num_ik_seeds=num_ik_seeds,
- num_trajopt_seeds=num_trajopt_seeds,
- num_graph_seeds=num_graph_seeds,
- batch_size=1,
- n_envs=1,
- n_goalset=1,
- )
- elif solve_type == ReacherSolveType.GOALSET:
- solve_state = ReacherSolveState(
- solve_type,
- num_ik_seeds=num_ik_seeds,
- num_trajopt_seeds=num_trajopt_seeds,
- num_graph_seeds=num_graph_seeds,
- batch_size=1,
- n_envs=1,
- n_goalset=goal_pose.n_goalset,
- )
- elif solve_type == ReacherSolveType.BATCH:
- solve_state = ReacherSolveState(
- solve_type,
- num_ik_seeds=num_ik_seeds,
- num_trajopt_seeds=num_trajopt_seeds,
- num_graph_seeds=num_graph_seeds,
- batch_size=goal_pose.batch,
- n_envs=1,
- n_goalset=1,
- )
- elif solve_type == ReacherSolveType.BATCH_GOALSET:
- solve_state = ReacherSolveState(
- solve_type,
- num_ik_seeds=num_ik_seeds,
- num_trajopt_seeds=num_trajopt_seeds,
- num_graph_seeds=num_graph_seeds,
- batch_size=goal_pose.batch,
- n_envs=1,
- n_goalset=goal_pose.n_goalset,
- )
- elif solve_type == ReacherSolveType.BATCH_ENV:
- solve_state = ReacherSolveState(
- solve_type,
- num_ik_seeds=num_ik_seeds,
- num_trajopt_seeds=num_trajopt_seeds,
- num_graph_seeds=num_graph_seeds,
- batch_size=goal_pose.batch,
- n_envs=goal_pose.batch,
- n_goalset=1,
- )
- elif solve_type == ReacherSolveType.BATCH_ENV_GOALSET:
- solve_state = ReacherSolveState(
- solve_type,
- num_ik_seeds=num_ik_seeds,
- num_trajopt_seeds=num_trajopt_seeds,
- num_graph_seeds=num_graph_seeds,
- batch_size=goal_pose.batch,
- n_envs=goal_pose.batch,
- n_goalset=goal_pose.n_goalset,
- )
- else:
- raise ValueError("Unsupported Solve type")
- return solve_state
-
- def _plan_attempts(
- self,
- solve_state: ReacherSolveState,
- start_state: JointState,
- goal_pose: Pose,
- plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
- link_poses: List[Pose] = None,
- ):
- start_time = time.time()
- if plan_config.pose_cost_metric is not None:
- valid_query = self.update_pose_cost_metric(
- plan_config.pose_cost_metric, start_state, goal_pose
- )
- if not valid_query:
- result = MotionGenResult(
- success=torch.as_tensor([False], device=self.tensor_args.device),
- valid_query=valid_query,
- status="Invalid Hold Partial Pose",
- )
- return result
- # if plan_config.enable_opt:
- self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size)
- if solve_state.batch_env:
- if solve_state.batch_size > self.world_coll_checker.n_envs:
- raise ValueError("Batch Env is less that goal batch")
- force_graph = plan_config.enable_graph
- partial_ik = plan_config.partial_ik_opt
- ik_fail_count = 0
- time_dict = {
- "solve_time": 0,
- "ik_time": 0,
- "graph_time": 0,
- "trajopt_time": 0,
- "trajopt_attempts": 0,
- }
- best_status = 0 # 0== None, 1==IK Fail, 2== Graph Fail, 3==Opt Fail
- if plan_config.finetune_dt_scale is None:
- plan_config.finetune_dt_scale = self.finetune_dt_scale
- for n in range(plan_config.max_attempts):
- log_info("MG Iter: " + str(n))
- result = self._plan_from_solve_state(
- solve_state,
- start_state,
- goal_pose,
- plan_config,
- link_poses,
- )
- time_dict["solve_time"] += result.solve_time
- time_dict["ik_time"] += result.ik_time
- time_dict["graph_time"] += result.graph_time
- time_dict["trajopt_time"] += result.trajopt_time
- time_dict["trajopt_attempts"] += result.trajopt_attempts
- if (
- result.status == "IK Fail" and plan_config.ik_fail_return is not None
- ): # IF IK fails the first time, we exist assuming the goal is not reachable
- ik_fail_count += 1
- best_status = max(best_status, 1)
-
- if ik_fail_count > plan_config.ik_fail_return:
- break
- if result.success[0].item():
- break
-
- if result.status == "Finetune Opt Fail":
- plan_config.finetune_dt_scale *= (
- plan_config.finetune_dt_decay**plan_config.finetune_attempts
- )
- plan_config.finetune_dt_scale = min(plan_config.finetune_dt_scale, 1.25)
- if plan_config.enable_graph_attempt is not None and (
- n >= plan_config.enable_graph_attempt - 1
- and result.status in ["Opt Fail", "Finetune Opt Fail"]
- and not plan_config.enable_graph
- ):
- plan_config.enable_graph = True
- plan_config.partial_ik_opt = False
- if plan_config.disable_graph_attempt is not None and (
- n >= plan_config.disable_graph_attempt - 1
- and result.status in ["Opt Fail", "Graph Fail", "Finetune Opt Fail"]
- and not force_graph
- ):
- plan_config.enable_graph = False
- plan_config.partial_ik_opt = partial_ik
- if result.status in ["Opt Fail"]:
- best_status = 3
- elif result.status in ["Graph Fail"]:
- best_status = 2
- if time.time() - start_time > plan_config.timeout:
- break
- if not result.valid_query:
- result.status = "Invalid Problem"
- break
- if n == 10:
- self.reset_seed()
- log_warn("Couldn't find solution with 10 attempts, resetting seeds")
-
- result.solve_time = time_dict["solve_time"]
- result.ik_time = time_dict["ik_time"]
- result.graph_time = time_dict["graph_time"]
- result.trajopt_time = time_dict["trajopt_time"]
- result.trajopt_attempts = time_dict["trajopt_attempts"]
- result.attempts = n + 1
- torch.cuda.synchronize()
- if plan_config.pose_cost_metric is not None:
- self.update_pose_cost_metric(PoseCostMetric.reset_metric())
- result.total_time = time.time() - start_time
- return result
-
- def _plan_batch_attempts(
- self,
- solve_state: ReacherSolveState,
- start_state: JointState,
- goal_pose: Pose,
- plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
- ):
- start_time = time.time()
- goal_pose = goal_pose.clone()
- if plan_config.pose_cost_metric is not None:
- valid_query = self.update_pose_cost_metric(
- plan_config.pose_cost_metric, start_state, goal_pose
- )
- if not valid_query:
- result = MotionGenResult(
- success=torch.as_tensor(
- [False for _ in solve_state.batch_size],
- device=self.motion_gen.tensor_args.device,
- ),
- valid_query=valid_query,
- status="Invalid Hold Partial Pose",
- )
- return result
-
- if solve_state.batch_env:
- if solve_state.batch_size > self.world_coll_checker.n_envs:
- raise ValueError("Batch Env is less that goal batch")
- if plan_config.enable_graph:
- raise ValueError("Graph Search / Geometric Planner not supported in batch_env mode")
-
- if plan_config.enable_graph or (
- plan_config.enable_graph_attempt is not None
- and plan_config.max_attempts >= plan_config.enable_graph_attempt
- ):
- log_warn("Batch mode enable graph is only supported with num_graph_seeds==1")
- plan_config.num_trajopt_seeds = 1
- plan_config.num_graph_seeds = 1
- solve_state.num_trajopt_seeds = 1
- solve_state.num_graph_seeds = 1
- self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size)
-
- ik_fail_count = 0
- force_graph = plan_config.enable_graph
- partial_ik = plan_config.partial_ik_opt
-
- time_dict = {
- "solve_time": 0,
- "ik_time": 0,
- "graph_time": 0,
- "trajopt_time": 0,
- "trajopt_attempts": 0,
- }
- best_result = None
-
- for n in range(plan_config.max_attempts):
- result = self._plan_from_solve_state_batch(
- solve_state, start_state, goal_pose, plan_config
- )
-
- time_dict["solve_time"] += result.solve_time
- time_dict["ik_time"] += result.ik_time
-
- time_dict["graph_time"] += result.graph_time
- time_dict["trajopt_time"] += result.trajopt_time
- time_dict["trajopt_attempts"] += result.trajopt_attempts
-
- # if not all have succeeded, store the successful ones and re attempt:
- # TODO: update stored based on error
- if best_result is None:
- best_result = result.clone()
- else:
- # get success idx:
- idx = torch.nonzero(result.success).reshape(-1)
- if len(idx) > 0:
- best_result.copy_idx(idx, result)
-
- if (
- result.status == "IK Fail" and plan_config.ik_fail_return is not None
- ): # IF IK fails the first time, we exit assuming the goal is not reachable
- ik_fail_count += 1
- if ik_fail_count > plan_config.ik_fail_return:
- break
-
- if (
- torch.count_nonzero(best_result.success)
- >= goal_pose.batch * plan_config.success_ratio
- ): # we want 90% targets to succeed
- best_result.status = None
- break
- if plan_config.enable_graph_attempt is not None and (
- n >= plan_config.enable_graph_attempt - 1
- and result.status != "IK Fail"
- and not plan_config.enable_graph
- ):
- plan_config.enable_graph = True
- plan_config.partial_ik_opt = False
-
- if plan_config.disable_graph_attempt is not None and (
- n >= plan_config.disable_graph_attempt - 1
- and result.status in ["Opt Fail", "Graph Fail"]
- and not force_graph
- ):
- plan_config.enable_graph = False
- plan_config.partial_ik_opt = partial_ik
-
- if plan_config.fail_on_invalid_query:
- if not result.valid_query:
- best_result.valid_query = False
- best_result.status = "Invalid Problem"
- break
- if time.time() - start_time > plan_config.timeout:
- break
- best_result.solve_time = time_dict["solve_time"]
- best_result.ik_time = time_dict["ik_time"]
- best_result.graph_time = time_dict["graph_time"]
- best_result.trajopt_time = time_dict["trajopt_time"]
- best_result.trajopt_attempts = time_dict["trajopt_attempts"]
- best_result.attempts = n + 1
- torch.cuda.synchronize()
- if plan_config.pose_cost_metric is not None:
- self.update_pose_cost_metric(PoseCostMetric.reset_metric())
- best_result.total_time = time.time() - start_time
- return best_result
-
def plan_single(
self,
start_state: JointState,
@@ -1388,6 +1433,27 @@ class MotionGen(MotionGenConfig):
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
link_poses: List[Pose] = None,
) -> MotionGenResult:
+ """Plan a single motion to reach a goal pose from a start joint state.
+
+ This also supports reaching target poses for different links in the robot by providing
+ goal poses for each link in the link_poses argument. The link_poses argument should contain
+ a pose for each link specified in :attr:`MotionGen.kinematics`. Constrained planning is
+ supported by calling :meth:`MotionGen.update_pose_cost_metric` before calling this method.
+ See :ref:`tut_constrained_planning` for more details.
+
+ Args:
+ start_state: Start joint state of the robot. When planning from a non-static state, i.e,
+ when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt` to
+ False.
+ goal_pose: Goal pose for the end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for each link in the robot when planning for multiple links.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to see if the query was successful.
+ """
+ log_info("Planning for Single Goal: " + str(goal_pose.batch))
solve_state = self._get_solve_state(
ReacherSolveType.SINGLE, plan_config, goal_pose, start_state
)
@@ -1408,6 +1474,25 @@ class MotionGen(MotionGenConfig):
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
link_poses: List[Pose] = None,
) -> MotionGenResult:
+ """Plan a single motion to reach a goal from set of poses, from a start joint state.
+
+ Use this when planning to reach a grasp from a set of possible grasps. This method will
+ try to reach the closest goal pose from the set of goal poses at every iteration of
+ optimization during IK and trajectory optimization. This method only supports goalset for
+ main end-effector (i.e., `goal_pose`). Use single Pose target for links in `link_poses`.
+
+ Args:
+ start_state: Start joint state of the robot. When planning from a non-static state,
+ i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt`
+ to False.
+ goal_pose: Goal pose for the end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for each link in the robot when planning for multiple links.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to see if the query was successful.
+ """
solve_state = self._get_solve_state(
ReacherSolveType.GOALSET, plan_config, goal_pose, start_state
)
@@ -1426,7 +1511,22 @@ class MotionGen(MotionGenConfig):
start_state: JointState,
goal_pose: Pose,
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Dict[str, List[Pose]] = None,
) -> MotionGenResult:
+ """Plan motions to reach a batch of goal poses from a batch of start joint states.
+
+ Args:
+ start_state: Start joint states of the robot. When planning from a non-static state,
+ i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt`
+ to False.
+ goal_pose: Goal poses for the end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for each link in the robot when planning for multiple links.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to check which indices of the batch were successful.
+ """
solve_state = self._get_solve_state(
ReacherSolveType.BATCH, plan_config, goal_pose, start_state
)
@@ -1436,6 +1536,7 @@ class MotionGen(MotionGenConfig):
start_state,
goal_pose,
plan_config,
+ link_poses=link_poses,
)
return result
@@ -1444,7 +1545,23 @@ class MotionGen(MotionGenConfig):
start_state: JointState,
goal_pose: Pose,
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Dict[str, List[Pose]] = None,
) -> MotionGenResult:
+ """Plan motions to reach a batch of poses (goalset) from a batch of start joint states.
+
+ Args:
+ start_state: Start joint states of the robot. When planning from a non-static state,
+ i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt`
+ to False.
+ goal_pose: Goal poses for the end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for each link in the robot when planning for multiple links.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to check which indices of the batch were successful.
+ """
+
solve_state = self._get_solve_state(
ReacherSolveType.BATCH_GOALSET, plan_config, goal_pose, start_state
)
@@ -1454,24 +1571,7 @@ class MotionGen(MotionGenConfig):
start_state,
goal_pose,
plan_config,
- )
- return result
-
- def plan_batch_env_goalset(
- self,
- start_state: JointState,
- goal_pose: Pose,
- plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
- ) -> MotionGenResult:
- solve_state = self._get_solve_state(
- ReacherSolveType.BATCH_ENV_GOALSET, plan_config, goal_pose, start_state
- )
-
- result = self._plan_batch_attempts(
- solve_state,
- start_state,
- goal_pose,
- plan_config,
+ link_poses=link_poses,
)
return result
@@ -1480,10 +1580,85 @@ class MotionGen(MotionGenConfig):
start_state: JointState,
goal_pose: Pose,
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Dict[str, List[Pose]] = None,
) -> MotionGenResult:
+ """Plan motions to reach (batch) poses in different collision environments.
+
+ Args:
+ start_state: Start joint states of the robot. When planning from a non-static state,
+ i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt`
+ to False.
+ goal_pose: Goal poses for the end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for each link in the robot when planning for multiple links.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to check which indices of the batch were successful.
+ """
+ if plan_config.enable_graph:
+ log_info(
+ "Batch env mode does not support graph search, setting "
+ + "MotionGenPlanConfig.enable_graph=False"
+ )
+ plan_config.enable_graph = False
+
+ if plan_config.enable_graph_attempt is not None:
+ log_info(
+ "Batch env mode does not support graph search, setting "
+ + "MotionGenPlanConfig.enable_graph_attempt=None"
+ )
+ plan_config.enable_graph_attempt = None
solve_state = self._get_solve_state(
ReacherSolveType.BATCH_ENV, plan_config, goal_pose, start_state
)
+ result = self._plan_batch_attempts(
+ solve_state,
+ start_state,
+ goal_pose,
+ plan_config,
+ link_poses=link_poses,
+ )
+ return result
+
+ def plan_batch_env_goalset(
+ self,
+ start_state: JointState,
+ goal_pose: Pose,
+ plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Dict[str, List[Pose]] = None,
+ ) -> MotionGenResult:
+ """Plan motions to reach (batch) goalset poses in different collision environments.
+
+ Args:
+ start_state: Start joint states of the robot. When planning from a non-static state,
+ i.e, when velocity or acceleration is non-zero, set :attr:`MotionGen.optimize_dt`
+ to False.
+ goal_pose: Goal poses for the end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for each link in the robot when planning for multiple links.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to check which indices of the batch were successful.
+ """
+
+ if plan_config.enable_graph:
+ log_info(
+ "Batch env mode does not support graph search, setting "
+ + "MotionGenPlanConfig.enable_graph=False"
+ )
+ plan_config.enable_graph = False
+
+ if plan_config.enable_graph_attempt is not None:
+ log_info(
+ "Batch env mode does not support graph search, setting "
+ + "MotionGenPlanConfig.enable_graph_attempt=None"
+ )
+ plan_config.enable_graph_attempt = None
+ solve_state = self._get_solve_state(
+ ReacherSolveType.BATCH_ENV_GOALSET, plan_config, goal_pose, start_state
+ )
result = self._plan_batch_attempts(
solve_state,
start_state,
@@ -1492,786 +1667,85 @@ class MotionGen(MotionGenConfig):
)
return result
- def _plan_from_solve_state(
- self,
- solve_state: ReacherSolveState,
- start_state: JointState,
- goal_pose: Pose,
- plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
- link_poses: Optional[Dict[str, Pose]] = None,
- ) -> MotionGenResult:
- trajopt_seed_traj = None
- trajopt_seed_success = None
- trajopt_newton_iters = None
- graph_success = 0
- if len(start_state.shape) == 1:
- log_error("Joint state should be not a vector (dof) should be (bxdof)")
- # plan ik:
-
- ik_result = self._solve_ik_from_solve_state(
- goal_pose,
- solve_state,
- start_state,
- plan_config.use_nn_ik_seed,
- plan_config.partial_ik_opt,
- link_poses,
- )
-
- if not plan_config.enable_graph and plan_config.partial_ik_opt:
- ik_result.success[:] = True
-
- # check for success:
- result = MotionGenResult(
- ik_result.success.view(-1)[0:1], # This is not true for batch mode
- ik_time=ik_result.solve_time,
- solve_time=ik_result.solve_time,
- )
-
- if self.store_debug_in_result:
- result.debug_info = {"ik_result": ik_result}
- ik_success = torch.count_nonzero(ik_result.success)
- if ik_success == 0:
- result.status = "IK Fail"
- return result
-
- # do graph search:
- with profiler.record_function("motion_gen/post_ik"):
- ik_out_seeds = solve_state.num_trajopt_seeds
- if plan_config.enable_graph:
- ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success)
-
- goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof)[
- :ik_success
- ]
- start_config = tensor_repeat_seeds(start_state.position, ik_out_seeds)
- if plan_config.enable_opt:
- self._trajopt_goal_config[:] = ik_result.solution
-
- # do graph search:
- if plan_config.enable_graph:
- interpolation_steps = None
- if plan_config.enable_opt:
- interpolation_steps = self.trajopt_solver.action_horizon
- log_info("MG: running GP")
- graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
- trajopt_seed_success = graph_result.success
-
- graph_success = torch.count_nonzero(graph_result.success).item()
- result.graph_time = graph_result.solve_time
- result.solve_time += graph_result.solve_time
- if graph_success > 0:
- result.graph_plan = graph_result.interpolated_plan
- result.interpolated_plan = graph_result.interpolated_plan
-
- result.used_graph = True
- if plan_config.enable_opt:
- trajopt_seed = (
- result.graph_plan.position.view(
- 1, # solve_state.batch_size,
- graph_success, # solve_state.num_trajopt_seeds,
- interpolation_steps,
- self._dof,
- )
- .transpose(0, 1)
- .contiguous()
- )
- trajopt_seed_traj = torch.zeros(
- (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof),
- device=self.tensor_args.device,
- dtype=self.tensor_args.dtype,
- )
- trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed
- trajopt_seed_success = ik_result.success.clone()
- trajopt_seed_success[ik_result.success] = graph_result.success
-
- trajopt_seed_success = trajopt_seed_success.view(
- solve_state.batch_size, solve_state.num_trajopt_seeds
- )
- trajopt_newton_iters = self.graph_trajopt_iters
- else:
- _, idx = torch.topk(
- graph_result.path_length[graph_result.success], k=1, largest=False
- )
- result.interpolated_plan = result.interpolated_plan[idx].squeeze(0)
- result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt)
- result.optimized_plan = result.interpolated_plan[
- : graph_result.path_buffer_last_tstep[idx.item()]
- ]
- idx = idx.view(-1) + self._batch_col
- result.position_error = ik_result.position_error[ik_result.success][
- graph_result.success
- ][idx]
- result.rotation_error = ik_result.rotation_error[ik_result.success][
- graph_result.success
- ][idx]
- result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[
- idx.item() : idx.item() + 1
- ]
- result.success = result.success.view(-1)[0:1]
- result.success[:] = True
- return result
- else:
- result.success[:] = False
- result.status = "Graph Fail"
- if not graph_result.valid_query:
- result.valid_query = False
- if self.store_debug_in_result:
- result.debug_info["graph_debug"] = graph_result.debug_info
- return result
- if plan_config.need_graph_success:
- return result
-
- # do trajopt::
-
- if plan_config.enable_opt:
- with profiler.record_function("motion_gen/setup_trajopt_seeds"):
- self._trajopt_goal_config[:, :ik_success] = goal_config
- retract_config = None
- if plan_config.use_start_state_as_retract:
- retract_config = start_state.position.clone()
- goal = Goal(
- goal_pose=goal_pose,
- current_state=start_state,
- links_goal_pose=link_poses,
- retract_state=retract_config,
- )
-
- if (
- trajopt_seed_traj is None
- or graph_success < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
- ):
- goal_config = self._trajopt_goal_config[0] # batch index == 0
-
- goal_state = JointState.from_position(
- goal_config,
- )
- seed_link_poses = None
- if link_poses is not None:
- seed_link_poses = {}
-
- for k in link_poses.keys():
- seed_link_poses[k] = link_poses[k].repeat_seeds(
- solve_state.num_trajopt_seeds
- )
- seed_goal = Goal(
- goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds),
- current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
- goal_state=goal_state,
- links_goal_pose=seed_link_poses,
- )
- if trajopt_seed_traj is not None:
- trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
- # batch, num_seeds, h, dof
- if (
- trajopt_seed_success.shape[1]
- < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
- ):
- trajopt_seed_success_new = torch.zeros(
- (1, solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds),
- device=self.tensor_args.device,
- dtype=torch.bool,
- )
- trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = (
- trajopt_seed_success
- )
- trajopt_seed_success = trajopt_seed_success_new
- # create seeds here:
- trajopt_seed_traj = self.trajopt_solver.get_seed_set(
- seed_goal,
- trajopt_seed_traj, # batch, num_seeds, h, dof
- num_seeds=self.noisy_trajopt_seeds,
- batch_mode=solve_state.batch_mode,
- seed_success=trajopt_seed_success,
- )
- trajopt_seed_traj = trajopt_seed_traj.view(
- solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
- solve_state.batch_size,
- self.trajopt_solver.action_horizon,
- self._dof,
- ).contiguous()
- if plan_config.enable_finetune_trajopt:
- og_value = self.trajopt_solver.interpolation_type
- self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
- with profiler.record_function("motion_gen/trajopt"):
- log_info("MG: running TO")
- traj_result = self._solve_trajopt_from_solve_state(
- goal,
- solve_state,
- trajopt_seed_traj,
- num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
- newton_iters=trajopt_newton_iters,
- return_all_solutions=plan_config.parallel_finetune
- and plan_config.enable_finetune_trajopt,
- )
- if plan_config.enable_finetune_trajopt:
- self.trajopt_solver.interpolation_type = og_value
- if self.store_debug_in_result:
- result.debug_info["trajopt_result"] = traj_result
- # run finetune
- if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
- with profiler.record_function("motion_gen/finetune_trajopt"):
- seed_traj = traj_result.raw_action.clone() # solution.position.clone()
- seed_traj = seed_traj.contiguous()
- og_solve_time = traj_result.solve_time
- seed_override = 1
- opt_dt = traj_result.optimized_dt
-
- if plan_config.parallel_finetune:
- opt_dt = torch.min(opt_dt[traj_result.success])
- seed_override = solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
-
- finetune_time = 0
- for k in range(plan_config.finetune_attempts):
- newton_iters = None
-
- scaled_dt = torch.clamp(
- opt_dt
- * plan_config.finetune_dt_scale
- * (plan_config.finetune_dt_decay ** (k)),
- self.trajopt_solver.interpolation_dt,
- )
- if self.optimize_dt:
- self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item())
-
- traj_result = self._solve_trajopt_from_solve_state(
- goal,
- solve_state,
- seed_traj,
- trajopt_instance=self.finetune_trajopt_solver,
- num_seeds_override=seed_override,
- newton_iters=newton_iters,
- )
- finetune_time += traj_result.solve_time
- if torch.count_nonzero(traj_result.success) > 0:
- break
- seed_traj = traj_result.optimized_seeds.detach().clone()
- newton_iters = 4
-
- traj_result.solve_time = finetune_time
-
- result.finetune_time = traj_result.solve_time
-
- traj_result.solve_time = og_solve_time
- if self.store_debug_in_result:
- result.debug_info["finetune_trajopt_result"] = traj_result
- elif plan_config.enable_finetune_trajopt:
- traj_result.success = traj_result.success[0:1]
- # if torch.count_nonzero(result.success) == 0:
- result.status = "Opt Fail"
- result.solve_time += traj_result.solve_time + result.finetune_time
- result.trajopt_time = traj_result.solve_time
- result.trajopt_attempts = 1
- result.success = traj_result.success
-
- if plan_config.enable_finetune_trajopt and torch.count_nonzero(result.success) == 0:
- result.status = "Finetune Opt Fail"
-
- result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory(
- 0, traj_result.path_buffer_last_tstep[0]
- )
- result.interpolation_dt = self.trajopt_solver.interpolation_dt
- result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
- result.position_error = traj_result.position_error
- result.rotation_error = traj_result.rotation_error
- result.optimized_dt = traj_result.optimized_dt
- result.optimized_plan = traj_result.solution
- result.goalset_index = traj_result.goalset_index
- return result
-
- def _plan_js_from_solve_state(
- self,
- solve_state: ReacherSolveState,
- start_state: JointState,
- goal_state: JointState,
- plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
- ) -> MotionGenResult:
- trajopt_seed_traj = None
- trajopt_seed_success = None
- trajopt_newton_iters = self.js_trajopt_solver.newton_iters
-
- graph_success = 0
- if len(start_state.shape) == 1:
- log_error("Joint state should be not a vector (dof) should be (bxdof)")
-
- result = MotionGenResult(cspace_error=torch.zeros((1), device=self.tensor_args.device))
- if self.store_debug_in_result:
- result.debug_info = {}
- # do graph search:
- if plan_config.enable_graph and False:
- start_config = torch.zeros(
- (solve_state.num_graph_seeds, self.js_trajopt_solver.dof),
- device=self.tensor_args.device,
- dtype=self.tensor_args.dtype,
- )
- goal_config = start_config.clone()
- start_config[:] = start_state.position
- goal_config[:] = goal_state.position
- interpolation_steps = None
- if plan_config.enable_opt:
- interpolation_steps = self.js_trajopt_solver.action_horizon
- log_info("MG: running GP")
- graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
- trajopt_seed_success = graph_result.success
-
- graph_success = torch.count_nonzero(graph_result.success).item()
- result.graph_time = graph_result.solve_time
- result.solve_time += graph_result.solve_time
- if graph_success > 0:
- result.graph_plan = graph_result.interpolated_plan
- result.interpolated_plan = graph_result.interpolated_plan
-
- result.used_graph = True
- if plan_config.enable_opt:
- trajopt_seed = (
- result.graph_plan.position.view(
- 1, # solve_state.batch_size,
- graph_success, # solve_state.num_trajopt_seeds,
- interpolation_steps,
- self._dof,
- )
- .transpose(0, 1)
- .contiguous()
- )
- trajopt_seed_traj = torch.zeros(
- (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof),
- device=self.tensor_args.device,
- dtype=self.tensor_args.dtype,
- )
- trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed
- trajopt_seed_success = graph_result.success
-
- trajopt_seed_success = trajopt_seed_success.view(
- 1, solve_state.num_trajopt_seeds
- )
- trajopt_newton_iters = self.graph_trajopt_iters
- else:
- _, idx = torch.topk(
- graph_result.path_length[graph_result.success], k=1, largest=False
- )
- result.interpolated_plan = result.interpolated_plan[idx].squeeze(0)
- result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt)
- result.optimized_plan = result.interpolated_plan[
- : graph_result.path_buffer_last_tstep[idx.item()]
- ]
- idx = idx.view(-1) + self._batch_col
- result.cspace_error = torch.zeros((1), device=self.tensor_args.device)
-
- result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[
- idx.item() : idx.item() + 1
- ]
- result.success = torch.as_tensor([True], device=self.tensor_args.device)
- return result
- else:
- result.success = torch.as_tensor([False], device=self.tensor_args.device)
- result.status = "Graph Fail"
- if not graph_result.valid_query:
- result.valid_query = False
- if self.store_debug_in_result:
- result.debug_info["graph_debug"] = graph_result.debug_info
- return result
- if plan_config.need_graph_success:
- return result
-
- # do trajopt:
- if plan_config.enable_opt:
- with profiler.record_function("motion_gen/setup_trajopt_seeds"):
-
- goal = Goal(
- current_state=start_state,
- goal_state=goal_state,
- )
-
- if trajopt_seed_traj is None or graph_success < solve_state.num_trajopt_seeds * 1:
- seed_goal = Goal(
- current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
- goal_state=goal_state.repeat_seeds(solve_state.num_trajopt_seeds),
- )
- if trajopt_seed_traj is not None:
- trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
- # batch, num_seeds, h, dof
- if trajopt_seed_success.shape[1] < self.js_trajopt_solver.num_seeds:
- trajopt_seed_success_new = torch.zeros(
- (1, solve_state.num_trajopt_seeds),
- device=self.tensor_args.device,
- dtype=torch.bool,
- )
- trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = (
- trajopt_seed_success
- )
- trajopt_seed_success = trajopt_seed_success_new
- # create seeds here:
- trajopt_seed_traj = self.js_trajopt_solver.get_seed_set(
- seed_goal,
- trajopt_seed_traj, # batch, num_seeds, h, dof
- num_seeds=1,
- batch_mode=False,
- seed_success=trajopt_seed_success,
- )
- trajopt_seed_traj = (
- trajopt_seed_traj.view(
- self.js_trajopt_solver.num_seeds * 1,
- 1,
- self.trajopt_solver.action_horizon,
- self._dof,
- )
- .contiguous()
- .clone()
- )
- if plan_config.enable_finetune_trajopt:
- og_value = self.trajopt_solver.interpolation_type
- self.js_trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
- with profiler.record_function("motion_gen/trajopt"):
- log_info("MG: running TO")
- traj_result = self._solve_trajopt_from_solve_state(
- goal,
- solve_state,
- trajopt_seed_traj,
- num_seeds_override=solve_state.num_trajopt_seeds,
- newton_iters=trajopt_newton_iters + 2,
- return_all_solutions=plan_config.enable_finetune_trajopt,
- trajopt_instance=self.js_trajopt_solver,
- )
- if plan_config.enable_finetune_trajopt:
- self.trajopt_solver.interpolation_type = og_value
- # self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate)
- if self.store_debug_in_result:
- result.debug_info["trajopt_result"] = traj_result
- if torch.count_nonzero(traj_result.success) == 0:
- result.status = "TrajOpt Fail"
- # run finetune
- if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
- with profiler.record_function("motion_gen/finetune_trajopt"):
- seed_traj = traj_result.raw_action.clone() # solution.position.clone()
- seed_traj = seed_traj.contiguous()
- og_solve_time = traj_result.solve_time
-
- scaled_dt = torch.clamp(
- torch.max(traj_result.optimized_dt[traj_result.success]),
- self.trajopt_solver.interpolation_dt,
- )
- og_dt = self.js_trajopt_solver.solver_dt.clone()
- self.js_trajopt_solver.update_solver_dt(scaled_dt.item())
- traj_result = self._solve_trajopt_from_solve_state(
- goal,
- solve_state,
- seed_traj,
- trajopt_instance=self.js_trajopt_solver,
- num_seeds_override=solve_state.num_trajopt_seeds,
- newton_iters=trajopt_newton_iters + 4,
- )
- self.js_trajopt_solver.update_solver_dt(og_dt)
-
- result.finetune_time = traj_result.solve_time
-
- traj_result.solve_time = og_solve_time
- if self.store_debug_in_result:
- result.debug_info["finetune_trajopt_result"] = traj_result
- if torch.count_nonzero(traj_result.success) == 0:
- result.status = "Finetune Fail"
- elif plan_config.enable_finetune_trajopt:
- traj_result.success = traj_result.success[0:1]
- result.solve_time += traj_result.solve_time + result.finetune_time
- result.trajopt_time = traj_result.solve_time
- result.trajopt_attempts = 1
- result.success = traj_result.success
-
- result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory(
- 0, traj_result.path_buffer_last_tstep[0]
- )
-
- result.interpolation_dt = self.trajopt_solver.interpolation_dt
- result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
- result.cspace_error = traj_result.cspace_error
- result.optimized_dt = traj_result.optimized_dt
- result.optimized_plan = traj_result.solution
- result.goalset_index = traj_result.goalset_index
-
- return result
-
- def _plan_from_solve_state_batch(
- self,
- solve_state: ReacherSolveState,
- start_state: JointState,
- goal_pose: Pose,
- plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
- ) -> MotionGenResult:
- self._trajopt_goal_config[:] = self.get_retract_config().view(1, 1, self._dof)
- trajopt_seed_traj = None
- trajopt_seed_success = None
- trajopt_newton_iters = None
- graph_success = 0
-
- # plan ik:
- ik_result = self._solve_ik_from_solve_state(
- goal_pose,
- solve_state,
- start_state,
- plan_config.use_nn_ik_seed,
- plan_config.partial_ik_opt,
- )
-
- if not plan_config.enable_graph and plan_config.partial_ik_opt:
- ik_result.success[:] = True
-
- # check for success:
- result = MotionGenResult(
- ik_result.success,
- position_error=ik_result.position_error,
- rotation_error=ik_result.rotation_error,
- ik_time=ik_result.solve_time,
- solve_time=ik_result.solve_time,
- debug_info={},
- # goalset_index=ik_result.goalset_index,
- )
-
- ik_success = torch.count_nonzero(ik_result.success)
- if ik_success == 0:
- result.status = "IK Fail"
- result.success = result.success[:, 0]
- return result
-
- # do graph search:
- ik_out_seeds = solve_state.num_trajopt_seeds
- if plan_config.enable_graph:
- ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success)
-
- # if not plan_config.enable_opt and plan_config.enable_graph:
- # self.graph_planner.interpolation_steps = self.interpolation_steps
- # self.graph_planner.interpolation_type = self.interpolation_type
- # elif plan_config.enable_graph:
- # self.graph_planner.interpolation_steps = self.trajopt_solver.traj_tsteps
- # self.graph_planner.interpolation_type = InterpolateType.LINEAR
- goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof)
-
- # get shortest path
- if plan_config.enable_graph:
- interpolation_steps = None
- if plan_config.enable_opt:
- interpolation_steps = self.trajopt_solver.action_horizon
-
- start_graph_state = start_state.repeat_seeds(ik_out_seeds)
- start_config = start_graph_state.position[ik_result.success.view(-1)].view(
- -1, self.ik_solver.dof
- )
- graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
- graph_success = torch.count_nonzero(graph_result.success).item()
-
- result.graph_time = graph_result.solve_time
- result.solve_time += graph_result.solve_time
- if graph_success > 0:
- # path = graph_result.interpolated_plan
- result.graph_plan = graph_result.interpolated_plan
- result.interpolated_plan = graph_result.interpolated_plan
- result.used_graph = True
-
- if plan_config.enable_opt:
- trajopt_seed = result.graph_plan.position.view(
- graph_success, # solve_state.num_trajopt_seeds,
- interpolation_steps,
- self._dof,
- ).contiguous()
- trajopt_seed_traj = torch.zeros(
- (1, trajopt_seed.shape[0], self.trajopt_solver.action_horizon, self._dof),
- device=self.tensor_args.device,
- dtype=self.tensor_args.dtype,
- )
- trajopt_seed_traj[0, :, :interpolation_steps, :] = trajopt_seed
- trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
- trajopt_seed_success = ik_result.success.clone()
- trajopt_seed_success[ik_result.success] = graph_result.success
-
- trajopt_seed_success = trajopt_seed_success.view(
- solve_state.num_trajopt_seeds, solve_state.batch_size
- )
- trajopt_newton_iters = self.graph_trajopt_iters
-
- else:
- ik_success = ik_result.success.view(-1).clone()
-
- # only some might be successful:
-
- g_dim = torch.nonzero(ik_success).view(-1)[graph_result.success]
-
- self._batch_graph_search_buffer.copy_at_index(
- graph_result.interpolated_plan, g_dim
- )
-
- # result.graph_plan = JointState.from_position(
- # self._batch_graph_search_buffer,
- # joint_names=graph_result.interpolated_plan.joint_names,
- # )
- result.interpolated_plan = self._batch_graph_search_buffer
- g_dim = g_dim.cpu().squeeze().tolist()
- for x, x_val in enumerate(g_dim):
- self._batch_path_buffer_last_tstep[x_val] = (
- graph_result.path_buffer_last_tstep[x]
- )
- result.path_buffer_last_tstep = self._batch_path_buffer_last_tstep
- result.optimized_plan = result.interpolated_plan
- result.optimized_dt = torch.as_tensor(
- [
- self.interpolation_dt
- for i in range(result.interpolated_plan.position.shape[0])
- ],
- device=self.tensor_args.device,
- dtype=self.tensor_args.dtype,
- )
- result.success = result.success.view(-1).clone()
- result.success[ik_success][graph_result.success] = True
- return result
-
- else:
- result.success[:] = False
- result.success = result.success[:, 0]
- result.status = "Graph Fail"
- if not graph_result.valid_query:
- result.valid_query = False
- if self.store_debug_in_result:
- result.debug_info = {"graph_debug": graph_result.debug_info}
- return result
-
- if plan_config.enable_opt:
- # get goal configs based on ik success:
- self._trajopt_goal_config[ik_result.success] = goal_config
-
- goal_config = self._trajopt_goal_config # batch index == 0
-
- goal = Goal(
- goal_pose=goal_pose,
- current_state=start_state,
- )
- # generate seeds:
- if trajopt_seed_traj is None or (
- plan_config.enable_graph and graph_success < solve_state.batch_size
- ):
- seed_goal = Goal(
- goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds),
- current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
- goal_state=JointState.from_position(goal_config.view(-1, self._dof)),
- )
- if trajopt_seed_traj is not None:
- trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
-
- # create seeds here:
- trajopt_seed_traj = self.trajopt_solver.get_seed_set(
- seed_goal,
- trajopt_seed_traj, # batch, num_seeds, h, dof
- num_seeds=1,
- batch_mode=solve_state.batch_mode,
- seed_success=trajopt_seed_success,
- )
- trajopt_seed_traj = trajopt_seed_traj.view(
- solve_state.num_trajopt_seeds,
- solve_state.batch_size,
- self.trajopt_solver.action_horizon,
- self._dof,
- ).contiguous()
- if plan_config.enable_finetune_trajopt:
- og_value = self.trajopt_solver.interpolation_type
- self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
-
- traj_result = self._solve_trajopt_from_solve_state(
- goal,
- solve_state,
- trajopt_seed_traj,
- newton_iters=trajopt_newton_iters,
- return_all_solutions=True,
- )
-
- # output of traj result will have 1 solution per batch
-
- # run finetune opt on 1 solution per batch:
- if plan_config.enable_finetune_trajopt:
- self.trajopt_solver.interpolation_type = og_value
- # self.trajopt_solver.compute_metrics(not og_evaluate, og_evaluate)
- if self.store_debug_in_result:
- result.debug_info["trajopt_result"] = traj_result
-
- # run finetune
- if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
- with profiler.record_function("motion_gen/finetune_trajopt"):
- seed_traj = traj_result.raw_action.clone() # solution.position.clone()
- seed_traj = seed_traj.contiguous()
- og_solve_time = traj_result.solve_time
-
- scaled_dt = torch.clamp(
- torch.max(traj_result.optimized_dt[traj_result.success])
- * self.finetune_dt_scale,
- self.trajopt_solver.interpolation_dt,
- )
- self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item())
-
- traj_result = self._solve_trajopt_from_solve_state(
- goal,
- solve_state,
- seed_traj,
- trajopt_instance=self.finetune_trajopt_solver,
- num_seeds_override=solve_state.num_trajopt_seeds,
- )
-
- result.finetune_time = traj_result.solve_time
-
- traj_result.solve_time = og_solve_time
- if self.store_debug_in_result:
- result.debug_info["finetune_trajopt_result"] = traj_result
- elif plan_config.enable_finetune_trajopt:
- traj_result.success = traj_result.success[:, 0]
-
- result.success = traj_result.success
-
- result.interpolated_plan = traj_result.interpolated_solution
- result.solve_time += traj_result.solve_time
- result.trajopt_time = traj_result.solve_time
- result.position_error = traj_result.position_error
- result.rotation_error = traj_result.rotation_error
- result.cspace_error = traj_result.cspace_error
- result.goalset_index = traj_result.goalset_index
- result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
- result.optimized_plan = traj_result.solution
- result.optimized_dt = traj_result.optimized_dt
- if torch.count_nonzero(traj_result.success) == 0:
- result.status = "Opt Fail"
- result.success[:] = False
- if self.store_debug_in_result:
- result.debug_info = {"trajopt_result": traj_result}
- return result
-
def compute_kinematics(self, state: JointState) -> KinematicModelState:
+ """Compute kinematics for a given joint state.
+
+ Args:
+ state: Joint state of the robot. Only :attr:`JointState.position` is used.
+
+ Returns:
+ KinematicModelState: Kinematic state of the robot.
+ """
out = self.rollout_fn.compute_kinematics(state)
return out
@property
- def kinematics(self):
+ def kinematics(self) -> CudaRobotModel:
+ """Returns the shared kinematics model of the robot."""
return self.rollout_fn.kinematics
@property
- def dof(self):
+ def dof(self) -> int:
+ """Returns the number of controlled degrees of freedom of the robot."""
return self.kinematics.get_dof()
+ @property
+ def collision_cache(self) -> Dict[str, int]:
+ """Returns the collision cache created by the world collision checker."""
+ return self.world_coll_checker.cache
+
def check_constraints(self, state: JointState) -> RolloutMetrics:
+ """Compute IK constraints for a given joint state.
+
+ Args:
+ state: Joint state of the robot.
+
+ Returns:
+ RolloutMetrics: Metrics for the joint state.
+ """
metrics = self.ik_solver.check_constraints(state)
return metrics
def update_world(self, world: WorldConfig):
+ """Update the world representation for collision checking.
+
+ This allows for updating the world representation as long as the new world representation
+ does not have a larger number of obstacles than the :attr:`MotionGen.collision_cache` as
+ created during initialization of :class:`MotionGenConfig`. Updating the world also
+ invalidates the cached roadmaps in the graph planner. See :ref:`world_collision` for more
+ details.
+
+ Args:
+ world: New world configuration for collision checking.
+ """
self.world_coll_checker.load_collision_model(world, fix_cache_reference=self.use_cuda_graph)
self.graph_planner.reset_graph()
- return True
def clear_world_cache(self):
+ """Remove all collision objects from collision cache."""
+
self.world_coll_checker.clear_cache()
def reset(self, reset_seed=True):
+ """Reset the motion generation module.
+
+ Args:
+ reset_seed: Reset the random seed generator. Resetting this can be time consuming, if
+ deterministic results are not required, set this to False.
+ """
self.graph_planner.reset_buffer()
if reset_seed:
self.reset_seed()
def reset_seed(self):
+ """Reset the random seed generators in all sub-modules of motion generation."""
self.rollout_fn.reset_seed()
self.ik_solver.reset_seed()
self.graph_planner.reset_seed()
self.trajopt_solver.reset_seed()
- def get_retract_config(self):
+ def get_retract_config(self) -> T_DOF:
+ """Returns the retract/home configuration of the robot."""
return self.rollout_fn.dynamics_model.retract_config
def warmup(
@@ -2285,6 +1759,23 @@ class MotionGen(MotionGenConfig):
warmup_joint_index: int = 0,
warmup_joint_delta: float = 0.1,
):
+ """Warmup planning methods for motion generation.
+
+ Args:
+ enable_graph: Enable graph search for warmup.
+ batch: Number of problem queries for batch mode. If None, single query is run.
+ warmup_js_trajopt: Warmup joint space planning in addition to Cartesian planning.
+ batch_env_mode: Enable batch world environment mode for warmup. Only used when batch is
+ not None.
+ parallel_finetune: Run finetuning trajectory optimization in parallel for warmup. Leave
+ this to True for most cases.
+ n_goalset: Number of goal poses to use for warmup. If -1, single goal pose is used. Set
+ this to the largest number of goals you plan to use with
+ :meth:`MotionGen.plan_goalset`. After warmup, you can use smaller number of goals
+ and the method will internally pad the extra goals with the first goal.
+ warmup_joint_index: Index of the joint to perturb for warmup.
+ warmup_joint_delta: Delta to perturb the joint for warmup.
+ """
log_info("Warmup")
if warmup_js_trajopt:
start_state = JointState.from_position(
@@ -2391,7 +1882,8 @@ class MotionGen(MotionGenConfig):
MotionGenPlanConfig(
max_attempts=10,
enable_finetune_trajopt=True,
- enable_graph=enable_graph,
+ enable_graph=False,
+ enable_graph_attempt=None,
),
)
else:
@@ -2402,6 +1894,7 @@ class MotionGen(MotionGenConfig):
max_attempts=10,
enable_finetune_trajopt=True,
enable_graph=enable_graph,
+ enable_graph_attempt=None if not enable_graph else 20,
),
)
else:
@@ -2420,8 +1913,7 @@ class MotionGen(MotionGenConfig):
MotionGenPlanConfig(
max_attempts=10,
enable_finetune_trajopt=True,
- enable_graph=enable_graph,
- enable_graph_attempt=20,
+ enable_graph=False,
),
)
else:
@@ -2432,12 +1924,11 @@ class MotionGen(MotionGenConfig):
max_attempts=10,
enable_finetune_trajopt=True,
enable_graph=enable_graph,
- enable_graph_attempt=20,
+ enable_graph_attempt=None if not enable_graph else 20,
),
)
log_info("Warmup complete")
- return True
def plan_single_js(
self,
@@ -2445,7 +1936,23 @@ class MotionGen(MotionGenConfig):
goal_state: JointState,
plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
) -> MotionGenResult:
- # NOTE: currently runs only one seed
+ """Plan a single motion to reach a goal joint state from a start joint state.
+
+ This method uses trajectory optimization to find a collision-free path between the start
+ and goal joint states. If trajectory optimization fails, it uses a graph planner to find
+ a collision-free path to the goal joint state. The graph plan is then used as a seed for
+ trajectory optimization.
+
+ Args:
+ start_state: Start joint state of the robot.
+ goal_state: Goal joint state of the robot.
+ plan_config: Planning parameters for motion generation.
+
+ Returns:
+ MotionGenResult: Result of motion generation. Check :attr:`MotionGenResult.success`
+ attribute to see if the query was successful.
+ """
+
time_dict = {
"solve_time": 0,
"ik_time": 0,
@@ -2495,84 +2002,11 @@ class MotionGen(MotionGenConfig):
result.solve_time = result.trajopt_time + result.graph_time + result.finetune_time
result.total_time = result.solve_time
result.attempts = n
-
- return result
-
- def plan(
- self,
- start_state: JointState,
- goal_pose: Pose,
- enable_graph: bool = True,
- enable_opt: bool = True,
- use_nn_ik_seed: bool = False,
- need_graph_success: bool = False,
- max_attempts: int = 100,
- timeout: float = 10.0,
- enable_graph_attempt: Optional[int] = None,
- disable_graph_attempt: Optional[int] = None,
- trajopt_attempts: int = 1,
- ik_fail_return: Optional[int] = None,
- partial_ik_opt: bool = False,
- num_ik_seeds: Optional[int] = None,
- num_graph_seeds: Optional[int] = None,
- num_trajopt_seeds: Optional[int] = None,
- ):
- plan_config = MotionGenPlanConfig(
- enable_graph,
- enable_opt,
- use_nn_ik_seed,
- need_graph_success,
- max_attempts,
- timeout,
- enable_graph_attempt,
- disable_graph_attempt,
- ik_fail_return,
- partial_ik_opt,
- num_ik_seeds,
- num_graph_seeds,
- num_trajopt_seeds,
- )
- result = self.plan_single(start_state, goal_pose, plan_config)
- return result
-
- def batch_plan(
- self,
- start_state: JointState,
- goal_pose: Pose,
- enable_graph: bool = True,
- enable_opt: bool = True,
- use_nn_ik_seed: bool = False,
- need_graph_success: bool = False,
- max_attempts: int = 1,
- timeout: float = 10.0,
- enable_graph_attempt: Optional[int] = None,
- disable_graph_attempt: Optional[int] = None,
- success_ratio: float = 1.0,
- ik_fail_return: Optional[int] = None,
- fail_on_invalid_query: bool = False,
- partial_ik_opt: bool = False,
- num_ik_seeds: Optional[int] = None,
- num_graph_seeds: Optional[int] = None,
- num_trajopt_seeds: Optional[int] = None,
- ):
- plan_config = MotionGenPlanConfig(
- enable_graph,
- enable_opt,
- use_nn_ik_seed,
- need_graph_success,
- max_attempts,
- timeout,
- enable_graph_attempt,
- disable_graph_attempt,
- ik_fail_return,
- partial_ik_opt,
- num_ik_seeds,
- num_graph_seeds,
- num_trajopt_seeds,
- success_ratio=success_ratio,
- fail_on_invalid_query=fail_on_invalid_query,
- )
- result = self.plan_batch(start_state, goal_pose, plan_config)
+ if plan_config.time_dilation_factor is not None and torch.count_nonzero(result.success) > 0:
+ result.retime_trajectory(
+ plan_config.time_dilation_factor,
+ interpolation_kind=self.js_trajopt_solver.interpolation_type,
+ )
return result
def solve_trajopt(
@@ -2582,6 +2016,17 @@ class MotionGen(MotionGenConfig):
return_all_solutions: bool = False,
num_seeds: Optional[int] = None,
):
+ """Solve trajectory optimization for a given goal.
+
+ Args:
+ goal: Goal for trajectory optimization.
+ act_seed: Seed for trajectory optimization.
+ return_all_solutions: Return results for all seeds in trajectory optimization.
+ num_seeds: Override number of seeds to use for trajectory optimization.
+
+ Returns:
+ TrajOptResult: Result of trajectory optimization.
+ """
result = self.trajopt_solver.solve(
goal, act_seed, return_all_solutions=return_all_solutions, num_seeds=num_seeds
)
@@ -2591,6 +2036,17 @@ class MotionGen(MotionGenConfig):
self,
in_js: JointState,
):
+ """Get controlled joint state from input joint state.
+
+ This is used to get the joint state for only joints that are optimization variables. This
+ also re-orders the joints to match the order of optimization variables.
+
+ Args:
+ in_js: Input joint state.
+
+ Returns:
+ JointState: Active joint state.
+ """
opt_jnames = self.rollout_fn.joint_names
opt_js = in_js.get_ordered_joint_state(opt_jnames)
return opt_js
@@ -2601,6 +2057,19 @@ class MotionGen(MotionGenConfig):
start_state: Optional[JointState] = None,
goal_pose: Optional[Pose] = None,
) -> bool:
+ """Update the pose cost metric for :ref:`tut_constrained_planning`.
+
+ Only supports for the main end-effector. Does not support for multiple links that are
+ specified with `link_poses` in planning methods.
+
+ Args:
+ metric: Type and parameters for pose constraint to add.
+ start_state: Start joint state for the constraint.
+ goal_pose: Goal pose for the constraint.
+
+ Returns:
+ bool: True if the constraint can be added, False otherwise.
+ """
# check if constraint is valid:
if metric.hold_partial_pose and metric.offset_tstep_fraction < 0.0:
start_pose = self.compute_kinematics(start_state).ee_pose.clone()
@@ -2648,10 +2117,11 @@ class MotionGen(MotionGenConfig):
for rollout in rollouts
if isinstance(rollout, ArmReacher)
]
- torch.cuda.synchronize()
+ torch.cuda.synchronize(device=self.tensor_args.device)
return True
def get_all_rollout_instances(self) -> List[RolloutBase]:
+ """Get all rollout instances used across components in motion generation."""
if self._rollout_list is None:
self._rollout_list = (
self.ik_solver.get_all_rollout_instances()
@@ -2663,6 +2133,7 @@ class MotionGen(MotionGenConfig):
return self._rollout_list
def get_all_solver_rollout_instances(self) -> List[RolloutBase]:
+ """Get all rollout instances in solvers (IK, TrajOpt)."""
if self._solver_rollout_list is None:
self._solver_rollout_list = (
self.ik_solver.solver.get_all_rollout_instances()
@@ -2673,6 +2144,7 @@ class MotionGen(MotionGenConfig):
return self._solver_rollout_list
def get_all_pose_solver_rollout_instances(self) -> List[RolloutBase]:
+ """Get all rollout instances in solvers (IK, TrajOpt) that support Cartesian cost terms."""
if self._pose_solver_rollout_list is None:
self._pose_solver_rollout_list = (
self.ik_solver.solver.get_all_rollout_instances()
@@ -2682,6 +2154,14 @@ class MotionGen(MotionGenConfig):
return self._pose_solver_rollout_list
def get_all_kinematics_instances(self) -> List[CudaRobotModel]:
+ """Get all kinematics instances used across components in motion generation.
+
+ This is deprecated. Use :meth:`MotionGen.kinematics` instead as MotionGen now uses a shared
+ kinematics instance across all components.
+
+ Returns:
+ List[CudaRobotModel]: Single kinematics instance, returned as a list for compatibility.
+ """
if self._kin_list is None:
self._kin_list = [
i.dynamics_model.robot_model for i in self.get_all_rollout_instances()
@@ -2699,13 +2179,36 @@ class MotionGen(MotionGenConfig):
world_objects_pose_offset: Optional[Pose] = None,
remove_obstacles_from_world_config: bool = False,
) -> None:
- """Attach an object from world to robot's link.
+ """Attach an object or objects from world to a robot's link.
+
+ This method assumes that the objects exist in the world configuration. If attaching
+ objects that are not in world, use :meth:`MotionGen.attach_external_objects_to_robot`.
Args:
- surface_sphere_radius: _description_. Defaults to None.
- sphere_tensor: _description_. Defaults to None.
- link_name: _description_. Defaults to "attached_object".
+ joint_state: Joint state of the robot.
+ object_names: Names of objects in the world to attach to the robot.
+ surface_sphere_radius: Radius (in meters) to use for points sampled on surface of the
+ object. A smaller radius will allow for generating motions very close to obstacles.
+ link_name: Name of the link (frame) to attach the objects to. The assumption is that
+ this link does not have any geometry and all spheres of this link represent
+ attached objects.
+ sphere_fit_type: Sphere fit algorithm to use. See :ref:`attach_object_note` for more
+ details. The default method :attr:`SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE`
+ voxelizes the volume of the objects and adds spheres representing the voxels, then
+ samples points on the surface of the object, adds :attr:`surface_sphere_radius` to
+ these points. This should be used for most cases.
+ voxelize_method: Method to use for voxelization, passed to
+ :py:func:`trimesh.voxel.creation.voxelize`.
+ world_objects_pose_offset: Offset to apply to the object poses before attaching to the
+ robot. This is useful when attaching an object that's in contact with the world.
+ The offset is applied in the world frame before attaching to the robot.
+ remove_obstacles_from_world_config: Remove the obstacles from the world cache after
+ attaching to the robot to reduce memory usage. Note that when an object is attached
+ to the robot, it's disabled in the world collision checker. This flag when enabled,
+ also removes the object from world cache. For most cases, this should be set to
+ False.
"""
+
log_info("MG: Attach objects to robot")
kin_state = self.compute_kinematics(joint_state)
ee_pose = kin_state.ee_pose # w_T_ee
@@ -2765,12 +2268,26 @@ class MotionGen(MotionGenConfig):
voxelize_method: str = "ray",
world_objects_pose_offset: Optional[Pose] = None,
) -> None:
- """Attach an object from world to robot's link.
+ """Attach external objects (not in world model) to a robot's link.
Args:
- surface_sphere_radius: _description_. Defaults to None.
- sphere_tensor: _description_. Defaults to None.
- link_name: _description_. Defaults to "attached_object".
+ joint_state: Joint state of the robot.
+ external_objects: List of external objects to attach to the robot.
+ surface_sphere_radius: Radius (in meters) to use for points sampled on surface of the
+ object. A smaller radius will allow for generating motions very close to obstacles.
+ link_name: Name of the link (frame) to attach the objects to. The assumption is that
+ this link does not have any geometry and all spheres of this link represent
+ attached objects.
+ sphere_fit_type: Sphere fit algorithm to use. See :ref:`attach_object_note` for more
+ details. The default method :attr:`SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE`
+ voxelizes the volume of the objects and adds spheres representing the voxels, then
+ samples points on the surface of the object, adds :attr:`surface_sphere_radius` to
+ these points. This should be used for most cases.
+ voxelize_method: Method to use for voxelization, passed to
+ :py:func:`trimesh.voxel.creation.voxelize`.
+ world_objects_pose_offset: Offset to apply to the object poses before attaching to the
+ robot. This is useful when attaching an object that's in contact with the world.
+ The offset is applied in the world frame before attaching to the robot.
"""
log_info("MG: Attach objects to robot")
if len(external_objects) == 0:
@@ -2815,9 +2332,28 @@ class MotionGen(MotionGenConfig):
self.attach_spheres_to_robot(sphere_tensor=sphere_tensor, link_name=link_name)
def add_camera_frame(self, camera_observation: CameraObservation, obstacle_name: str):
+ """Add camera frame to the world collision checker.
+
+ Only supported by :py:class:`~curobo.geom.sdf.world_blox.WorldBloxCollision`.
+
+ Args:
+ camera_observation: Camera observation to add to the world collision checker.
+ obstacle_name: Name of the obstacle/layer to add the camera frame to.
+ """
self.world_coll_checker.add_camera_frame(camera_observation, obstacle_name)
def process_camera_frames(self, obstacle_name: Optional[str] = None, process_aux: bool = False):
+ """Process camera frames for collision checking.
+
+ Only supported by :py:class:`~curobo.geom.sdf.world_blox.WorldBloxCollision`.
+
+ Args:
+ obstacle_name: Name of the obstacle/layer to process the camera frames for. If None,
+ processes camera frames on all supported obstacles.
+ process_aux: Process auxillary information such as mesh integration in nvblox. This is
+ not required for collision checking and is only needed for debugging. Default is
+ False to reduce computation time.
+ """
self.world_coll_checker.process_camera_frames(obstacle_name, process_aux=process_aux)
def attach_bounding_box_from_blox_to_robot(
@@ -2833,18 +2369,11 @@ class MotionGen(MotionGenConfig):
):
"""Attach the voxels in a blox layer to robot's link.
- NOTE: This is not currently implemented.
+ .. note::
+ This is not currently implemented.
- Args:
- joint_state (JointState): _description_
- bounding_box (Cuboid): _description_
- blox_layer_name (Optional[str], optional): _description_. Defaults to None.
- surface_sphere_radius (float, optional): _description_. Defaults to 0.001.
- link_name (str, optional): _description_. Defaults to "attached_object".
- sphere_fit_type (SphereFitType, optional): _description_. Defaults to SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE.
- voxelize_method (str, optional): _description_. Defaults to "ray".
- world_objects_pose_offset (Optional[Pose], optional): _description_. Defaults to None.
"""
+ log_error("Not implemented")
kin_state = self.compute_kinematics(joint_state)
ee_pose = kin_state.ee_pose # w_T_ee
if world_objects_pose_offset is not None:
@@ -2895,54 +2424,27 @@ class MotionGen(MotionGenConfig):
):
"""Attach an object to robot's link. The object to be attached is not in the world model.
- Args:
- joint_state (JointState): _description_
- obstacle (Obstacle): _description_
- surface_sphere_radius (float, optional): _description_. Defaults to 0.001.
- link_name (str, optional): _description_. Defaults to "attached_object".
- sphere_fit_type (SphereFitType, optional): _description_. Defaults to SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE.
- voxelize_method (str, optional): _description_. Defaults to "ray".
- world_objects_pose_offset (Optional[Pose], optional): _description_. Defaults to None.
+ Deprecated. Use :meth:`MotionGen.attach_external_objects_to_robot` instead.
+
"""
- log_info("MG: Attach objects to robot")
- kin_state = self.compute_kinematics(joint_state)
- ee_pose = kin_state.ee_pose # w_T_ee
- if world_objects_pose_offset is not None:
- # add offset from ee:
- ee_pose = world_objects_pose_offset.inverse().multiply(ee_pose)
- # new ee_pose:
- # w_T_ee = offset_T_w * w_T_ee
- # ee_T_w
- ee_pose = ee_pose.inverse() # ee_T_w to multiply all objects later
- max_spheres = self.robot_cfg.kinematics.kinematics_config.get_number_of_spheres(link_name)
- n_spheres = max_spheres
- sphere_tensor = torch.zeros((max_spheres, 4))
- sphere_tensor[:, 3] = -10.0
- sph_list = []
- if n_spheres == 0:
- log_error("MG: No spheres found")
- sph = obstacle.get_bounding_spheres(
- n_spheres,
- surface_sphere_radius,
- pre_transform_pose=ee_pose,
- tensor_args=self.tensor_args,
- fit_type=sphere_fit_type,
+ log_warn("Deprecated. Use attach_external_objects_to_robot instead")
+ return self.attach_external_objects_to_robot(
+ joint_state=joint_state,
+ external_objects=[obstacle],
+ surface_sphere_radius=surface_sphere_radius,
+ link_name=link_name,
+ sphere_fit_type=sphere_fit_type,
voxelize_method=voxelize_method,
+ world_objects_pose_offset=world_objects_pose_offset,
)
- sph_list += [s.position + [s.radius] for s in sph]
-
- log_info("MG: Computed spheres for attach objects to robot")
-
- spheres = self.tensor_args.to_device(torch.as_tensor(sph_list))
-
- if spheres.shape[0] > max_spheres:
- spheres = spheres[: spheres.shape[0]]
- sphere_tensor[: spheres.shape[0], :] = spheres.contiguous()
-
- self.attach_spheres_to_robot(sphere_tensor=sphere_tensor, link_name=link_name)
def detach_object_from_robot(self, link_name: str = "attached_object") -> None:
- return self.detach_spheres_from_robot(link_name)
+ """Detach object from robot's link.
+
+ Args:
+ link_name: Name of the link.
+ """
+ self.detach_spheres_from_robot(link_name)
def attach_spheres_to_robot(
self,
@@ -2953,100 +2455,47 @@ class MotionGen(MotionGenConfig):
"""Attach spheres to robot's link.
Args:
- sphere_radius: _description_. Defaults to None.
- sphere_tensor: _description_. Defaults to None.
- link_name: _description_. Defaults to "attached_object".
+ sphere_radius: Radius of the spheres. Set to None if :attr:`sphere_tensor` is provided.
+ sphere_tensor: Sphere x, y, z, r tensor.
+ link_name: Name of the link to attach the spheres to. Note that this link should
+ already have pre-allocated spheres.
"""
self.robot_cfg.kinematics.kinematics_config.attach_object(
sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name
)
def detach_spheres_from_robot(self, link_name: str = "attached_object") -> None:
+ """Detach spheres from a robot's link.
+
+ Args:
+ link_name: Name of the link.
+ """
self.robot_cfg.kinematics.kinematics_config.detach_object(link_name)
def get_full_js(self, active_js: JointState) -> JointState:
+ """Get full joint state from controlled joint state, appending locked joints.
+
+ Args:
+ active_js: Controlled joint state
+
+ Returns:
+ JointState: Full joint state.
+ """
return self.rollout_fn.get_full_dof_from_solution(active_js)
- def add_running_pose_constraint(
- self,
- lock_x: bool = False,
- lock_y: bool = False,
- lock_z: bool = False,
- lock_rx: bool = False,
- lock_ry: bool = False,
- lock_rz: bool = False,
- ):
- raise NotImplementedError()
-
- def remove_running_pose_constraint(self):
- raise NotImplementedError()
-
- def run_finetune_trajopt(
- self,
- start_state: JointState,
- goal_pose: Pose,
- traj_solution: JointState,
- traj_dt: Union[float, torch.Tensor, None],
- max_attempts: int = 1,
- ):
- # NOTE: Currently only works for single environment. Need to rewrite for all modes
- # finetunes solution
- if traj_dt is not None:
- self.finetune_trajopt_solver.update_solver_dt(traj_dt.item())
-
- # call trajopt with seed:
-
- # NOTE: currently runs only one seed
- time_dict = {
- "solve_time": 0,
- "ik_time": 0,
- "graph_time": 0,
- "trajopt_time": 0,
- "trajopt_attempts": 0,
- }
- result = None
- # goal_state = JointState.from_position(
- # traj_solution.position[..., -2:-1, :].clone(), joint_names=start_state.joint_names
- # )
- goal = Goal(
- goal_pose=goal_pose,
- # goal_state=goal_state,
- current_state=start_state,
- )
-
- for n in range(max_attempts):
- traj_result = self.finetune_trajopt_solver.solve_single(goal, traj_solution)
- time_dict["trajopt_time"] += traj_result.solve_time
- time_dict["trajopt_attempts"] = n
- if result is None:
- result = MotionGenResult(success=traj_result.success)
-
- if traj_result.success.item():
- break
-
- if self.store_debug_in_result:
- result.debug_info = {"trajopt_result": traj_result}
-
- result.position_error = traj_result.position_error
- result.rotation_error = traj_result.rotation_error
- result.cspace_error = traj_result.cspace_error
- result.optimized_dt = traj_result.optimized_dt
- result.interpolated_plan = traj_result.interpolated_solution
- result.optimized_plan = traj_result.solution
- result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
- result.trajopt_time = time_dict["trajopt_time"]
- return result
-
@property
def world_model(self) -> WorldConfig:
+ """Get the world model used for collision checking."""
return self.world_coll_checker.world_model
@property
def world_collision(self) -> WorldCollision:
+ """Get the shared instance of world collision checker."""
return self.world_coll_checker
@property
def project_pose_to_goal_frame(self) -> bool:
+ """Check if the pose cost metric is projected to goal frame."""
return self.trajopt_solver.rollout_fn.goal_cost.project_distance
def update_interpolation_type(
@@ -3055,6 +2504,13 @@ class MotionGen(MotionGenConfig):
update_graph: bool = True,
update_trajopt: bool = True,
):
+ """Update interpolation type for motion generation.
+
+ Args:
+ interpolation_type: Type of interpolation to use.
+ update_graph: Update graph planner with the new interpolation type.
+ update_trajopt: Update trajectory optimization solvers with the new interpolation type.
+ """
if update_graph:
self.graph_planner.interpolation_type = interpolation_type
if update_trajopt:
@@ -3065,6 +2521,17 @@ class MotionGen(MotionGenConfig):
def update_locked_joints(
self, lock_joints: Dict[str, float], robot_config_dict: Union[str, Dict[Any]]
):
+ """Update locked joints in the robot configuration.
+
+ Use this function to update the joint values of currently locked joints between
+ planning calls. This function can also be used to change which joints are locked, however
+ this is only supported when the number of locked joints is the same as the original
+ robot configuration as the kinematics tensors are pre-allocated.
+
+ Args:
+ lock_joints: Dictionary of joint names and values to lock.
+ robot_config_dict: Robot configuration dictionary or path to robot configuration file.
+ """
if isinstance(robot_config_dict, str):
robot_config_dict = load_yaml(join_path(get_robot_configs_path(), robot_config_dict))[
"robot_cfg"
@@ -3073,11 +2540,1358 @@ class MotionGen(MotionGenConfig):
robot_config_dict = robot_config_dict["robot_cfg"]
robot_config_dict["kinematics"]["lock_joints"] = lock_joints
robot_cfg = RobotConfig.from_dict(robot_config_dict, self.tensor_args)
-
- # make sure the new robot config and current have the same joint limits:
- new_joint_limits = robot_cfg.kinematics.get_joint_limits()
- current_joint_limits = self.robot_cfg.kinematics.get_joint_limits()
- # if new_joint_limits != current_joint_limits:
- # log_error("Joint limits are different, reinstance motion gen")
-
self.kinematics.update_kinematics_config(robot_cfg.kinematics.kinematics_config)
+
+ @profiler.record_function("motion_gen/ik")
+ def _solve_ik_from_solve_state(
+ self,
+ goal_pose: Pose,
+ solve_state: ReacherSolveState,
+ start_state: JointState,
+ use_nn_seed: bool,
+ partial_ik_opt: bool,
+ link_poses: Optional[Dict[str, Pose]] = None,
+ ) -> IKResult:
+ """Solve inverse kinematics from solve state, used by motion generation planning call.
+
+ Args:
+ goal_pose: Goal Pose for the end-effector.
+ solve_state: Solve state for motion generation.
+ start_state: Start joint configuration of the robot.
+ use_nn_seed: Use seed from a neural network. Not implemented.
+ partial_ik_opt: Only run 50 iterations of inverse kinematics.
+ link_poses: Goal Poses of any other link in the robot that was specified in
+ :meth:`curobo.types.robot.RobotConfig.kinematics.link_names`.
+
+ Returns:
+ IKResult: Result of inverse kinematics.
+ """
+ newton_iters = None
+ if partial_ik_opt:
+ newton_iters = self.partial_ik_iters
+ ik_result = self.ik_solver.solve_any(
+ solve_state.solve_type,
+ goal_pose,
+ start_state.position.view(-1, self._dof),
+ start_state.position.view(-1, 1, self._dof),
+ solve_state.num_trajopt_seeds,
+ solve_state.num_ik_seeds,
+ use_nn_seed,
+ newton_iters,
+ link_poses,
+ )
+ return ik_result
+
+ @profiler.record_function("motion_gen/trajopt_solve")
+ def _solve_trajopt_from_solve_state(
+ self,
+ goal: Goal,
+ solve_state: ReacherSolveState,
+ act_seed: Optional[JointState] = None,
+ use_nn_seed: bool = False,
+ return_all_solutions: bool = False,
+ seed_success: Optional[torch.Tensor] = None,
+ newton_iters: Optional[int] = None,
+ trajopt_instance: Optional[TrajOptSolver] = None,
+ num_seeds_override: Optional[int] = None,
+ ) -> TrajOptResult:
+ """Solve trajectory optimization from solve state, used by motion generation planning call.
+
+ Args:
+ goal: Goal containing Pose/Joint targets, current robot state and any other information.
+ solve_state: Solve state for motion generation.
+ act_seed: Seed to run trajectory optimization.
+ use_nn_seed: Use neural network seed for solving trajectory optimization. This is not
+ implemented.
+ return_all_solutions: Return all solutions found by trajectory optimization.
+ seed_success: Success tensor for seeds.
+ newton_iters: Override Newton iterations to run for solving trajectory optimization.
+ trajopt_instance: Instance of TrajOptSolver to use for solving trajectory optimization.
+ num_seeds_override: Override number of seeds to use for solving trajectory optimization.
+
+ Returns:
+ TrajOptResult: Result of trajectory optimization.
+ """
+ if trajopt_instance is None:
+ trajopt_instance = self.trajopt_solver
+ if num_seeds_override is None:
+ num_seeds_override = solve_state.num_trajopt_seeds
+ traj_result = trajopt_instance.solve_any(
+ solve_state.solve_type,
+ goal,
+ act_seed,
+ use_nn_seed,
+ return_all_solutions,
+ num_seeds_override,
+ seed_success,
+ newton_iters=newton_iters,
+ )
+ return traj_result
+
+ def _get_solve_state(
+ self,
+ solve_type: ReacherSolveType,
+ plan_config: MotionGenPlanConfig,
+ goal_pose: Pose,
+ start_state: JointState,
+ ) -> ReacherSolveState:
+ """Generate solve state for motion generation based on planning type and configuration.
+
+ MotionGen creates a :class:`ReacherSolveState` for every planning call to keep track of
+ planning parameters such as number of seeds, batch size, solve type, etc. This solve state
+ is then compared with existing solve state to determine if solvers (IK, TrajOpt) need to be
+ re-initialized. Note that changing solve state is not supported when
+ :attr:`MotionGen.use_cuda_graph` is enabled.
+
+ Args:
+ solve_type: Type of reacher problem to solve.
+ plan_config: Planning configuration for motion generation.
+ goal_pose: Goal Pose for the end-effector.
+ start_state: Start joint configuration of the robot.
+
+ Raises:
+ ValueError: If the solve type is not supported.
+
+ Returns:
+ ReacherSolveState: Solve state for motion generation.
+ """
+
+ num_ik_seeds = (
+ self.ik_seeds if plan_config.num_ik_seeds is None else plan_config.num_ik_seeds
+ )
+ num_trajopt_seeds = (
+ self.trajopt_seeds
+ if plan_config.num_trajopt_seeds is None
+ else plan_config.num_trajopt_seeds
+ )
+
+ num_graph_seeds = (
+ self.trajopt_seeds if plan_config.num_graph_seeds is None else num_trajopt_seeds
+ )
+ solve_state = None
+ if solve_type == ReacherSolveType.SINGLE:
+ solve_state = ReacherSolveState(
+ solve_type,
+ num_ik_seeds=num_ik_seeds,
+ num_trajopt_seeds=num_trajopt_seeds,
+ num_graph_seeds=num_graph_seeds,
+ batch_size=1,
+ n_envs=1,
+ n_goalset=1,
+ )
+ elif solve_type == ReacherSolveType.GOALSET:
+ solve_state = ReacherSolveState(
+ solve_type,
+ num_ik_seeds=num_ik_seeds,
+ num_trajopt_seeds=num_trajopt_seeds,
+ num_graph_seeds=num_graph_seeds,
+ batch_size=1,
+ n_envs=1,
+ n_goalset=goal_pose.n_goalset,
+ )
+ elif solve_type == ReacherSolveType.BATCH:
+ solve_state = ReacherSolveState(
+ solve_type,
+ num_ik_seeds=num_ik_seeds,
+ num_trajopt_seeds=num_trajopt_seeds,
+ num_graph_seeds=num_graph_seeds,
+ batch_size=goal_pose.batch,
+ n_envs=1,
+ n_goalset=1,
+ )
+ elif solve_type == ReacherSolveType.BATCH_GOALSET:
+ solve_state = ReacherSolveState(
+ solve_type,
+ num_ik_seeds=num_ik_seeds,
+ num_trajopt_seeds=num_trajopt_seeds,
+ num_graph_seeds=num_graph_seeds,
+ batch_size=goal_pose.batch,
+ n_envs=1,
+ n_goalset=goal_pose.n_goalset,
+ )
+ elif solve_type == ReacherSolveType.BATCH_ENV:
+ solve_state = ReacherSolveState(
+ solve_type,
+ num_ik_seeds=num_ik_seeds,
+ num_trajopt_seeds=num_trajopt_seeds,
+ num_graph_seeds=num_graph_seeds,
+ batch_size=goal_pose.batch,
+ n_envs=goal_pose.batch,
+ n_goalset=1,
+ )
+ elif solve_type == ReacherSolveType.BATCH_ENV_GOALSET:
+ solve_state = ReacherSolveState(
+ solve_type,
+ num_ik_seeds=num_ik_seeds,
+ num_trajopt_seeds=num_trajopt_seeds,
+ num_graph_seeds=num_graph_seeds,
+ batch_size=goal_pose.batch,
+ n_envs=goal_pose.batch,
+ n_goalset=goal_pose.n_goalset,
+ )
+ else:
+ raise ValueError("Unsupported Solve type")
+ return solve_state
+
+ def _plan_attempts(
+ self,
+ solve_state: ReacherSolveState,
+ start_state: JointState,
+ goal_pose: Pose,
+ plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: List[Pose] = None,
+ ):
+ """Call many planning attempts for a given reacher solve state.
+
+ Args:
+ solve_state: Reacher solve state for planning.
+ start_state: Start joint state for planning.
+ goal_pose: Goal pose to reach for end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for other links in the robot.
+
+ Returns:
+ MotionGenResult: Result of planning.
+ """
+ start_time = time.time()
+ if plan_config.pose_cost_metric is not None:
+ valid_query = self.update_pose_cost_metric(
+ plan_config.pose_cost_metric, start_state, goal_pose
+ )
+ if not valid_query:
+ result = MotionGenResult(
+ success=torch.as_tensor([False], device=self.tensor_args.device),
+ valid_query=valid_query,
+ status="Invalid Hold Partial Pose",
+ )
+ return result
+ self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size)
+ if solve_state.batch_env:
+ if solve_state.batch_size > self.world_coll_checker.n_envs:
+ log_error("Batch Env is less that goal batch")
+ force_graph = plan_config.enable_graph
+ partial_ik = plan_config.partial_ik_opt
+ ik_fail_count = 0
+ time_dict = {
+ "solve_time": 0,
+ "ik_time": 0,
+ "graph_time": 0,
+ "trajopt_time": 0,
+ "trajopt_attempts": 0,
+ }
+ best_status = 0
+ if plan_config.finetune_dt_scale is None:
+ plan_config.finetune_dt_scale = self.finetune_dt_scale
+ for n in range(plan_config.max_attempts):
+ log_info("MG Iter: " + str(n))
+ result = self._plan_from_solve_state(
+ solve_state,
+ start_state,
+ goal_pose,
+ plan_config,
+ link_poses,
+ )
+ time_dict["solve_time"] += result.solve_time
+ time_dict["ik_time"] += result.ik_time
+ time_dict["graph_time"] += result.graph_time
+ time_dict["trajopt_time"] += result.trajopt_time
+ time_dict["trajopt_attempts"] += result.trajopt_attempts
+ if (
+ result.status == MotionGenStatus.IK_FAIL and plan_config.ik_fail_return is not None
+ ): # IF IK fails the first time, we exist assuming the goal is not reachable
+ ik_fail_count += 1
+ best_status = max(best_status, 1)
+
+ if ik_fail_count > plan_config.ik_fail_return:
+ break
+ if result.success[0].item():
+ break
+
+ if result.status == MotionGenStatus.FINETUNE_TRAJOPT_FAIL:
+ plan_config.finetune_dt_scale *= (
+ plan_config.finetune_dt_decay**plan_config.finetune_attempts
+ )
+ plan_config.finetune_dt_scale = min(plan_config.finetune_dt_scale, 1.25)
+ if plan_config.enable_graph_attempt is not None and (
+ n >= plan_config.enable_graph_attempt - 1
+ and result.status
+ in [MotionGenStatus.TRAJOPT_FAIL, MotionGenStatus.FINETUNE_TRAJOPT_FAIL]
+ and not plan_config.enable_graph
+ ):
+ plan_config.enable_graph = True
+ plan_config.partial_ik_opt = False
+ if plan_config.disable_graph_attempt is not None and (
+ n >= plan_config.disable_graph_attempt - 1
+ and result.status
+ in [
+ MotionGenStatus.TRAJOPT_FAIL,
+ MotionGenStatus.GRAPH_FAIL,
+ MotionGenStatus.FINETUNE_TRAJOPT_FAIL,
+ ]
+ and not force_graph
+ ):
+ plan_config.enable_graph = False
+ plan_config.partial_ik_opt = partial_ik
+ if result.status in [MotionGenStatus.TRAJOPT_FAIL]:
+ best_status = 3
+ elif result.status in [MotionGenStatus.GRAPH_FAIL]:
+ best_status = 2
+ if time.time() - start_time > plan_config.timeout:
+ break
+ if not result.valid_query:
+ result.status = MotionGenStatus.INVALID_QUERY
+ break
+ if n == 10:
+ self.reset_seed()
+ log_warn("Couldn't find solution with 10 attempts, resetting seeds")
+
+ result.solve_time = time_dict["solve_time"]
+ result.ik_time = time_dict["ik_time"]
+ result.graph_time = time_dict["graph_time"]
+ result.trajopt_time = time_dict["trajopt_time"]
+ result.trajopt_attempts = time_dict["trajopt_attempts"]
+ result.attempts = n + 1
+ torch.cuda.synchronize(device=self.tensor_args.device)
+ if plan_config.pose_cost_metric is not None:
+ self.update_pose_cost_metric(PoseCostMetric.reset_metric())
+ if plan_config.time_dilation_factor is not None and torch.count_nonzero(result.success) > 0:
+ result.retime_trajectory(
+ plan_config.time_dilation_factor,
+ interpolation_kind=self.finetune_trajopt_solver.interpolation_type,
+ )
+
+ result.total_time = time.time() - start_time
+ return result
+
+ def _plan_batch_attempts(
+ self,
+ solve_state: ReacherSolveState,
+ start_state: JointState,
+ goal_pose: Pose,
+ plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Optional[Dict[str, Pose]] = None,
+ ):
+ """Plan batch attempts for a given reacher solve state.
+
+ Args:
+ solve_state: Reacher solve state for planning.
+ start_state: Start joint state for planning.
+ goal_pose: Goal pose to reach for end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for other links in the robot.
+
+ Returns:
+ MotionGenResult: Result of batched planning.
+ """
+ start_time = time.time()
+ goal_pose = goal_pose.clone()
+ if plan_config.pose_cost_metric is not None:
+ valid_query = self.update_pose_cost_metric(
+ plan_config.pose_cost_metric, start_state, goal_pose
+ )
+ if not valid_query:
+ result = MotionGenResult(
+ success=torch.as_tensor(
+ [False for _ in solve_state.batch_size],
+ device=self.motion_gen.tensor_args.device,
+ ),
+ valid_query=valid_query,
+ status="Invalid Hold Partial Pose",
+ )
+ return result
+
+ if solve_state.batch_env:
+ if solve_state.batch_size > self.world_coll_checker.n_envs:
+ log_error("Batch Env is less that goal batch")
+ if plan_config.enable_graph:
+ log_error("Graph Search / Geometric Planner not supported in batch_env mode")
+
+ if plan_config.enable_graph or (
+ plan_config.enable_graph_attempt is not None
+ and plan_config.max_attempts >= plan_config.enable_graph_attempt
+ ):
+ log_warn("Batch mode enable graph is only supported with num_graph_seeds==1")
+ plan_config.num_trajopt_seeds = 1
+ plan_config.num_graph_seeds = 1
+ solve_state.num_trajopt_seeds = 1
+ solve_state.num_graph_seeds = 1
+ self.update_batch_size(seeds=solve_state.num_trajopt_seeds, batch=solve_state.batch_size)
+
+ ik_fail_count = 0
+ force_graph = plan_config.enable_graph
+ partial_ik = plan_config.partial_ik_opt
+
+ time_dict = {
+ "solve_time": 0,
+ "ik_time": 0,
+ "graph_time": 0,
+ "trajopt_time": 0,
+ "trajopt_attempts": 0,
+ }
+ best_result = None
+
+ for n in range(plan_config.max_attempts):
+ result = self._plan_from_solve_state_batch(
+ solve_state,
+ start_state,
+ goal_pose,
+ plan_config,
+ link_poses=link_poses,
+ )
+
+ time_dict["solve_time"] += result.solve_time
+ time_dict["ik_time"] += result.ik_time
+
+ time_dict["graph_time"] += result.graph_time
+ time_dict["trajopt_time"] += result.trajopt_time
+ time_dict["trajopt_attempts"] += result.trajopt_attempts
+
+ # if not all have succeeded, store the successful ones and re attempt:
+ # TODO: update stored based on error
+ if best_result is None:
+ best_result = result.clone()
+ else:
+ # get success idx:
+ idx = torch.nonzero(result.success).reshape(-1)
+ if len(idx) > 0:
+ best_result.copy_idx(idx, result)
+
+ if (
+ result.status == MotionGenStatus.IK_FAIL and plan_config.ik_fail_return is not None
+ ): # IF IK fails the first time, we exit assuming the goal is not reachable
+ ik_fail_count += 1
+ if ik_fail_count > plan_config.ik_fail_return:
+ break
+
+ if (
+ torch.count_nonzero(best_result.success)
+ >= goal_pose.batch * plan_config.success_ratio
+ ): # we want 90% targets to succeed
+ best_result.status = None
+ break
+ if plan_config.enable_graph_attempt is not None and (
+ n >= plan_config.enable_graph_attempt - 1
+ and result.status != MotionGenStatus.IK_FAIL
+ and not plan_config.enable_graph
+ ):
+ plan_config.enable_graph = True
+ plan_config.partial_ik_opt = False
+
+ if plan_config.disable_graph_attempt is not None and (
+ n >= plan_config.disable_graph_attempt - 1
+ and result.status in [MotionGenStatus.TRAJOPT_FAIL, MotionGenStatus.GRAPH_FAIL]
+ and not force_graph
+ ):
+ plan_config.enable_graph = False
+ plan_config.partial_ik_opt = partial_ik
+
+ if plan_config.fail_on_invalid_query:
+ if not result.valid_query:
+ best_result.valid_query = False
+ best_result.status = "Invalid Problem"
+ break
+ if time.time() - start_time > plan_config.timeout:
+ break
+ best_result.solve_time = time_dict["solve_time"]
+ best_result.ik_time = time_dict["ik_time"]
+ best_result.graph_time = time_dict["graph_time"]
+ best_result.trajopt_time = time_dict["trajopt_time"]
+ best_result.trajopt_attempts = time_dict["trajopt_attempts"]
+ best_result.attempts = n + 1
+ torch.cuda.synchronize(device=self.tensor_args.device)
+ if plan_config.pose_cost_metric is not None:
+ self.update_pose_cost_metric(PoseCostMetric.reset_metric())
+
+ if plan_config.time_dilation_factor is not None and torch.count_nonzero(result.success) > 0:
+ result.retime_trajectory(
+ plan_config.time_dilation_factor,
+ interpolation_kind=self.finetune_trajopt_solver.interpolation_type,
+ )
+ best_result.total_time = time.time() - start_time
+ return best_result
+
+ def _plan_from_solve_state(
+ self,
+ solve_state: ReacherSolveState,
+ start_state: JointState,
+ goal_pose: Pose,
+ plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Optional[Dict[str, Pose]] = None,
+ ) -> MotionGenResult:
+ """Plan from a given reacher solve state.
+
+ Args:
+ solve_state: Reacher solve state for planning.
+ start_state: Start joint state for planning.
+ goal_pose: Goal pose to reach for end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for other links in the robot.
+
+ Returns:
+ MotionGenResult: Result of planning.
+ """
+ trajopt_seed_traj = None
+ trajopt_seed_success = None
+ trajopt_newton_iters = None
+ graph_success = 0
+
+ if len(start_state.shape) == 1:
+ log_error("Joint state should be not a vector (dof) should be (bxdof)")
+
+ if goal_pose.shape[0] != 1:
+ log_error(
+ "Goal position should be of shape [1, n_goalset, -1], current shape: "
+ + str(goal_pose.shape)
+ )
+ # plan ik:
+
+ ik_result = self._solve_ik_from_solve_state(
+ goal_pose,
+ solve_state,
+ start_state,
+ plan_config.use_nn_ik_seed,
+ plan_config.partial_ik_opt,
+ link_poses,
+ )
+
+ if not plan_config.enable_graph and plan_config.partial_ik_opt:
+ ik_result.success[:] = True
+
+ # check for success:
+ result = MotionGenResult(
+ ik_result.success.view(-1)[0:1], # This is not true for batch mode
+ ik_time=ik_result.solve_time,
+ solve_time=ik_result.solve_time,
+ )
+
+ if self.store_debug_in_result:
+ result.debug_info = {"ik_result": ik_result}
+ ik_success = torch.count_nonzero(ik_result.success)
+ if ik_success == 0:
+ result.status = MotionGenStatus.IK_FAIL
+ return result
+
+ # do graph search:
+ with profiler.record_function("motion_gen/post_ik"):
+ ik_out_seeds = solve_state.num_trajopt_seeds
+ if plan_config.enable_graph:
+ ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success)
+
+ goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof)[
+ :ik_success
+ ]
+ start_config = tensor_repeat_seeds(start_state.position, ik_out_seeds)
+ if plan_config.enable_opt:
+ self._trajopt_goal_config[:] = ik_result.solution
+
+ # do graph search:
+ if plan_config.enable_graph:
+ interpolation_steps = None
+ if plan_config.enable_opt:
+ interpolation_steps = self.trajopt_solver.action_horizon
+ log_info("MG: running GP")
+ graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
+ trajopt_seed_success = graph_result.success
+
+ graph_success = torch.count_nonzero(graph_result.success).item()
+ result.graph_time = graph_result.solve_time
+ result.solve_time += graph_result.solve_time
+ if graph_success > 0:
+ log_info("MG: GP Success")
+ result.graph_plan = graph_result.interpolated_plan
+ result.interpolated_plan = graph_result.interpolated_plan
+
+ result.used_graph = True
+ if plan_config.enable_opt:
+ trajopt_seed = (
+ result.graph_plan.position.view(
+ 1, # solve_state.batch_size,
+ graph_success, # solve_state.num_trajopt_seeds,
+ interpolation_steps,
+ self._dof,
+ )
+ .transpose(0, 1)
+ .contiguous()
+ )
+ trajopt_seed_traj = torch.zeros(
+ (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof),
+ device=self.tensor_args.device,
+ dtype=self.tensor_args.dtype,
+ )
+ trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed
+ trajopt_seed_success = ik_result.success.clone()
+ trajopt_seed_success[ik_result.success] = graph_result.success
+
+ trajopt_seed_success = trajopt_seed_success.view(
+ solve_state.batch_size, solve_state.num_trajopt_seeds
+ )
+ trajopt_newton_iters = self.graph_trajopt_iters
+ else:
+ _, idx = torch.topk(
+ graph_result.path_length[graph_result.success], k=1, largest=False
+ )
+ result.interpolated_plan = result.interpolated_plan[idx].squeeze(0)
+ result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt)
+ result.optimized_plan = result.interpolated_plan[
+ : graph_result.path_buffer_last_tstep[idx.item()]
+ ]
+ idx = idx.view(-1) + self._batch_col
+ result.position_error = ik_result.position_error[ik_result.success][
+ graph_result.success
+ ][idx]
+ result.rotation_error = ik_result.rotation_error[ik_result.success][
+ graph_result.success
+ ][idx]
+ result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[
+ idx.item() : idx.item() + 1
+ ]
+ result.success = result.success.view(-1)[0:1]
+ result.success[:] = True
+ return result
+ else:
+ result.success[:] = False
+ result.status = MotionGenStatus.GRAPH_FAIL
+ if not graph_result.valid_query:
+ result.valid_query = False
+ if self.store_debug_in_result:
+ result.debug_info["graph_debug"] = graph_result.debug_info
+ return result
+ if plan_config.need_graph_success:
+ return result
+
+ # do trajopt::
+
+ if plan_config.enable_opt:
+ with profiler.record_function("motion_gen/setup_trajopt_seeds"):
+ self._trajopt_goal_config[:, :ik_success] = goal_config
+ retract_config = None
+ if plan_config.use_start_state_as_retract:
+ retract_config = start_state.position.clone()
+ goal = Goal(
+ goal_pose=goal_pose,
+ current_state=start_state,
+ links_goal_pose=link_poses,
+ retract_state=retract_config,
+ )
+
+ if (
+ trajopt_seed_traj is None
+ or graph_success < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
+ ):
+ goal_config = self._trajopt_goal_config[0] # batch index == 0
+
+ goal_state = JointState.from_position(
+ goal_config,
+ )
+ seed_link_poses = None
+ if link_poses is not None:
+ seed_link_poses = {}
+
+ for k in link_poses.keys():
+ seed_link_poses[k] = link_poses[k].repeat_seeds(
+ solve_state.num_trajopt_seeds
+ )
+ if goal_pose.shape[0] != 1:
+ log_error(
+ "Batch size of goal pose should be 1, current shape: "
+ + str(goal_pose.shape)
+ )
+ seed_goal = Goal(
+ goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds),
+ current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
+ goal_state=goal_state,
+ links_goal_pose=seed_link_poses,
+ )
+ if trajopt_seed_traj is not None:
+ trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
+ # batch, num_seeds, h, dof
+ if (
+ trajopt_seed_success.shape[1]
+ < solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
+ ):
+ trajopt_seed_success_new = torch.zeros(
+ (1, solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds),
+ device=self.tensor_args.device,
+ dtype=torch.bool,
+ )
+ trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = (
+ trajopt_seed_success
+ )
+ trajopt_seed_success = trajopt_seed_success_new
+ # create seeds here:
+ trajopt_seed_traj = self.trajopt_solver.get_seed_set(
+ seed_goal,
+ trajopt_seed_traj, # batch, num_seeds, h, dof
+ num_seeds=self.noisy_trajopt_seeds,
+ batch_mode=solve_state.batch_mode,
+ seed_success=trajopt_seed_success,
+ )
+ trajopt_seed_traj = trajopt_seed_traj.view(
+ solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
+ solve_state.batch_size,
+ self.trajopt_solver.action_horizon,
+ self._dof,
+ ).contiguous()
+ if plan_config.enable_finetune_trajopt:
+ og_value = self.trajopt_solver.interpolation_type
+ self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
+ with profiler.record_function("motion_gen/trajopt"):
+ log_info("MG: running TO")
+ traj_result = self._solve_trajopt_from_solve_state(
+ goal,
+ solve_state,
+ trajopt_seed_traj,
+ num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds,
+ newton_iters=trajopt_newton_iters,
+ return_all_solutions=plan_config.parallel_finetune
+ and plan_config.enable_finetune_trajopt,
+ )
+ if plan_config.enable_finetune_trajopt:
+ self.trajopt_solver.interpolation_type = og_value
+ if self.store_debug_in_result:
+ result.debug_info["trajopt_result"] = traj_result
+ # run finetune
+ if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
+ with profiler.record_function("motion_gen/finetune_trajopt"):
+ seed_traj = traj_result.raw_action.clone() # solution.position.clone()
+ seed_traj = seed_traj.contiguous()
+ og_solve_time = traj_result.solve_time
+ seed_override = 1
+ opt_dt = traj_result.optimized_dt
+
+ if plan_config.parallel_finetune:
+ opt_dt = torch.min(opt_dt[traj_result.success])
+ seed_override = solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds
+
+ finetune_time = 0
+ for k in range(plan_config.finetune_attempts):
+ newton_iters = None
+
+ scaled_dt = torch.clamp(
+ opt_dt
+ * plan_config.finetune_dt_scale
+ * (plan_config.finetune_dt_decay ** (k)),
+ self.trajopt_solver.minimum_trajectory_dt,
+ )
+ if self.optimize_dt:
+ self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item())
+
+ traj_result = self._solve_trajopt_from_solve_state(
+ goal,
+ solve_state,
+ seed_traj,
+ trajopt_instance=self.finetune_trajopt_solver,
+ num_seeds_override=seed_override,
+ newton_iters=newton_iters,
+ )
+ finetune_time += traj_result.solve_time
+ if torch.count_nonzero(traj_result.success) > 0:
+ break
+ seed_traj = traj_result.optimized_seeds.detach().clone()
+ newton_iters = 4
+
+ traj_result.solve_time = finetune_time
+
+ result.finetune_time = traj_result.solve_time
+
+ traj_result.solve_time = og_solve_time
+ if self.store_debug_in_result:
+ result.debug_info["finetune_trajopt_result"] = traj_result
+ elif plan_config.enable_finetune_trajopt:
+ traj_result.success = traj_result.success[0:1]
+ # if torch.count_nonzero(result.success) == 0:
+ result.status = MotionGenStatus.TRAJOPT_FAIL
+ result.solve_time += traj_result.solve_time + result.finetune_time
+ result.trajopt_time = traj_result.solve_time
+ result.trajopt_attempts = 1
+ result.success = traj_result.success
+
+ if plan_config.enable_finetune_trajopt and torch.count_nonzero(result.success) == 0:
+ result.status = MotionGenStatus.FINETUNE_TRAJOPT_FAIL
+
+ result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory(
+ 0, traj_result.path_buffer_last_tstep[0]
+ )
+ result.interpolation_dt = self.trajopt_solver.interpolation_dt
+ result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
+ result.position_error = traj_result.position_error
+ result.rotation_error = traj_result.rotation_error
+ result.optimized_dt = traj_result.optimized_dt
+ result.optimized_plan = traj_result.solution
+ result.goalset_index = traj_result.goalset_index
+ return result
+
+ def _plan_js_from_solve_state(
+ self,
+ solve_state: ReacherSolveState,
+ start_state: JointState,
+ goal_state: JointState,
+ plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ ) -> MotionGenResult:
+ """Plan from a given reacher solve state for joint state.
+
+ Args:
+ solve_state: Reacher solve state for planning.
+ start_state: Start joint state for planning.
+ goal_state: Goal joint state to reach.
+ plan_config: Planning parameters for motion generation.
+
+ Returns:
+ MotionGenResult: Result of planning.
+ """
+ trajopt_seed_traj = None
+ trajopt_seed_success = None
+ trajopt_newton_iters = self.js_trajopt_solver.newton_iters
+
+ graph_success = 0
+ if len(start_state.shape) == 1:
+ log_error("Joint state should be not a vector (dof) should be (bxdof)")
+
+ result = MotionGenResult(cspace_error=torch.zeros((1), device=self.tensor_args.device))
+ if self.store_debug_in_result:
+ result.debug_info = {}
+ # do graph search:
+ if plan_config.enable_graph:
+ start_config = torch.zeros(
+ (solve_state.num_graph_seeds, self.js_trajopt_solver.dof),
+ device=self.tensor_args.device,
+ dtype=self.tensor_args.dtype,
+ )
+ goal_config = start_config.clone()
+ start_config[:] = start_state.position
+ goal_config[:] = goal_state.position
+ interpolation_steps = None
+ if plan_config.enable_opt:
+ interpolation_steps = self.js_trajopt_solver.action_horizon
+ log_info("MG: running GP")
+ graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
+ trajopt_seed_success = graph_result.success
+
+ graph_success = torch.count_nonzero(graph_result.success).item()
+ result.graph_time = graph_result.solve_time
+ result.solve_time += graph_result.solve_time
+ if graph_success > 0:
+ result.graph_plan = graph_result.interpolated_plan
+ result.interpolated_plan = graph_result.interpolated_plan
+
+ result.used_graph = True
+ if plan_config.enable_opt:
+ trajopt_seed = (
+ result.graph_plan.position.view(
+ 1, # solve_state.batch_size,
+ graph_success, # solve_state.num_trajopt_seeds,
+ interpolation_steps,
+ self._dof,
+ )
+ .transpose(0, 1)
+ .contiguous()
+ )
+ trajopt_seed_traj = torch.zeros(
+ (trajopt_seed.shape[0], 1, self.trajopt_solver.action_horizon, self._dof),
+ device=self.tensor_args.device,
+ dtype=self.tensor_args.dtype,
+ )
+ trajopt_seed_traj[:, :, :interpolation_steps, :] = trajopt_seed
+ trajopt_seed_success = graph_result.success
+
+ trajopt_seed_success = trajopt_seed_success.view(
+ 1, solve_state.num_trajopt_seeds
+ )
+ trajopt_newton_iters = self.graph_trajopt_iters
+ else:
+ _, idx = torch.topk(
+ graph_result.path_length[graph_result.success], k=1, largest=False
+ )
+ result.interpolated_plan = result.interpolated_plan[idx].squeeze(0)
+ result.optimized_dt = self.tensor_args.to_device(self.interpolation_dt)
+ result.optimized_plan = result.interpolated_plan[
+ : graph_result.path_buffer_last_tstep[idx.item()]
+ ]
+ idx = idx.view(-1) + self._batch_col
+ result.cspace_error = torch.zeros((1), device=self.tensor_args.device)
+
+ result.path_buffer_last_tstep = graph_result.path_buffer_last_tstep[
+ idx.item() : idx.item() + 1
+ ]
+ result.success = torch.as_tensor([True], device=self.tensor_args.device)
+ return result
+ else:
+ result.success = torch.as_tensor([False], device=self.tensor_args.device)
+ result.status = MotionGenStatus.GRAPH_FAIL
+ if not graph_result.valid_query:
+ result.valid_query = False
+ if self.store_debug_in_result:
+ result.debug_info["graph_debug"] = graph_result.debug_info
+ return result
+ if plan_config.need_graph_success:
+ return result
+
+ # do trajopt:
+ if plan_config.enable_opt:
+ with profiler.record_function("motion_gen/setup_trajopt_seeds"):
+
+ goal = Goal(
+ current_state=start_state,
+ goal_state=goal_state,
+ )
+
+ if trajopt_seed_traj is None or graph_success < solve_state.num_trajopt_seeds * 1:
+ seed_goal = Goal(
+ current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
+ goal_state=goal_state.repeat_seeds(solve_state.num_trajopt_seeds),
+ )
+ if trajopt_seed_traj is not None:
+ trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
+ # batch, num_seeds, h, dof
+ if trajopt_seed_success.shape[1] < self.js_trajopt_solver.num_seeds:
+ trajopt_seed_success_new = torch.zeros(
+ (1, solve_state.num_trajopt_seeds),
+ device=self.tensor_args.device,
+ dtype=torch.bool,
+ )
+ trajopt_seed_success_new[0, : trajopt_seed_success.shape[1]] = (
+ trajopt_seed_success
+ )
+ trajopt_seed_success = trajopt_seed_success_new
+ # create seeds here:
+ trajopt_seed_traj = self.js_trajopt_solver.get_seed_set(
+ seed_goal,
+ trajopt_seed_traj, # batch, num_seeds, h, dof
+ num_seeds=1,
+ batch_mode=False,
+ seed_success=trajopt_seed_success,
+ )
+ trajopt_seed_traj = (
+ trajopt_seed_traj.view(
+ self.js_trajopt_solver.num_seeds * 1,
+ 1,
+ self.trajopt_solver.action_horizon,
+ self._dof,
+ )
+ .contiguous()
+ .clone()
+ )
+ if plan_config.enable_finetune_trajopt:
+ og_value = self.trajopt_solver.interpolation_type
+ self.js_trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
+ with profiler.record_function("motion_gen/trajopt"):
+ log_info("MG: running TO")
+ traj_result = self._solve_trajopt_from_solve_state(
+ goal,
+ solve_state,
+ trajopt_seed_traj,
+ num_seeds_override=solve_state.num_trajopt_seeds,
+ newton_iters=trajopt_newton_iters + 2,
+ return_all_solutions=plan_config.enable_finetune_trajopt,
+ trajopt_instance=self.js_trajopt_solver,
+ )
+ if plan_config.enable_finetune_trajopt:
+ self.trajopt_solver.interpolation_type = og_value
+ if self.store_debug_in_result:
+ result.debug_info["trajopt_result"] = traj_result
+ if torch.count_nonzero(traj_result.success) == 0:
+ result.status = MotionGenStatus.TRAJOPT_FAIL
+ # run finetune
+ if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
+ with profiler.record_function("motion_gen/finetune_trajopt"):
+ seed_traj = traj_result.raw_action.clone() # solution.position.clone()
+ seed_traj = seed_traj.contiguous()
+ og_solve_time = traj_result.solve_time
+
+ scaled_dt = torch.clamp(
+ torch.max(traj_result.optimized_dt[traj_result.success]),
+ self.trajopt_solver.minimum_trajectory_dt,
+ )
+ og_dt = self.js_trajopt_solver.solver_dt.clone()
+ self.js_trajopt_solver.update_solver_dt(scaled_dt.item())
+ traj_result = self._solve_trajopt_from_solve_state(
+ goal,
+ solve_state,
+ seed_traj,
+ trajopt_instance=self.js_trajopt_solver,
+ num_seeds_override=solve_state.num_trajopt_seeds,
+ newton_iters=trajopt_newton_iters + 4,
+ )
+ self.js_trajopt_solver.update_solver_dt(og_dt)
+
+ result.finetune_time = traj_result.solve_time
+
+ traj_result.solve_time = og_solve_time
+ if self.store_debug_in_result:
+ result.debug_info["finetune_trajopt_result"] = traj_result
+ if torch.count_nonzero(traj_result.success) == 0:
+ result.status = MotionGenStatus.FINETUNE_TRAJOPT_FAIL
+ elif plan_config.enable_finetune_trajopt:
+ traj_result.success = traj_result.success[0:1]
+ result.solve_time += traj_result.solve_time + result.finetune_time
+ result.trajopt_time = traj_result.solve_time
+ result.trajopt_attempts = 1
+ result.success = traj_result.success
+
+ result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory(
+ 0, traj_result.path_buffer_last_tstep[0]
+ )
+
+ result.interpolation_dt = self.trajopt_solver.interpolation_dt
+ result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
+ result.cspace_error = traj_result.cspace_error
+ result.optimized_dt = traj_result.optimized_dt
+ result.optimized_plan = traj_result.solution
+ result.goalset_index = traj_result.goalset_index
+
+ return result
+
+ def _plan_from_solve_state_batch(
+ self,
+ solve_state: ReacherSolveState,
+ start_state: JointState,
+ goal_pose: Pose,
+ plan_config: MotionGenPlanConfig = MotionGenPlanConfig(),
+ link_poses: Optional[Dict[str, Pose]] = None,
+ ) -> MotionGenResult:
+ """Plan from a given reacher solve state in batch mode.
+
+ Args:
+ solve_state: Reacher solve state for planning.
+ start_state: Start joint state for planning.
+ goal_pose: Goal poses to reach for end-effector.
+ plan_config: Planning parameters for motion generation.
+ link_poses: Goal poses for other links in the robot.
+
+ Returns:
+ MotionGenResult: Result of planning.
+ """
+ self._trajopt_goal_config[:] = self.get_retract_config().view(1, 1, self._dof)
+ trajopt_seed_traj = None
+ trajopt_seed_success = None
+ trajopt_newton_iters = None
+ graph_success = 0
+
+ # plan ik:
+ ik_result = self._solve_ik_from_solve_state(
+ goal_pose,
+ solve_state,
+ start_state,
+ plan_config.use_nn_ik_seed,
+ plan_config.partial_ik_opt,
+ link_poses,
+ )
+
+ if not plan_config.enable_graph and plan_config.partial_ik_opt:
+ ik_result.success[:] = True
+
+ # check for success:
+ result = MotionGenResult(
+ ik_result.success,
+ position_error=ik_result.position_error,
+ rotation_error=ik_result.rotation_error,
+ ik_time=ik_result.solve_time,
+ solve_time=ik_result.solve_time,
+ debug_info={},
+ # goalset_index=ik_result.goalset_index,
+ )
+
+ ik_success = torch.count_nonzero(ik_result.success)
+ if ik_success == 0:
+ result.status = MotionGenStatus.IK_FAIL
+ result.success = result.success[:, 0]
+ return result
+
+ # do graph search:
+ ik_out_seeds = solve_state.num_trajopt_seeds
+ if plan_config.enable_graph:
+ ik_out_seeds = min(solve_state.num_trajopt_seeds, ik_success)
+
+ # if not plan_config.enable_opt and plan_config.enable_graph:
+ # self.graph_planner.interpolation_steps = self.interpolation_steps
+ # self.graph_planner.interpolation_type = self.interpolation_type
+ # elif plan_config.enable_graph:
+ # self.graph_planner.interpolation_steps = self.trajopt_solver.traj_tsteps
+ # self.graph_planner.interpolation_type = InterpolateType.LINEAR
+ goal_config = ik_result.solution[ik_result.success].view(-1, self.ik_solver.dof)
+
+ # get shortest path
+ if plan_config.enable_graph:
+ interpolation_steps = None
+ if plan_config.enable_opt:
+ interpolation_steps = self.trajopt_solver.action_horizon
+
+ start_graph_state = start_state.repeat_seeds(ik_out_seeds)
+ start_config = start_graph_state.position[ik_result.success.view(-1)].view(
+ -1, self.ik_solver.dof
+ )
+ graph_result = self.graph_search(start_config, goal_config, interpolation_steps)
+ graph_success = torch.count_nonzero(graph_result.success).item()
+
+ result.graph_time = graph_result.solve_time
+ result.solve_time += graph_result.solve_time
+ if graph_success > 0:
+ # path = graph_result.interpolated_plan
+ result.graph_plan = graph_result.interpolated_plan
+ result.interpolated_plan = graph_result.interpolated_plan
+ result.used_graph = True
+
+ if plan_config.enable_opt:
+ trajopt_seed = result.graph_plan.position.view(
+ graph_success, # solve_state.num_trajopt_seeds,
+ interpolation_steps,
+ self._dof,
+ ).contiguous()
+ trajopt_seed_traj = torch.zeros(
+ (1, trajopt_seed.shape[0], self.trajopt_solver.action_horizon, self._dof),
+ device=self.tensor_args.device,
+ dtype=self.tensor_args.dtype,
+ )
+ trajopt_seed_traj[0, :, :interpolation_steps, :] = trajopt_seed
+ trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
+ trajopt_seed_success = ik_result.success.clone()
+ trajopt_seed_success[ik_result.success] = graph_result.success
+
+ trajopt_seed_success = trajopt_seed_success.view(
+ solve_state.num_trajopt_seeds, solve_state.batch_size
+ )
+ trajopt_newton_iters = self.graph_trajopt_iters
+
+ else:
+ ik_success = ik_result.success.view(-1).clone()
+
+ # only some might be successful:
+
+ g_dim = torch.nonzero(ik_success).view(-1)[graph_result.success]
+
+ self._batch_graph_search_buffer.copy_at_index(
+ graph_result.interpolated_plan, g_dim
+ )
+
+ # result.graph_plan = JointState.from_position(
+ # self._batch_graph_search_buffer,
+ # joint_names=graph_result.interpolated_plan.joint_names,
+ # )
+ result.interpolated_plan = self._batch_graph_search_buffer
+ g_dim = g_dim.cpu().squeeze().tolist()
+ for x, x_val in enumerate(g_dim):
+ self._batch_path_buffer_last_tstep[x_val] = (
+ graph_result.path_buffer_last_tstep[x]
+ )
+ result.path_buffer_last_tstep = self._batch_path_buffer_last_tstep
+ result.optimized_plan = result.interpolated_plan
+ result.optimized_dt = torch.as_tensor(
+ [
+ self.interpolation_dt
+ for i in range(result.interpolated_plan.position.shape[0])
+ ],
+ device=self.tensor_args.device,
+ dtype=self.tensor_args.dtype,
+ )
+ result.success = result.success.view(-1).clone()
+ result.success[ik_success][graph_result.success] = True
+ return result
+
+ else:
+ result.success[:] = False
+ result.success = result.success[:, 0]
+ result.status = MotionGenStatus.GRAPH_FAIL
+ if not graph_result.valid_query:
+ result.valid_query = False
+ if self.store_debug_in_result:
+ result.debug_info = {"graph_debug": graph_result.debug_info}
+ return result
+
+ if plan_config.enable_opt:
+ # get goal configs based on ik success:
+ self._trajopt_goal_config[ik_result.success] = goal_config
+
+ goal_config = self._trajopt_goal_config # batch index == 0
+
+ goal = Goal(
+ goal_pose=goal_pose,
+ current_state=start_state,
+ links_goal_pose=link_poses,
+ )
+ # generate seeds:
+ if trajopt_seed_traj is None or (
+ plan_config.enable_graph and graph_success < solve_state.batch_size
+ ):
+ seed_link_poses = None
+ if link_poses is not None:
+ seed_link_poses = {}
+
+ for k in link_poses.keys():
+ seed_link_poses[k] = link_poses[k].repeat_seeds(
+ solve_state.num_trajopt_seeds
+ )
+ seed_goal = Goal(
+ goal_pose=goal_pose.repeat_seeds(solve_state.num_trajopt_seeds),
+ current_state=start_state.repeat_seeds(solve_state.num_trajopt_seeds),
+ goal_state=JointState.from_position(goal_config.view(-1, self._dof)),
+ links_goal_pose=seed_link_poses,
+ )
+ if trajopt_seed_traj is not None:
+ trajopt_seed_traj = trajopt_seed_traj.transpose(0, 1).contiguous()
+
+ # create seeds here:
+ trajopt_seed_traj = self.trajopt_solver.get_seed_set(
+ seed_goal,
+ trajopt_seed_traj, # batch, num_seeds, h, dof
+ num_seeds=1,
+ batch_mode=solve_state.batch_mode,
+ seed_success=trajopt_seed_success,
+ )
+ trajopt_seed_traj = trajopt_seed_traj.view(
+ solve_state.num_trajopt_seeds,
+ solve_state.batch_size,
+ self.trajopt_solver.action_horizon,
+ self._dof,
+ ).contiguous()
+ if plan_config.enable_finetune_trajopt:
+ og_value = self.trajopt_solver.interpolation_type
+ self.trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA
+
+ traj_result = self._solve_trajopt_from_solve_state(
+ goal,
+ solve_state,
+ trajopt_seed_traj,
+ newton_iters=trajopt_newton_iters,
+ return_all_solutions=plan_config.enable_finetune_trajopt,
+ )
+
+ # output of traj result will have 1 solution per batch
+
+ # run finetune opt on 1 solution per batch:
+ if plan_config.enable_finetune_trajopt:
+ self.trajopt_solver.interpolation_type = og_value
+ if self.store_debug_in_result:
+ result.debug_info["trajopt_result"] = traj_result
+
+ # run finetune
+ if plan_config.enable_finetune_trajopt and torch.count_nonzero(traj_result.success) > 0:
+ with profiler.record_function("motion_gen/finetune_trajopt"):
+ seed_traj = traj_result.raw_action.clone() # solution.position.clone()
+ seed_traj = seed_traj.contiguous()
+ og_solve_time = traj_result.solve_time
+
+ scaled_dt = torch.clamp(
+ torch.max(traj_result.optimized_dt[traj_result.success])
+ * self.finetune_dt_scale,
+ self.trajopt_solver.minimum_trajectory_dt,
+ )
+ self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item())
+
+ traj_result = self._solve_trajopt_from_solve_state(
+ goal,
+ solve_state,
+ seed_traj,
+ trajopt_instance=self.finetune_trajopt_solver,
+ num_seeds_override=solve_state.num_trajopt_seeds,
+ )
+
+ result.finetune_time = traj_result.solve_time
+
+ traj_result.solve_time = og_solve_time
+ if self.store_debug_in_result:
+ result.debug_info["finetune_trajopt_result"] = traj_result
+ elif plan_config.enable_finetune_trajopt and len(traj_result.success.shape) == 2:
+ traj_result.success = traj_result.success[:, 0]
+
+ result.success = traj_result.success
+
+ result.interpolated_plan = traj_result.interpolated_solution
+ result.solve_time += traj_result.solve_time
+ result.trajopt_time = traj_result.solve_time
+ result.position_error = traj_result.position_error
+ result.rotation_error = traj_result.rotation_error
+ result.cspace_error = traj_result.cspace_error
+ result.goalset_index = traj_result.goalset_index
+ result.path_buffer_last_tstep = traj_result.path_buffer_last_tstep
+ result.optimized_plan = traj_result.solution
+ result.optimized_dt = traj_result.optimized_dt
+ if torch.count_nonzero(traj_result.success) == 0:
+ result.status = MotionGenStatus.TRAJOPT_FAIL
+ result.success[:] = False
+ if self.store_debug_in_result:
+ result.debug_info = {"trajopt_result": traj_result}
+ return result
+
+ def plan(
+ self,
+ start_state: JointState,
+ goal_pose: Pose,
+ enable_graph: bool = True,
+ enable_opt: bool = True,
+ use_nn_ik_seed: bool = False,
+ need_graph_success: bool = False,
+ max_attempts: int = 100,
+ timeout: float = 10.0,
+ enable_graph_attempt: Optional[int] = None,
+ disable_graph_attempt: Optional[int] = None,
+ trajopt_attempts: int = 1,
+ ik_fail_return: Optional[int] = None,
+ partial_ik_opt: bool = False,
+ num_ik_seeds: Optional[int] = None,
+ num_graph_seeds: Optional[int] = None,
+ num_trajopt_seeds: Optional[int] = None,
+ ):
+ """Deprecated method. Use :meth:`MotionGen.plan_single` or others instead."""
+
+ log_warn("Deprecated method. Use MotionGen.plan_single or others instead.")
+ plan_config = MotionGenPlanConfig(
+ enable_graph,
+ enable_opt,
+ use_nn_ik_seed,
+ need_graph_success,
+ max_attempts,
+ timeout,
+ enable_graph_attempt,
+ disable_graph_attempt,
+ ik_fail_return,
+ partial_ik_opt,
+ num_ik_seeds,
+ num_graph_seeds,
+ num_trajopt_seeds,
+ )
+ result = self.plan_single(start_state, goal_pose, plan_config)
+ return result
+
+ def batch_plan(
+ self,
+ start_state: JointState,
+ goal_pose: Pose,
+ enable_graph: bool = True,
+ enable_opt: bool = True,
+ use_nn_ik_seed: bool = False,
+ need_graph_success: bool = False,
+ max_attempts: int = 1,
+ timeout: float = 10.0,
+ enable_graph_attempt: Optional[int] = None,
+ disable_graph_attempt: Optional[int] = None,
+ success_ratio: float = 1.0,
+ ik_fail_return: Optional[int] = None,
+ fail_on_invalid_query: bool = False,
+ partial_ik_opt: bool = False,
+ num_ik_seeds: Optional[int] = None,
+ num_graph_seeds: Optional[int] = None,
+ num_trajopt_seeds: Optional[int] = None,
+ ):
+ """Deprecated method. Use :meth:`MotionGen.plan_batch` or others instead."""
+
+ log_warn("Deprecated method. Use MotionGen.plan_batch or others instead.")
+
+ plan_config = MotionGenPlanConfig(
+ enable_graph,
+ enable_opt,
+ use_nn_ik_seed,
+ need_graph_success,
+ max_attempts,
+ timeout,
+ enable_graph_attempt,
+ disable_graph_attempt,
+ ik_fail_return,
+ partial_ik_opt,
+ num_ik_seeds,
+ num_graph_seeds,
+ num_trajopt_seeds,
+ success_ratio=success_ratio,
+ fail_on_invalid_query=fail_on_invalid_query,
+ )
+ result = self.plan_batch(start_state, goal_pose, plan_config)
+ return result
diff --git a/src/curobo/wrap/reacher/mpc.py b/src/curobo/wrap/reacher/mpc.py
index 378756d..28919b5 100644
--- a/src/curobo/wrap/reacher/mpc.py
+++ b/src/curobo/wrap/reacher/mpc.py
@@ -8,7 +8,31 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
+"""
+This module contains :meth:`MpcSolver` that provides a high-level interface to for model
+predictive control (MPC) for reaching Cartesian poses and also joint configurations while
+avoiding obstacles. The solver uses Model Predictive Path Integral (MPPI) optimization as the
+solver. MPC only optimizes locally so the robot can get stuck near joint limits or behind
+obstacles. To generate global trajectories, use
+:py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`.
+A python example is available at :ref:`python_mpc_example`.
+
+
+
+.. note::
+ Gradient-based MPC is also implemented with L-BFGS but is highly experimental and not
+ recommended for real robots.
+
+
+.. raw:: html
+
+
+
+
+
+
+"""
# Standard Library
import time
@@ -19,6 +43,7 @@ from typing import Dict, Optional, Union
import torch
# CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.geom.sdf.utils import create_collision_checker
from curobo.geom.sdf.world import CollisionCheckerType, WorldCollision, WorldCollisionConfig
from curobo.geom.types import WorldConfig
@@ -44,22 +69,31 @@ from curobo.wrap.wrap_mpc import WrapConfig, WrapMpc
@dataclass
class MpcSolverConfig:
+ """Configuration dataclass for MPC."""
+
+ #: MPC Solver.
solver: WrapMpc
+
+ #: World Collision Checker.
world_coll_checker: Optional[WorldCollision] = None
+
+ #: Numeric precision and device to run computations.
tensor_args: TensorDeviceType = TensorDeviceType()
+
+ #: Capture full step in MPC as a single CUDA graph. This is not supported currently.
use_cuda_graph_full_step: bool = False
@staticmethod
def load_from_robot_config(
robot_cfg: Union[Union[str, dict], RobotConfig],
- world_cfg: Union[Union[str, dict], WorldConfig],
+ world_model: Union[Union[str, dict], WorldConfig],
base_cfg: Optional[dict] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
compute_metrics: bool = True,
use_cuda_graph: Optional[bool] = None,
particle_opt_iters: Optional[int] = None,
self_collision_check: bool = True,
- collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE,
+ collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH,
use_es: Optional[bool] = None,
es_learning_rate: Optional[float] = 0.01,
use_cuda_graph_metrics: bool = False,
@@ -74,9 +108,56 @@ class MpcSolverConfig:
use_lbfgs: bool = False,
use_mppi: bool = True,
):
+ """Create an MPC solver configuration from robot and world configuration.
+
+ Args:
+ robot_cfg: Robot configuration. Can be a path to a YAML file or a dictionary or
+ an instance of :class:`~curobo.types.robot.RobotConfig`.
+ world_model: World configuration. Can be a path to a YAML file or a dictionary or
+ an instance of :class:`~curobo.geom.types.WorldConfig`.
+ base_cfg: Base configuration for the solver. This file is used to check constraints
+ and convergence. If None, the default configuration from ``base_cfg.yml`` is used.
+ tensor_args: Numeric precision and device to run computations.
+ compute_metrics: Compute metrics on MPC step.
+ use_cuda_graph: Use CUDA graph for the optimization step.
+ particle_opt_iters: Number of iterations for the particle optimization.
+ self_collision_check: Enable self-collision check during MPC optimization.
+ collision_checker_type: Type of collision checker to use. See :ref:`world_collision`.
+ use_es: Use Evolution Strategies (ES) solver for MPC. Highly experimental.
+ es_learning_rate: Learning rate for ES solver.
+ use_cuda_graph_metrics: Use CUDA graph for computing metrics.
+ store_rollouts: Store rollouts information for debugging. This will also store the
+ trajectory taken by the end-effector across the horizon.
+ use_cuda_graph_full_step: Capture full step in MPC as a single CUDA graph. This is
+ experimental and might not work reliably.
+ sync_cuda_time: Synchronize CUDA device with host using
+ :py:func:`torch.cuda.synchronize` before calculating compute time.
+ collision_cache: Cache of obstacles to create to load obstacles between planning calls.
+ An example: ``{"obb": 10, "mesh": 10}``, to create a cache of 10 cuboids and 10
+ meshes.
+ n_collision_envs: Number of collision environments to create for batched planning
+ across different environments. Only used for :py:meth:`MpcSolver.setup_solve_batch_env`
+ and :py:meth:`MpcSolver.setup_solve_batch_env_goalset`.
+ collision_activation_distance: Distance in meters to activate collision cost. A good
+ value to start with is 0.01 meters. Increase the distance if the robot needs to
+ stay further away from obstacles.
+ world_coll_checker: Instance of world collision checker to use for MPC. Leaving this to
+ None will create a new instance of world collision checker using the provided
+ :attr:`world_model`.
+ step_dt: Time step to use between each step in the trajectory. If None, the default
+ time step from the configuration~(`particle_mpc.yml` or `gradient_mpc.yml`)
+ is used. This dt should match the control frequency at which you are sending
+ commands to the robot. This dt should also be greater than than the compute
+ time for a single step.
+ use_lbfgs: Use L-BFGS solver for MPC. Highly experimental.
+ use_mppi: Use MPPI solver for MPC.
+
+ Returns:
+ MpcSolverConfig: Configuration for the MPC solver.
+ """
+
if use_cuda_graph_full_step:
log_error("use_cuda_graph_full_step currently is not supported")
- raise ValueError("use_cuda_graph_full_step currently is not supported")
task_file = "particle_mpc.yml"
config_data = load_yaml(join_path(get_task_configs_path(), task_file))
@@ -108,14 +189,14 @@ class MpcSolverConfig:
if isinstance(robot_cfg, dict):
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
- if isinstance(world_cfg, str):
- world_cfg = load_yaml(join_path(get_world_configs_path(), world_cfg))
+ if isinstance(world_model, str):
+ world_model = load_yaml(join_path(get_world_configs_path(), world_model))
- if world_coll_checker is None and world_cfg is not None:
- world_cfg = WorldCollisionConfig.load_from_dict(
- base_cfg["world_collision_checker_cfg"], world_cfg, tensor_args
+ if world_coll_checker is None and world_model is not None:
+ world_model = WorldCollisionConfig.load_from_dict(
+ base_cfg["world_collision_checker_cfg"], world_model, tensor_args
)
- world_coll_checker = create_collision_checker(world_cfg)
+ world_coll_checker = create_collision_checker(world_model)
grad_config_data = None
if use_lbfgs:
grad_config_data = load_yaml(join_path(get_task_configs_path(), "gradient_mpc.yml"))
@@ -134,7 +215,7 @@ class MpcSolverConfig:
base_cfg["constraint"],
base_cfg["convergence"],
base_cfg["world_collision_checker_cfg"],
- world_cfg,
+ world_model,
world_coll_checker=world_coll_checker,
tensor_args=tensor_args,
)
@@ -172,7 +253,7 @@ class MpcSolverConfig:
base_cfg["constraint"],
base_cfg["convergence"],
base_cfg["world_collision_checker_cfg"],
- world_cfg,
+ world_model,
world_coll_checker=world_coll_checker,
tensor_args=tensor_args,
)
@@ -201,13 +282,42 @@ class MpcSolverConfig:
class MpcSolver(MpcSolverConfig):
- """Model Predictive Control Solver for Arm Reacher task.
+ """High-level interface for Model Predictive Control (MPC).
- Args:
- MpcSolverConfig: _description_
+ MPC can reach Cartesian poses and joint configurations while avoiding obstacles. The solver
+ uses Model Predictive Path Integral (MPPI) optimization as the solver. MPC only optimizes
+ locally so the robot can get stuck near joint limits or behind obstacles. To generate global
+ trajectories, use :py:meth:`~curobo.wrap.reacher.motion_gen.MotionGen`.
+
+ See :ref:`python_mpc_example` for an example. This MPC solver implementation can be used in the
+ following steps:
+
+ 1. Create a :py:class:`~curobo.rollout.rollout_base.Goal` object with the target pose or joint
+ configuration.
+ 2. Create a goal buffer for the problem type using :meth:`setup_solve_single`,
+ :meth:`setup_solve_goalset`, :meth:`setup_solve_batch`, :meth:`setup_solve_batch_goalset`,
+ :meth:`setup_solve_batch_env`, or :meth:`setup_solve_batch_env_goalset`. Pass the goal
+ object from the previous step to this function. This function will update the internal
+ solve state of MPC and also the goal for MPC. An augmented goal buffer is returned.
+ 3. Call :meth:`step` with the current joint state to get the next action.
+ 4. To change the goal, create a :py:class:`~curobo.types.math.Pose` object with new pose or
+ :py:class:`~curobo.types.state.JointState` object with new joint configuration. Then
+ copy the target into the augmented goal buffer using
+ ``goal_buffer.goal_pose.copy_(new_pose)`` or ``goal_buffer.goal_state.copy_(new_state)``.
+ 5. Call :meth:`update_goal` with the augmented goal buffer to update the goal for MPC.
+ 6. Call :meth:`step` with the current joint state to get the next action.
+
+ To dynamically change the type of goal reached between pose and joint configuration targets,
+ create the goal object in step 1 with both targets and then use :meth:`enable_cspace_cost` and
+ :meth:`enable_pose_cost` to enable or disable reaching joint configuration cost and pose cost.
"""
def __init__(self, config: MpcSolverConfig) -> None:
+ """Initializes the MPC solver.
+
+ Args:
+ config: Configuration parameters for MPC.
+ """
super().__init__(**vars(config))
self.tensor_args = self.solver.rollout_fn.tensor_args
self._goal_buffer = Goal()
@@ -222,15 +332,326 @@ class MpcSolver(MpcSolverConfig):
self._cu_step_graph = None
self._cu_result = None
- def _update_batch_size(self, batch_size):
- if self.batch_size != batch_size:
- self.batch_size = batch_size
+ def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
+ """Creates a goal buffer to solve for a robot to reach target pose or joint configuration.
- def update_goal_buffer(
+ Args:
+ goal: goal object containing target pose or joint configuration.
+ num_seeds: Number of seeds to use in the solver.
+
+ Returns:
+ Goal: Instance of augmented goal buffer.
+ """
+ solve_state = ReacherSolveState(
+ ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1
+ )
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+
+ self.update_goal(goal_buffer)
+
+ return goal_buffer
+
+ def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
+ """Creates a goal buffer to solve for a robot to reach a pose in a set of poses.
+
+ Args:
+ goal: goal object containing target goalset or joint configuration.
+ num_seeds: Number of seeds to use in the solver.
+
+ Returns:
+ Goal: Instance of augmented goal buffer.
+ """
+ solve_state = ReacherSolveState(
+ ReacherSolveType.GOALSET,
+ num_mpc_seeds=num_seeds,
+ batch_size=1,
+ n_envs=1,
+ n_goalset=goal.n_goalset,
+ )
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+ self.update_goal(goal_buffer)
+ return goal_buffer
+
+ def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
+ """Creates a goal buffer to solve for a batch of robots to reach targets.
+
+ Args:
+ goal: goal object containing target poses or joint configurations.
+ num_seeds: Number of seeds to use in the solver.
+
+ Returns:
+ Goal: Instance of augmented goal buffer.
+ """
+ solve_state = ReacherSolveState(
+ ReacherSolveType.BATCH,
+ num_mpc_seeds=num_seeds,
+ batch_size=goal.batch,
+ n_envs=1,
+ n_goalset=1,
+ )
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+ self.update_goal(goal_buffer)
+
+ return goal_buffer
+
+ def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
+ """Creates a goal buffer to solve for a batch of robots to reach a set of poses.
+
+ Args:
+ goal: goal object containing target goalset or joint configurations.
+ num_seeds: Number of seeds to use in the solver.
+
+ Returns:
+ Goal: Instance of augmented goal buffer.
+ """
+ solve_state = ReacherSolveState(
+ ReacherSolveType.BATCH_GOALSET,
+ num_mpc_seeds=num_seeds,
+ batch_size=goal.batch,
+ n_envs=1,
+ n_goalset=goal.n_goalset,
+ )
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+ self.update_goal(goal_buffer)
+
+ return goal_buffer
+
+ def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
+ """Creates a goal buffer to solve for a batch of robots in different collision worlds.
+
+ Args:
+ goal: goal object containing target poses or joint configurations.
+ num_seeds: Number of seeds to use in the solver.
+
+ Returns:
+ Goal: Instance of augmented goal buffer.
+ """
+ solve_state = ReacherSolveState(
+ ReacherSolveType.BATCH_ENV,
+ num_mpc_seeds=num_seeds,
+ batch_size=goal.batch,
+ n_envs=1,
+ n_goalset=1,
+ )
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+ self.update_goal(goal_buffer)
+
+ return goal_buffer
+
+ def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
+ """Creates a goal buffer to solve for a batch of robots in different collision worlds.
+
+ Args:
+ goal: goal object containing target goalset or joint configurations.
+ num_seeds: Number of seeds to use in the solver.
+
+ Returns:
+ Goal: Instance of augmented goal buffer.
+ """
+ solve_state = ReacherSolveState(
+ ReacherSolveType.BATCH_ENV_GOALSET,
+ num_mpc_seeds=num_seeds,
+ batch_size=goal.batch,
+ n_envs=1,
+ n_goalset=goal.n_goalset,
+ )
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+ self.update_goal(goal_buffer)
+
+ return goal_buffer
+
+ def step(
+ self,
+ current_state: JointState,
+ shift_steps: int = 1,
+ seed_traj: Optional[JointState] = None,
+ max_attempts: int = 1,
+ ):
+ """Solve for the next action given the current state.
+
+ Args:
+ current_state: Current joint state of the robot.
+ shift_steps: Number of steps to shift the trajectory.
+ seed_traj: Initial trajectory to seed the optimization. If None, the solver
+ uses the solution from the previous step.
+ max_attempts: Maximum number of attempts to solve the problem.
+
+ Returns:
+ WrapResult: Result of the optimization.
+ """
+ converged = True
+
+ for _ in range(max_attempts):
+ result = self._step_once(current_state.clone(), shift_steps, seed_traj)
+ if (
+ torch.count_nonzero(torch.isnan(result.action.position)) == 0
+ and torch.count_nonzero(~result.metrics.feasible) == 0
+ ):
+ converged = True
+ break
+ self.reset()
+ if not converged:
+ result.action.copy_(current_state)
+ log_warn("MPC didn't converge")
+
+ return result
+
+ def update_goal(self, goal: Goal):
+ """Update the goal for MPC.
+
+ Args:
+ goal: goal object containing target pose or joint configuration. This goal instance
+ should be created using one of the setup_solve functions.
+ """
+ self.solver.update_params(goal)
+
+ def reset(self):
+ """Reset the solver."""
+ # reset warm start
+ self.solver.reset()
+
+ def reset_cuda_graph(self):
+ """Reset captured CUDA graph. This does not work."""
+ self.solver.reset_cuda_graph()
+
+ def enable_cspace_cost(self, enable=True):
+ """Enable or disable reaching joint configuration cost in the solver.
+
+ Args:
+ enable: Enable or disable reaching joint configuration cost. When False, cspace cost
+ is disabled.
+ """
+ self.solver.safety_rollout.enable_cspace_cost(enable)
+ for opt in self.solver.optimizers:
+ opt.rollout_fn.enable_cspace_cost(enable)
+
+ def enable_pose_cost(self, enable=True):
+ """Enable or disable reaching pose cost in the solver.
+
+ Args:
+ enable: Enable or disable reaching pose cost. When False, pose cost is disabled.
+ """
+ self.solver.safety_rollout.enable_pose_cost(enable)
+ for opt in self.solver.optimizers:
+ opt.rollout_fn.enable_pose_cost(enable)
+
+ def get_active_js(
+ self,
+ in_js: JointState,
+ ):
+ """Get controlled joints indexed in MPC order from the input joint state.
+
+ Args:
+ in_js: Input joint state.
+
+ Returns:
+ JointState: Joint state with controlled joints.
+ """
+
+ opt_jnames = self.rollout_fn.joint_names
+ opt_js = in_js.get_ordered_joint_state(opt_jnames)
+ return opt_js
+
+ def update_world(self, world: WorldConfig):
+ """Update the collision world for the solver.
+
+ This allows for updating the world representation as long as the new world representation
+ does not have a larger number of obstacles than the :attr:`MpcSolver.collision_cache` as
+ created during initialization of :class:`MpcSolverConfig`.
+
+ Args:
+ world: New collision world configuration. See :ref:`world_collision` for more details.
+ """
+ self.world_coll_checker.load_collision_model(world)
+
+ def get_visual_rollouts(self):
+ """Get rollouts for debugging."""
+ return self.solver.optimizers[0].get_rollouts()
+
+ @property
+ def joint_names(self):
+ """Get the ordered joint names of the robot."""
+ return self.rollout_fn.joint_names
+
+ @property
+ def collision_cache(self) -> Dict[str, int]:
+ """Returns the collision cache created by the world collision checker."""
+ return self.world_coll_checker.cache
+
+ @property
+ def kinematics(self) -> CudaRobotModel:
+ """Get kinematics instance of the robot."""
+ return self.solver.safety_rollout.dynamics_model.robot_model
+
+ @property
+ def world_collision(self) -> WorldCollision:
+ """Get the world collision checker."""
+ return self.world_coll_checker
+
+ @property
+ def rollout_fn(self) -> ArmReacher:
+ """Get the rollout function."""
+ return self.solver.safety_rollout
+
+ def _step_once(
+ self,
+ current_state: JointState,
+ shift_steps: int = 1,
+ seed_traj: Optional[JointState] = None,
+ ) -> WrapResult:
+ """Solve for the next action given the current state.
+
+ Args:
+ current_state: Current joint state of the robot.
+ shift_steps: Number of steps to shift the trajectory.
+ seed_traj: Initial trajectory to seed the optimization. If None, the solver
+ uses the solution from the previous step.
+
+ Returns:
+ WrapResult: Result of the optimization.
+ """
+ # Create cuda graph for whole solve step including computation of metrics
+ # Including updation of goal buffers
+
+ if self._solve_state is None:
+ log_error("Need to first setup solve state before calling solve()")
+
+ if self.use_cuda_graph_full_step:
+ st_time = time.time()
+ if not self._cu_step_init:
+ self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
+ self._cu_state_in.copy_(current_state)
+ if seed_traj is not None:
+ self._cu_seed.copy_(seed_traj)
+ self._cu_step_graph.replay()
+ result = self._cu_result.clone()
+ torch.cuda.synchronize(device=self.tensor_args.device)
+ result.solve_time = time.time() - st_time
+ else:
+ self._step_goal_buffer.current_state.copy_(current_state)
+ result = self._solve_from_solve_state(
+ self._solve_state,
+ self._step_goal_buffer,
+ shift_steps,
+ seed_traj,
+ )
+
+ return result
+
+ def _update_solve_state_and_goal_buffer(
self,
solve_state: ReacherSolveState,
goal: Goal,
) -> Goal:
+ """Update solve state and goal for MPC.
+
+ Args:
+ solve_state: New solve state.
+ goal: New goal buffer.
+
+ Returns:
+ Goal: Updated goal buffer.
+ """
self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal(
goal,
self._solve_state,
@@ -250,71 +671,64 @@ class MpcSolver(MpcSolverConfig):
)
return self._goal_buffer
- def step(
+ def _update_batch_size(self, batch_size: int):
+ """Update the batch size of the solver.
+
+ Args:
+ batch_size: Number of problems to solve in parallel.
+ """
+ if self.batch_size != batch_size:
+ self.batch_size = batch_size
+
+ def _solve_from_solve_state(
self,
- current_state: JointState,
- shift_steps: int = 1,
- seed_traj: Optional[JointState] = None,
- max_attempts: int = 1,
- ):
- converged = True
-
- for _ in range(max_attempts):
- result = self.step_once(current_state.clone(), shift_steps, seed_traj)
- if (
- torch.count_nonzero(torch.isnan(result.action.position)) == 0
- and torch.max(torch.abs(result.action.position)) < 10.0
- and torch.count_nonzero(~result.metrics.feasible) == 0
- ):
- converged = True
- break
- self.reset()
- if not converged:
- result.action.copy_(current_state)
- log_warn("NOT CONVERGED")
-
- return result
-
- def step_once(
- self,
- current_state: JointState,
+ solve_state: ReacherSolveState,
+ goal: Goal,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
) -> WrapResult:
- # Create cuda graph for whole solve step including computation of metrics
- # Including updation of goal buffers
+ """Solve for the next action given the current state.
- if self._solve_state is None:
- log_error("Need to first setup solve state before calling solve()")
+ Args:
+ solve_state: solve state object containing information about the current MPC problem.
+ goal: goal object containing target pose or joint configuration.
+ shift_steps: Number of steps to shift the trajectory before optimization.
+ seed_traj: Initial trajectory to seed the optimization. If None, the solver
+ uses the solution from the previous step.
- if self.use_cuda_graph_full_step:
- st_time = time.time()
- if not self._cu_step_init:
- self._initialize_cuda_graph_step(current_state, shift_steps, seed_traj)
- self._cu_state_in.copy_(current_state)
- if seed_traj is not None:
- self._cu_seed.copy_(seed_traj)
- self._cu_step_graph.replay()
- result = self._cu_result.clone()
- torch.cuda.synchronize()
- result.solve_time = time.time() - st_time
- else:
- self._step_goal_buffer.current_state.copy_(current_state)
- result = self._solve_from_solve_state(
- self._solve_state,
- self._step_goal_buffer,
- shift_steps,
- seed_traj,
- )
+ Returns:
+ WrapResult: Result of the optimization.
+ """
+ if solve_state.batch_env:
+ if solve_state.batch_size > self.world_coll_checker.n_envs:
+ log_error("Batch Env is less that goal batch")
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
+
+ if seed_traj is not None:
+ self.solver.update_init_seed(seed_traj)
+
+ result = self.solver.solve(goal_buffer, seed_traj, shift_steps)
+ result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action)
return result
- def _step(
+ def _mpc_step(
self,
current_state: JointState,
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
):
+ """One step function that is used to create a CUDA graph.
+
+ Args:
+ current_state: Current joint state of the robot.
+ shift_steps: Number of steps to shift the trajectory.
+ seed_traj: Initial trajectory to seed the optimization. If None, the solver
+ uses the solution from the previous step.
+
+ Returns:
+ WrapResult: Result of the optimization.
+ """
self._step_goal_buffer.current_state.copy_(current_state)
result = self._solve_from_solve_state(
self._solve_state,
@@ -331,6 +745,14 @@ class MpcSolver(MpcSolverConfig):
shift_steps: int = 1,
seed_traj: Optional[JointState] = None,
):
+ """Create a CUDA graph for the full step of MPC.
+
+ Args:
+ current_state: Current joint state of the robot.
+ shift_steps: Number of steps to shift the trajectory.
+ seed_traj: Initial trajectory to seed the optimization. If None, the solver
+ uses the solution from the previous step.
+ """
log_info("MpcSolver: Creating Cuda Graph")
self._cu_state_in = current_state.clone()
if seed_traj is not None:
@@ -339,7 +761,7 @@ class MpcSolver(MpcSolverConfig):
s.wait_stream(torch.cuda.current_stream(device=self.tensor_args.device))
with torch.cuda.stream(s):
for _ in range(3):
- self._cu_result = self._step(
+ self._cu_result = self._mpc_step(
self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed
)
torch.cuda.current_stream(device=self.tensor_args.device).wait_stream(s)
@@ -350,142 +772,3 @@ class MpcSolver(MpcSolverConfig):
self._cu_state_in, shift_steps=shift_steps, seed_traj=self._cu_seed
)
self._cu_step_init = True
-
- def setup_solve_single(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
- solve_state = ReacherSolveState(
- ReacherSolveType.SINGLE, num_mpc_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1
- )
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- return goal_buffer
-
- def setup_solve_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
- solve_state = ReacherSolveState(
- ReacherSolveType.GOALSET,
- num_mpc_seeds=num_seeds,
- batch_size=1,
- n_envs=1,
- n_goalset=goal.n_goalset,
- )
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- return goal_buffer
-
- def setup_solve_batch(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
- solve_state = ReacherSolveState(
- ReacherSolveType.BATCH,
- num_mpc_seeds=num_seeds,
- batch_size=goal.batch,
- n_envs=1,
- n_goalset=1,
- )
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- return goal_buffer
-
- def setup_solve_batch_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
- solve_state = ReacherSolveState(
- ReacherSolveType.BATCH_GOALSET,
- num_mpc_seeds=num_seeds,
- batch_size=goal.batch,
- n_envs=1,
- n_goalset=goal.n_goalset,
- )
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- return goal_buffer
-
- def setup_solve_batch_env(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
- solve_state = ReacherSolveState(
- ReacherSolveType.BATCH_ENV,
- num_mpc_seeds=num_seeds,
- batch_size=goal.batch,
- n_envs=1,
- n_goalset=1,
- )
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- return goal_buffer
-
- def setup_solve_batch_env_goalset(self, goal: Goal, num_seeds: Optional[int] = None) -> Goal:
- solve_state = ReacherSolveState(
- ReacherSolveType.BATCH_ENV_GOALSET,
- num_mpc_seeds=num_seeds,
- batch_size=goal.batch,
- n_envs=1,
- n_goalset=goal.n_goalset,
- )
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- return goal_buffer
-
- def _solve_from_solve_state(
- self,
- solve_state: ReacherSolveState,
- goal: Goal,
- shift_steps: int = 1,
- seed_traj: Optional[JointState] = None,
- ) -> WrapResult:
- if solve_state.batch_env:
- if solve_state.batch_size > self.world_coll_checker.n_envs:
- raise ValueError("Batch Env is less that goal batch")
-
- goal_buffer = self.update_goal_buffer(solve_state, goal)
- # NOTE: implement initialization from seed set here:
- if seed_traj is not None:
- self.solver.update_init_seed(seed_traj)
-
- result = self.solver.solve(goal_buffer, seed_traj, shift_steps)
- result.js_action = self.rollout_fn.get_full_dof_from_solution(result.action)
- return result
-
- def fn(self):
- # this will run one step of optimization and get new command
- pass
-
- def update_goal(self, goal: Goal):
- self.solver.update_params(goal)
- return True
-
- def reset(self):
- # reset warm start
- self.solver.reset()
- pass
-
- def reset_cuda_graph(self):
- self.solver.reset_cuda_graph()
-
- @property
- def rollout_fn(self):
- return self.solver.safety_rollout
-
- def enable_cspace_cost(self, enable=True):
- self.solver.safety_rollout.enable_cspace_cost(enable)
- for opt in self.solver.optimizers:
- opt.rollout_fn.enable_cspace_cost(enable)
-
- def enable_pose_cost(self, enable=True):
- self.solver.safety_rollout.enable_pose_cost(enable)
- for opt in self.solver.optimizers:
- opt.rollout_fn.enable_pose_cost(enable)
-
- def get_active_js(
- self,
- in_js: JointState,
- ):
- opt_jnames = self.rollout_fn.joint_names
- opt_js = in_js.get_ordered_joint_state(opt_jnames)
- return opt_js
-
- @property
- def joint_names(self):
- return self.rollout_fn.joint_names
-
- def update_world(self, world: WorldConfig):
- self.world_coll_checker.load_collision_model(world)
- return True
-
- def get_visual_rollouts(self):
- return self.solver.optimizers[0].get_rollouts()
-
- @property
- def kinematics(self):
- return self.solver.safety_rollout.dynamics_model.robot_model
-
- @property
- def world_collision(self):
- return self.world_coll_checker
diff --git a/src/curobo/wrap/reacher/trajopt.py b/src/curobo/wrap/reacher/trajopt.py
index 8ce1dbf..6bcbd60 100644
--- a/src/curobo/wrap/reacher/trajopt.py
+++ b/src/curobo/wrap/reacher/trajopt.py
@@ -8,13 +8,29 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
+"""
+Trajectory optimization module contains :meth:`TrajOptSolver` class which can optimize for
+minimum-jerk trajectories by first running a particle-based solver (MPPI) and then refining with
+a gradient-based solver (L-BFGS). The module also provides linear interpolation functions for
+generating seeds from start joint configuration to goal joint configurations. To generate
+trajectories for reaching Cartesian poses or joint configurations, use the higher-level wrapper
+:py:class:`~curobo.wrap.reacher.motion_gen.MotionGen`.
+Trajectory Optimization uses joint positions as optimization variables with cost terms for
+avoiding world collisions, self-collisions, and joint limits. The joint velocities, accelerations,
+and jerks are computed using five point stencil. A squared l2-norm cost term on joint accelerations
+and jerks is used to encourage smooth trajectories. A cost term for the terminal state to reach
+either a Cartesian pose or joint configuration is also used. Read :ref:`research_page` for
+more details.
+"""
+
+from __future__ import annotations
# Standard Library
import math
import time
from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
# Third Party
import torch
@@ -52,6 +68,8 @@ from curobo.wrap.wrap_base import WrapBase, WrapConfig, WrapResult
@dataclass
class TrajOptSolverConfig:
+ """Configuration parameters for TrajOptSolver."""
+
robot_config: RobotConfig
solver: WrapBase
rollout_fn: ArmReacher
@@ -98,6 +116,7 @@ class TrajOptSolverConfig:
interpolation_type: InterpolateType = InterpolateType.LINEAR_CUDA,
interpolation_steps: int = 10000,
interpolation_dt: float = 0.01,
+ minimum_trajectory_dt: Optional[float] = None,
use_cuda_graph: bool = True,
self_collision_check: bool = False,
self_collision_opt: bool = True,
@@ -130,8 +149,156 @@ class TrajOptSolverConfig:
optimize_dt: bool = True,
project_pose_to_goal_frame: bool = True,
):
- # NOTE: Don't have default optimize_dt, instead read from a configuration file.
- # use default values, disable environment collision checking
+ """Load TrajOptSolver configuration from robot configuration.
+
+ Args:
+ robot_cfg: Robot configuration to use for motion generation. This can be a path to a
+ yaml file, a dictionary, or an instance of :class:`RobotConfig`. See
+ :ref:`available_robot_list` for a list of available robots. You can also create a
+ a configuration file for your robot using :ref:`tut_robot_configuration`.
+ world_model: World configuration to use for motion generation. This can be a path to a
+ yaml file, a dictionary, or an instance of :class:`WorldConfig`. See
+ :ref:`world_collision` for more details.
+ tensor_args: Numerical precision and compute device to use for motion generation.
+ position_threshold: Position threshold between target position and reached position in
+ meters. 0.005 is a good value for most tasks (5mm).
+ rotation_threshold: Rotation threshold between target orientation and reached
+ orientation. The metric is q^T * q, where q is the quaternion difference between
+ target and reached orientation. The metric is not easy to interpret and a future
+ release will provide a more intuitive metric. For now, use 0.05 as a good value.
+ cspace_threshold: Joint space threshold in radians for revolute joints and meters for
+ linear joints between reached joint configuration and target joint configuration
+ used to measure success. Default of 0.05 has been found to be a good value for most
+ cases.
+ world_coll_checker: Instance of world collision checker to use for motion generation.
+ Leaving this to None will create a new instance of world collision checker using
+ the provided attr:`world_model`.
+ base_cfg_file: Base configuration file containing convergence and constraint criteria
+ to measure success.
+ particle_file: Optimizer configuration file to use for particle-based
+ optimization during trajectory optimization.
+ gradient_file: Optimizer configuration file to use for gradient-based
+ optimization during trajectory optimization.
+ trajopt_tsteps: Number of waypoints to use for trajectory optimization. Default of 32
+ is found to be a good number for most cases.
+ interpolation_type: Interpolation type to use for generating dense waypoints from
+ optimized trajectory. Default of
+ :py:attr:`curobo.util.trajectory.InterpolateType.LINEAR_CUDA` is found to be a
+ good choice for most cases. Other suitable options for real robot execution are
+ :py:attr:`curobo.util.trajectory.InterpolateType.QUINTIC` and
+ :py:attr:`curobo.util.trajectory.InterpolateType.CUBIC`.
+ interpolation_steps: Buffer size to use for storing interpolated trajectory. Default of
+ 5000 is found to be a good number for most cases.
+ interpolation_dt: Time step in seconds to use for generating interpolated trajectory
+ from optimized trajectory. Change this if you want to generate a trajectory with
+ a fixed timestep between waypoints.
+ minimum_trajectory_dt: Minimum time step in seconds allowed for trajectory
+ optimization.
+ use_cuda_graph: Record compute ops as cuda graphs and replay recorded graphs where
+ implemented. This can speed up execution by upto 10x. Default of True is
+ recommended. Enabling this will prevent changing solve type or batch size
+ after the first call to the solver.
+ self_collision_check: Enable self collision checks for generated motions. Default of
+ True is recommended. Set this to False to debug planning failures. Setting this to
+ False will also set self_collision_opt to False.
+ self_collision_opt: Enable self collision cost during optimization (IK, TrajOpt).
+ Default of True is recommended.
+ grad_trajopt_iters: Number of L-BFGS iterations to run trajectory optimization.
+ num_seeds: Number of seeds to use for trajectory optimization per problem.
+ seed_ratio: Ratio of linear and bias seeds to use for trajectory optimization.
+ Linear seed will generate a linear interpolated trajectory from start state
+ to IK solutions. Bias seed will add a mid-waypoint through the retract
+ configuration. Default of 1.0 linear and 0.0 bias is recommended. This can be
+ changed to 0.5 linear and 0.5 bias, along with changing num_seeds to 2.
+ trajopt_particle_opt: Enable particle-based optimization during trajectory
+ optimization. Default of True is recommended as particle-based optimization moves
+ the interpolated seeds away from bad local minima.
+ collision_checker_type: Type of collision checker to use for motion generation. Default
+ of CollisionCheckerType.MESH supports world represented by Cuboids and Meshes. See
+ :ref:`world_collision` for more details.
+ traj_evaluator_config: Configuration for trajectory evaluator. Default of None will
+ create a new instance of TrajEvaluatorConfig. After trajectory optimization across
+ many seeds, the best trajectory is selected based on this configuration. This
+ evaluator also checks if the optimized dt is within
+ :py:attr:`curobo.wrap.reacher.evaluator.TrajEvaluatorConfig.max_dt`. This check is
+ needed to measure smoothness of the optimized trajectory as bad trajectories can
+ have very high dt to fit within velocity, acceleration, and jerk limits.
+ traj_evaluator: Instance of trajectory evaluator to use for trajectory optimization.
+ Default of None will create a new instance of TrajEvaluator. In case you want to
+ use a custom evaluator, you can create a child instance of TrajEvaluator and
+ pass it.
+ minimize_jerk: Minimize jerk as regularization during trajectory optimizaiton.
+ use_gradient_descent: Use gradient descent instead of L-BFGS for trajectory
+ optimization. Default of False is recommended. Set to True for debugging gradients
+ of new cost terms.
+ collision_cache: Cache of obstacles to create to load obstacles between planning calls.
+ An example: ``{"obb": 10, "mesh": 10}``, to create a cache of 10 cuboids and 10
+ meshes.
+ n_collision_envs: Number of collision environments to create for batched optimization
+ across different environments. Only used for
+ :py:meth:`TrajOptSolver.solve_batch_env`and
+ :py:meth:`TrajOptSolver.solve_batch_env_goalset`.
+ n_collision_envs: Number of collision environments to create for batched planning
+ across different environments. Only used for :py:meth:`MotionGen.plan_batch_env`
+ and :py:meth:`MotionGen.plan_batch_env_goalset`.
+ use_es: Use Evolution Strategies for optimization. Default of None will use MPPI.
+ es_learning_rate: Learning rate to use for Evolution Strategies.
+ use_fixed_samples: Use fixed samples for MPPI. Setting to False will increase compute
+ time as new samples are generated for each iteration of MPPI.
+ aux_rollout: Rollout instance to use for auxiliary rollouts.
+ evaluate_interpolated_trajectory: Evaluate interpolated trajectory after optimization.
+ Default of True is recommended to ensure the optimized trajectory is not passing
+ through very thin obstacles.
+ fixed_iters: Use fixed number of iterations of L-BFGS for trajectory
+ optimization. Default of None will use the setting from the optimizer
+ configuration. In most cases, fixed iterations of solvers are run as current
+ solvers treat constraints as costs and there is no guarantee that the constraints
+ will be satisfied. Instead of checking constraints between iterations of a solver
+ and exiting, it's computationally cheaper to run a fixed number of iterations. In
+ addition, running fixed iterations of solvers is more robust to outlier problems.
+ store_debug: Store debugging information such as values of optimization
+ variables in TrajOpt result. Setting this to True will set :attr:`use_cuda_graph`
+ to False.
+ sync_cuda_time: Synchronize with host using :py:func:`torch.cuda.synchronize` before
+ measuring compute time.
+ collision_activation_distance: Distance in meters to activate collision cost. A good
+ value to start with is 0.01 meters. Increase the distance if the robot needs to
+ stay further away from obstacles.
+ trajopt_dt: Time step in seconds to use for trajectory optimization. A good value to
+ start with is 0.15 seconds. This value is used to compute velocity, acceleration,
+ and jerk values for waypoints through finite difference.
+ trim_steps: Trim waypoints from optimized trajectory. The optimized trajectory will
+ contain the start state at index 0 and have the last two waypoints be the same
+ as T-2 as trajectory optimization implicitly optimizes for zero acceleration and
+ velocity at the last waypoint. An example: ``[1,-2]`` will trim the first waypoint
+ and last 3 waypoints from the optimized trajectory.
+ store_debug_in_result: Store debugging information in MotionGenResult. This value is
+ set to True if either store_ik_debug or store_trajopt_debug is set to True.
+ smooth_weight: Override smooth weight for trajectory optimization. It's not recommended
+ to set this value for most cases.
+ state_finite_difference_mode: Finite difference mode to use for computing velocity,
+ acceleration, and jerk values. Default of None will use the setting from the
+ optimizer configuration file. The default finite difference method is a five
+ point stencil to compute the derivatives as this is accurate and provides
+ faster convergence compared to backward or central difference methods.
+ filter_robot_command: Filter generated trajectory to remove finite difference
+ artifacts. Default of True is recommended.
+ optimize_dt: Optimize dt during trajectory optimization. Default of True is
+ recommended to find time-optimal trajectories. Setting this to False will use the
+ provided :attr:`trajopt_dt` for trajectory optimization. Setting to False is
+ required when optimizing from a non-static start state.
+ project_pose_to_goal_frame: Project pose to goal frame when calculating distance
+ between reached and goal pose. Use this to constrain motion to specific axes
+ either in the global frame or the goal frame.
+
+ Returns:
+ TrajOptSolverConfig: Trajectory optimization configuration.
+ """
+
+ if minimum_trajectory_dt is None:
+ minimum_trajectory_dt = interpolation_dt
+ elif minimum_trajectory_dt < interpolation_dt:
+ log_error("minimum_trajectory_dt cannot be lower than interpolation_dt")
if isinstance(robot_cfg, str):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"]
@@ -291,7 +458,6 @@ class TrajOptSolverConfig:
elif use_particle_opt:
mppi_cfg = ParallelMPPIConfig(**config_dict)
parallel_mppi = ParallelMPPI(mppi_cfg)
-
config_dict = LBFGSOptConfig.create_data_dict(
grad_config_data["lbfgs"], arm_rollout_grad, tensor_args
)
@@ -323,7 +489,7 @@ class TrajOptSolverConfig:
trajopt = WrapBase(cfg)
if traj_evaluator_config is None:
traj_evaluator_config = TrajEvaluatorConfig.from_basic(
- min_dt=interpolation_dt,
+ min_dt=minimum_trajectory_dt,
dof=robot_cfg.kinematics.dof,
tensor_args=tensor_args,
)
@@ -358,7 +524,9 @@ class TrajOptSolverConfig:
@dataclass
-class TrajResult(Sequence):
+class TrajOptResult(Sequence):
+ """Trajectory optimization result."""
+
success: T_BValue_bool
goal: Goal
solution: JointState
@@ -379,9 +547,15 @@ class TrajResult(Sequence):
goalset_index: Optional[torch.Tensor] = None
optimized_seeds: Optional[torch.Tensor] = None
- def __getitem__(self, idx):
- # position_error = rotation_error = cspace_error = path_buffer_last_tstep = None
- # metrics = interpolated_solution = None
+ def __getitem__(self, idx: int) -> TrajOptResult:
+ """Get item at index.
+
+ Args:
+ idx: Index of the item to get.
+
+ Returns:
+ TrajOptResult: Trajectory optimization result at the given index.
+ """
d_list = [
self.interpolated_solution,
@@ -394,7 +568,7 @@ class TrajResult(Sequence):
]
idx_vals = list_idx_if_not_none(d_list, idx)
- return TrajResult(
+ return TrajOptResult(
goal=self.goal[idx],
solution=self.solution[idx],
success=self.success[idx],
@@ -411,15 +585,38 @@ class TrajResult(Sequence):
optimized_seeds=self.optimized_seeds,
)
- def __len__(self):
+ def __len__(self) -> int:
+ """Get length of the TrajOptResult."""
return self.success.shape[0]
+@dataclass
+class TrajResult(TrajOptResult):
+ """Deprecated: Use TrajOptResult instead of TrajResult"""
+
+ def __post_init__(self):
+ """post-init function for TrajResult"""
+ log_warn("Deprecated: Use TrajOptResult instead of TrajResult")
+
+
class TrajOptSolver(TrajOptSolverConfig):
+ """Trajectory Optimization Solver class for generating minimum-jerk trajectories.
+
+ Trajectory Optimization uses joint positions as optimization variables with cost terms for
+ avoiding world collisions, self-collisions, and joint limits. The joint velocities, accelerations,
+ and jerks are computed using five point stencil. A squared l2-norm cost term on joint accelerations
+ and jerks is used to encourage smooth trajectories. A cost term for the terminal state to reach
+ either a Cartesian pose or joint configuration is also used. Read :ref:`research_page` for
+ more details.
+ """
+
def __init__(self, config: TrajOptSolverConfig) -> None:
+ """Initialize TrajOptSolver with configuration parameters.
+
+ Args:
+ config: Configuration parameters for TrajOptSolver.
+ """
super().__init__(**vars(config))
- self.dof = self.rollout_fn.d_action
- self.action_horizon = self.rollout_fn.action_horizon
self.delta_vec = interpolate_kernel(2, self.action_horizon, self.tensor_args).unsqueeze(0)
self.waypoint_delta_vec = interpolate_kernel(
@@ -427,13 +624,11 @@ class TrajOptSolver(TrajOptSolverConfig):
).unsqueeze(0)
assert self.action_horizon / 2 != 0.0
self.solver.update_nproblems(self.num_seeds)
- self._max_joint_vel = (
- self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[1, :].reshape(
- 1, 1, self.dof
- )
- ) - 0.02
- self._max_joint_acc = self.rollout_fn.state_bounds.acceleration[1, :] - 0.02
- self._max_joint_jerk = self.rollout_fn.state_bounds.jerk[1, :] - 0.02
+ self._max_joint_vel = self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[
+ 1, :
+ ].reshape(1, 1, self.dof)
+ self._max_joint_acc = self.rollout_fn.state_bounds.acceleration[1, :]
+ self._max_joint_jerk = self.rollout_fn.state_bounds.jerk[1, :]
self._num_seeds = -1
self._col = None
if self.traj_evaluator is None:
@@ -450,6 +645,13 @@ class TrajOptSolver(TrajOptSolverConfig):
self._rollout_list = None
def get_all_rollout_instances(self) -> List[RolloutBase]:
+ """Get all rollout instances in the solver.
+
+ Useful to update parameters across all rollout instances.
+
+ Returns:
+ List[RolloutBase]: List of all rollout instances.
+ """
if self._rollout_list is None:
self._rollout_list = [
self.rollout_fn,
@@ -458,32 +660,67 @@ class TrajOptSolver(TrajOptSolverConfig):
return self._rollout_list
def get_all_kinematics_instances(self) -> List[CudaRobotModel]:
+ """Get all kinematics instances used across components in motion generation.
+
+ This is deprecated. Use :meth:`TrajOptSolver.kinematics` instead as TrajOptSolver now uses
+ a shared kinematics instance across all components.
+
+ Returns:
+ List[CudaRobotModel]: Single kinematics instance, returned as a list for compatibility.
+ """
+ log_warn(
+ "Deprecated: Use TrajOptSolver.kinematics instead as TrajOptSolver now uses a "
+ + "shared kinematics instance across all components."
+ )
if self._kin_list is None:
self._kin_list = [
i.dynamics_model.robot_model for i in self.get_all_rollout_instances()
]
return self._kin_list
- def attach_object_to_robot(
+ def attach_spheres_to_robot(
self,
sphere_radius: float,
sphere_tensor: Optional[torch.Tensor] = None,
link_name: str = "attached_object",
) -> None:
- for k in self.get_all_kinematics_instances():
- k.attach_object(
- sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name
- )
+ """Attach spheres to robot for collision checking.
- def detach_object_from_robot(self, link_name: str = "attached_object") -> None:
- for k in self.get_all_kinematics_instances():
- k.detach_object(link_name)
+ To fit spheres to an obstacle, see
+ :py:meth:`~curobo.geom.types.Obstacle.get_bounding_spheres`
- def update_goal_buffer(
+ Args:
+ sphere_radius: Radius of the spheres. Set to None if :attr:`sphere_tensor` is provided.
+ sphere_tensor: Sphere x, y, z, r tensor.
+ link_name: Name of the link to attach the spheres to. Note that this link should
+ already have pre-allocated spheres.
+ """
+ self.kinematics.attach_object(
+ sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name
+ )
+
+ def detach_spheres_from_robot(self, link_name: str = "attached_object") -> None:
+ """Detach spheres from robot.
+
+ Args:
+ link_name: Name of the link to detach the spheres from.
+ """
+ self.kinematics.detach_object(link_name)
+
+ def _update_solve_state_and_goal_buffer(
self,
solve_state: ReacherSolveState,
goal: Goal,
):
+ """Update goal buffer and solve state of current trajectory optimization problem.
+
+ Args:
+ solve_state: New solve state.
+ goal: New goal buffer.
+
+ Returns:
+ Goal: Updated goal buffer with augmentations for new solve state.
+ """
self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal(
goal,
self._solve_state,
@@ -514,7 +751,28 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: Optional[int] = None,
seed_success: Optional[torch.Tensor] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem with any solve type.
+
+ Args:
+ solve_type: Type of solve to perform.
+ goal: Goal to reach.
+ seed_traj: Seed trajectory to start optimization from. This should be
+ of shape [num_seeds, batch_size, action_horizon, dof]. If None, linearly
+ interpolated seeds from current state to goal state are used. If goal.goal_state
+ is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. This cannot be changed after the
+ first call to solve as CUDA graph re-creation is currently not supported.
+ seed_success: Success of seeds. This is used to filter out successful seeds from
+ :attr:`seed_traj`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number
+ of iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if solve_type == ReacherSolveType.SINGLE:
return self.solve_single(
goal,
@@ -574,7 +832,7 @@ class TrajOptSolver(TrajOptSolverConfig):
newton_iters=newton_iters,
)
- def solve_from_solve_state(
+ def _solve_from_solve_state(
self,
solve_state: ReacherSolveState,
goal: Goal,
@@ -585,13 +843,39 @@ class TrajOptSolver(TrajOptSolverConfig):
seed_success: Optional[torch.Tensor] = None,
newton_iters: Optional[int] = None,
):
+ """Solve trajectory optimization problem with a given solve state.
+
+ Args:
+ solve_state: Solve state for the optimization problem.
+ goal: Goal object containing target pose or joint configuration.
+ seed_traj: Seed trajectory to start optimization from. This should be of
+ shape [num_seeds, batch_size, action_horizon, dof]. If None, linearly
+ interpolated seeds from current state to goal state are used. If goal.goal_state
+ is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. This cannot be changed after the
+ first call to solve as CUDA graph re-creation is currently not supported.
+ seed_success: Success of seeds. This is used to filter out successful seeds from
+ :attr:`seed_traj`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if solve_state.batch_env:
if solve_state.batch_size > self.world_coll_checker.n_envs:
- raise ValueError("Batch Env is less that goal batch")
+ log_error("Batch Env is less that goal batch")
if newton_iters is not None:
self.solver.newton_optimizer.outer_iters = newton_iters
self.solver.newton_optimizer.fixed_iters = True
- goal_buffer = self.update_goal_buffer(solve_state, goal)
+ log_info("TrajOpt: solving with Pose batch:" + str(goal.batch))
+ goal_buffer = self._update_solve_state_and_goal_buffer(solve_state, goal)
# if self.evaluate_interpolated_trajectory:
self.interpolate_rollout.update_params(goal_buffer)
# get seeds:
@@ -628,7 +912,31 @@ class TrajOptSolver(TrajOptSolverConfig):
return_all_solutions: bool = False,
num_seeds: Optional[int] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem for a single goal.
+
+ This will use multiple seeds internally and return the best solution.
+
+ Args:
+ goal: Goal to reach.
+ seed_traj: Seed trajectory to start optimization from. This should be of shape
+ [:attr:`num_seeds`, 1, :attr:`TrajOptSolver.action_horizon`,
+ :attr:`TrajOptSolver.dof`]. If None, linearly interpolated seeds from current state
+ to goal state are used. If goal.goal_state is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. If None, the number of seeds
+ is set to :attr:`TrajOptSolver.num_seeds`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if num_seeds is None:
num_seeds = self.num_seeds
solve_state = ReacherSolveState(
@@ -639,7 +947,7 @@ class TrajOptSolver(TrajOptSolverConfig):
n_goalset=1,
)
- return self.solve_from_solve_state(
+ return self._solve_from_solve_state(
solve_state,
goal,
seed_traj,
@@ -657,7 +965,29 @@ class TrajOptSolver(TrajOptSolverConfig):
return_all_solutions: bool = False,
num_seeds: Optional[int] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem that uses goalset to represent Pose target.
+
+ Args:
+ goal: Goal to reach.
+ seed_traj: Seed trajectory to start optimization from. This should be of shape
+ [:attr:`num_seeds`, 1, :attr:`TrajOptSolver.action_horizon`,
+ :attr:`TrajOptSolver.dof`]. If None, linearly interpolated seeds from current state
+ to goal state are used. If goal.goal_state is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. If None, the number of seeds
+ is set to :attr:`TrajOptSolver.num_seeds`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if num_seeds is None:
num_seeds = self.num_seeds
solve_state = ReacherSolveState(
@@ -667,7 +997,7 @@ class TrajOptSolver(TrajOptSolverConfig):
n_envs=1,
n_goalset=goal.n_goalset,
)
- return self.solve_from_solve_state(
+ return self._solve_from_solve_state(
solve_state,
goal,
seed_traj,
@@ -686,7 +1016,31 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: Optional[int] = None,
seed_success: Optional[torch.Tensor] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem for a batch of goals.
+
+ Args:
+ goal: Batch of goals to reach, this includes batch of current states.
+ seed_traj: Seed trajectory to start optimization from. This should be of shape
+ [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If
+ None, linearly interpolated seeds from current state to goal state are used. If
+ goal.goal_state is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. If None, the number of seeds
+ is set to :attr:`TrajOptSolver.num_seeds`.
+ seed_success: Success of seeds. This is used to filter out successful seeds from
+ :attr:`seed_traj`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if num_seeds is None:
num_seeds = self.num_seeds
solve_state = ReacherSolveState(
@@ -696,7 +1050,7 @@ class TrajOptSolver(TrajOptSolverConfig):
n_envs=1,
n_goalset=1,
)
- return self.solve_from_solve_state(
+ return self._solve_from_solve_state(
solve_state,
goal,
seed_traj,
@@ -716,7 +1070,31 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: Optional[int] = None,
seed_success: Optional[torch.Tensor] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem for a batch of Poses with goalset.
+
+ Args:
+ goal: Batch of goals to reach, this includes batch of current states.
+ seed_traj: Seed trajectory to start optimization from. This should be of shape
+ [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If
+ None, linearly interpolated seeds from current state to goal state are used. If
+ goal.goal_state is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. If None, the number of seeds
+ is set to :attr:`TrajOptSolver.num_seeds`.
+ seed_success: Success of seeds. This is used to filter out successful seeds from
+ :attr:`seed_traj`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if num_seeds is None:
num_seeds = self.num_seeds
solve_state = ReacherSolveState(
@@ -726,7 +1104,7 @@ class TrajOptSolver(TrajOptSolverConfig):
n_envs=1,
n_goalset=goal.n_goalset,
)
- return self.solve_from_solve_state(
+ return self._solve_from_solve_state(
solve_state,
goal,
seed_traj,
@@ -745,7 +1123,31 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: Optional[int] = None,
seed_success: Optional[torch.Tensor] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem in a batch of environments.
+
+ Args:
+ goal: Batch of goals to reach, this includes batch of current states.
+ seed_traj: Seed trajectory to start optimization from. This should be of shape
+ [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If
+ None, linearly interpolated seeds from current state to goal state are used. If
+ goal.goal_state is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. If None, the number of seeds
+ is set to :attr:`TrajOptSolver.num_seeds`.
+ seed_success: Success of seeds. This is used to filter out successful seeds from
+ :attr:`seed_traj`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if num_seeds is None:
num_seeds = self.num_seeds
solve_state = ReacherSolveState(
@@ -755,7 +1157,7 @@ class TrajOptSolver(TrajOptSolverConfig):
n_envs=goal.batch,
n_goalset=1,
)
- return self.solve_from_solve_state(
+ return self._solve_from_solve_state(
solve_state,
goal,
seed_traj,
@@ -775,7 +1177,31 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: Optional[int] = None,
seed_success: Optional[torch.Tensor] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
+ ) -> TrajOptResult:
+ """Solve trajectory optimization problem in a batch of environments with goalset.
+
+ Args:
+ goal: Batch of goals to reach, this includes batch of current states.
+ seed_traj: Seed trajectory to start optimization from. This should be of shape
+ [:attr:`num_seeds`, :attr:`goal.batch`, :attr:`TrajOptSolver.action_horizon`]. If
+ None, linearly interpolated seeds from current state to goal state are used. If
+ goal.goal_state is empty, random seeds are generated.
+ use_nn_seed: Use neural network seed for optimization. This is not implemented.
+ return_all_solutions: Return solutions for all seeds.
+ num_seeds: Number of seeds to use for optimization. If None, the number of seeds
+ is set to :attr:`TrajOptSolver.num_seeds`.
+ seed_success: Success of seeds. This is used to filter out successful seeds from
+ :attr:`seed_traj`.
+ newton_iters: Number of iterations to run LBFGS optimization. If None, the number of
+ iterations is set to the default value in :attr:`TrajOptSolver.newton_iters`. This
+ is the outer iterations, where each outer iteration will run 25 inner iterations
+ of LBFGS optimization captured in a CUDA-Graph. Total number of optimization
+ iterations is 25 * outer_iters. The number of inner iterations can be changed
+ with :py:attr:`curobo.opt.newton.lbfgs.LBFGSOptConfig.inner_iters`.
+
+ Returns:
+ TrajOptResult: Result of the trajectory optimization.
+ """
if num_seeds is None:
num_seeds = self.num_seeds
solve_state = ReacherSolveState(
@@ -785,7 +1211,7 @@ class TrajOptSolver(TrajOptSolverConfig):
n_envs=goal.batch,
n_goalset=goal.n_goalset,
)
- return self.solve_from_solve_state(
+ return self._solve_from_solve_state(
solve_state,
goal,
seed_traj,
@@ -803,21 +1229,11 @@ class TrajOptSolver(TrajOptSolverConfig):
return_all_solutions: bool = False,
num_seeds: Optional[int] = None,
newton_iters: Optional[int] = None,
- ) -> TrajResult:
- """Only for single goal
-
- Args:
- goal (Goal): _description_
- seed_traj (Optional[JointState], optional): _description_. Defaults to None.
- use_nn_seed (bool, optional): _description_. Defaults to False.
-
- Raises:
- NotImplementedError: _description_
-
- Returns:
- TrajResult: _description_
- """
- log_warn("TrajOpt.solve() is deprecated, use TrajOpt.solve_single or others instead")
+ ) -> TrajOptResult:
+ """Deprecated: Use :meth:`TrajOptSolver.solve_single` or others instead."""
+ log_warn(
+ "TrajOptSolver.solve is deprecated, use TrajOptSolver.solve_single or others instead"
+ )
if goal.goal_pose.batch == 1 and goal.goal_pose.n_goalset == 1:
return self.solve_single(
goal,
@@ -858,12 +1274,29 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: int,
batch_mode: bool = False,
):
+ """Get result from the optimization problem.
+
+ Args:
+ result: _description_
+ return_all_solutions: _description_
+ goal: _description_
+ seed_traj: _description_
+ num_seeds: _description_
+ batch_mode: _description_
+
+ Raises:
+ log_error: _description_
+ ValueError: _description_
+
+ Returns:
+ _description_
+ """
st_time = time.time()
if self.trim_steps is not None:
result.action = result.action.trim_trajectory(self.trim_steps[0], self.trim_steps[1])
interpolated_trajs, last_tstep, opt_dt = self.get_interpolated_trajectory(result.action)
if self.sync_cuda_time:
- torch.cuda.synchronize()
+ torch.cuda.synchronize(device=self.tensor_args.device)
interpolation_time = time.time() - st_time
if self.evaluate_interpolated_trajectory:
with profiler.record_function("trajopt/evaluate_interpolated"):
@@ -876,7 +1309,6 @@ class TrajOptSolver(TrajOptSolverConfig):
result.metrics.rotation_error = metrics.rotation_error
result.metrics.cspace_error = metrics.cspace_error
result.metrics.goalset_index = metrics.goalset_index
-
st_time = time.time()
if result.metrics.cspace_error is None and result.metrics.position_error is None:
raise log_error("convergence check requires either goal_pose or goal_state")
@@ -890,6 +1322,7 @@ class TrajOptSolver(TrajOptSolverConfig):
self.rotation_threshold,
self.cspace_threshold,
)
+
if False:
feasible = torch.all(result.metrics.feasible, dim=-1)
@@ -905,7 +1338,7 @@ class TrajOptSolver(TrajOptSolverConfig):
success = torch.logical_and(feasible, converge)
if return_all_solutions:
- traj_result = TrajResult(
+ traj_result = TrajOptResult(
success=success,
goal=goal,
solution=result.action.scale_by_dt(self.solver_dt_tensor, opt_dt.view(-1, 1, 1)),
@@ -976,7 +1409,7 @@ class TrajOptSolver(TrajOptSolverConfig):
interpolated_traj = interpolated_trajs[idx]
if self.sync_cuda_time:
- torch.cuda.synchronize()
+ torch.cuda.synchronize(device=self.tensor_args.device)
if len(best_act_seq.shape) == 3:
opt_dt_v = opt_dt.view(-1, 1, 1)
else:
@@ -990,8 +1423,7 @@ class TrajOptSolver(TrajOptSolverConfig):
"interpolation_time": interpolation_time,
"select_time": select_time,
}
-
- traj_result = TrajResult(
+ traj_result = TrajOptResult(
success=success,
goal=goal,
solution=opt_solution,
@@ -1022,21 +1454,11 @@ class TrajOptSolver(TrajOptSolverConfig):
use_nn_seed: bool = False,
return_all_solutions: bool = False,
num_seeds: Optional[int] = None,
- ) -> TrajResult:
- """Only for single goal
-
- Args:
- goal (Goal): _description_
- seed_traj (Optional[JointState], optional): _description_. Defaults to None.
- use_nn_seed (bool, optional): _description_. Defaults to False.
-
- Raises:
- NotImplementedError: _description_
-
- Returns:
- TrajResult: _description_
- """
- log_warn("TrajOpt.solve() is deprecated, use TrajOpt.solve_single or others instead")
+ ) -> TrajOptResult:
+ """Deprecated: Use :meth:`TrajOptSolver.solve_batch` or others instead."""
+ log_warn(
+ "TrajOptSolver.batch_solve is deprecated, use TrajOptSolver.solve_batch or others instead"
+ )
if goal.n_goalset == 1:
return self.solve_batch(
goal, seed_traj, use_nn_seed, return_all_solutions, num_seeds, seed_success
@@ -1046,7 +1468,16 @@ class TrajOptSolver(TrajOptSolverConfig):
goal, seed_traj, use_nn_seed, return_all_solutions, num_seeds, seed_success
)
- def get_linear_seed(self, start_state, goal_state):
+ def get_linear_seed(self, start_state, goal_state) -> torch.Tensor:
+ """Get linearly interpolated seeds from start states to goal states.
+
+ Args:
+ start_state: start state of the robot.
+ goal_state: goal state of the robot.
+
+ Returns:
+ torch.Tensor: Linearly interpolated seeds.
+ """
start_q = start_state.position.view(-1, 1, self.dof)
end_q = goal_state.position.view(-1, 1, self.dof)
edges = torch.cat((start_q, end_q), dim=1)
@@ -1054,13 +1485,29 @@ class TrajOptSolver(TrajOptSolverConfig):
seed = self.delta_vec @ edges
return seed
- def get_start_seed(self, start_state):
+ def get_start_seed(self, start_state) -> torch.Tensor:
+ """Get trajectory seeds with start state repeated.
+
+ Args:
+ start_state: start state of the robot.
+
+ Returns:
+ torch.Tensor: Trajectory seeds with start state repeated.
+ """
start_q = start_state.position.view(-1, 1, self.dof)
edges = torch.cat((start_q, start_q), dim=1)
seed = self.delta_vec @ edges
return seed
- def _get_seed_numbers(self, num_seeds):
+ def _get_seed_numbers(self, num_seeds: int) -> Dict[str, int]:
+ """Get number of seeds for each seed type.
+
+ Args:
+ num_seeds: Total number of seeds to generate.
+
+ Returns:
+ Dict[str, int]: Number of seeds for each seed type.
+ """
n_seeds = {"linear": 0, "bias": 0, "start": 0, "goal": 0}
k = n_seeds.keys
t_seeds = 0
@@ -1082,10 +1529,21 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds: Optional[int] = None,
batch_mode: bool = False,
):
- # if batch_mode:
+ """Get seed set for optimization.
+
+ Args:
+ goal: Goal object containing target pose or joint configuration.
+ batch: _description_
+ h: _description_
+ seed_traj: _description_
+ dofseed_success: _description_
+ n_seedsnum_seeds: _description_
+ batch_mode: _description_
+
+ Returns:
+ _description_
+ """
total_seeds = goal.batch * num_seeds
- # else:
- # total_seeds = num_seeds
if isinstance(seed_traj, JointState):
seed_traj = seed_traj.position
@@ -1093,7 +1551,6 @@ class TrajOptSolver(TrajOptSolverConfig):
if goal.goal_state is not None and self.use_cspace_seed:
# get linear seed
seed_traj = self.get_seeds(goal.current_state, goal.goal_state, num_seeds=num_seeds)
- # .view(batch_size, self.num_seeds, self.action_horizon, self.dof)
else:
# get start repeat seed:
log_info("No goal state found, using current config to seed")
@@ -1109,14 +1566,50 @@ class TrajOptSolver(TrajOptSolverConfig):
new_seeds = self.get_seeds(
goal.current_state, goal.goal_state, num_seeds - seed_traj.shape[0]
)
- seed_traj = torch.cat((seed_traj, new_seeds), dim=0) # n_seed, batch, h, dof
+ seed_traj = torch.cat((seed_traj, new_seeds), dim=0)
- seed_traj = seed_traj.view(
- total_seeds, self.action_horizon, self.dof
- ) # n_seeds,batch, h, dof
+ if len(seed_traj.shape) == 3:
+ if (
+ seed_traj.shape[0] != total_seeds
+ or seed_traj.shape[1] != self.action_horizon
+ or seed_traj.shape[2] != self.dof
+ ):
+ log_error(
+ "Seed traj shape should be [num_seeds * batch, action_horizon, dof]"
+ + " current shape is "
+ + str(seed_traj.shape)
+ )
+ elif len(seed_traj.shape) == 4:
+ if (
+ seed_traj.shape[0] * seed_traj.shape[1] != total_seeds
+ or seed_traj.shape[2] != self.action_horizon
+ or seed_traj.shape[3] != self.dof
+ ):
+ log_error(
+ "Seed traj shape should be [num_seeds, batch, action_horizon, dof]"
+ + " or [1, num_seeds * batch, action_horizon, dof]"
+ + " current shape is "
+ + str(seed_traj.shape)
+ )
+ else:
+ log_error("Seed traj shape should have 3 or 4 dimensions: " + str(seed_traj.shape))
+ seed_traj = seed_traj.view(total_seeds, self.action_horizon, self.dof)
return seed_traj
- def get_seeds(self, start_state, goal_state, num_seeds=None):
+ def get_seeds(
+ self, start_state: JointState, goal_state: JointState, num_seeds: int = None
+ ) -> torch.Tensor:
+ """Get seed trajectories for optimization.
+
+ Args:
+ start_state: Start state of the robot.
+ goal_state: Goal state of the robot.
+ num_seeds: Number of seeds to generate. If None, the number of seeds is set to
+ :attr:`TrajOptSolver.num_seeds`.
+
+ Returns:
+ torch.Tensor: Seed trajectories of shape [num_seeds, batch, action_horizon, dof]
+ """
# repeat seeds:
if num_seeds is None:
num_seeds = self.num_seeds
@@ -1152,24 +1645,45 @@ class TrajOptSolver(TrajOptSolverConfig):
1, n_seeds["goal"], 1, 1
)
seed_set.append(bias_seeds)
- all_seeds = torch.cat(
- seed_set, dim=1
- ) # .#transpose(0,1).contiguous() # n_seed, batch, h, dof
+ all_seeds = torch.cat(seed_set, dim=1)
return all_seeds
- def get_bias_seed(self, start_state, goal_state):
+ def get_bias_seed(self, start_state, goal_state) -> torch.Tensor:
+ """Get seed trajectories that pass through the retract configuration at mid-waypoint.
+
+ Args:
+ start_state: start state of the robot.
+ goal_state: goal state of the robot.
+
+ Returns:
+ torch.Tensor: Seed trajectories of shape [num_seeds * batch, action_horizon, dof].
+ """
start_q = start_state.position.view(-1, 1, self.dof)
end_q = goal_state.position.view(-1, 1, self.dof)
bias_q = self.bias_node.view(-1, 1, self.dof).repeat(start_q.shape[0], 1, 1)
edges = torch.cat((start_q, bias_q, end_q), dim=1)
seed = self.waypoint_delta_vec @ edges
- # print(seed)
return seed
@profiler.record_function("trajopt/interpolation")
- def get_interpolated_trajectory(self, traj_state: JointState):
+ def get_interpolated_trajectory(
+ self, traj_state: JointState
+ ) -> Tuple[JointState, torch.Tensor]:
+ """Get interpolated trajectory from optimized trajectories.
+
+ This function will first find the optimal dt for each trajectory in the batch by scaling
+ the trajectories to joint limits. Then it will interpolate the trajectory using the optimal
+ dt.
+
+ Args:
+ traj_state: Optimized trajectories of shape [num_seeds * batch, action_horizon, dof].
+
+ Returns:
+ Tuple[JointState, torch.Tensor, torch.Tensor]: Interpolated trajectory, last time step
+ for each trajectory in batch, optimal dt for each trajectory in batch.
+ """
# do interpolation:
if (
self._interpolated_traj_buffer is None
@@ -1182,24 +1696,34 @@ class TrajOptSolver(TrajOptSolverConfig):
self._interpolated_traj_buffer.joint_names = self.rollout_fn.joint_names
state, last_tstep, opt_dt = get_batch_interpolated_trajectory(
traj_state,
+ self.solver_dt_tensor,
self.interpolation_dt,
self._max_joint_vel,
self._max_joint_acc,
self._max_joint_jerk,
- self.solver_dt_tensor,
kind=self.interpolation_type,
tensor_args=self.tensor_args,
out_traj_state=self._interpolated_traj_buffer,
min_dt=self.traj_evaluator_config.min_dt,
+ max_dt=self.traj_evaluator_config.max_dt,
optimize_dt=self.optimize_dt,
)
-
return state, last_tstep, opt_dt
def calculate_trajectory_dt(
self,
trajectory: JointState,
+ epsilon: float = 1e-6,
) -> torch.Tensor:
+ """Calculate the optimal dt for a given trajectory by scaling it to joint limits.
+
+ Args:
+ trajectory: Trajectory to calculate optimal dt for.
+ epsilon: Small value to improve numerical stability.
+
+ Returns:
+ torch.Tensor: Optimal trajectory dt.
+ """
opt_dt = calculate_dt_no_clamp(
trajectory.velocity,
trajectory.acceleration,
@@ -1207,42 +1731,78 @@ class TrajOptSolver(TrajOptSolverConfig):
self._max_joint_vel,
self._max_joint_acc,
self._max_joint_jerk,
+ epsilon=epsilon,
)
return opt_dt
def reset_seed(self):
+ """Reset the seed for random number generators in MPPI and rollout functions."""
self.solver.reset_seed()
def reset_cuda_graph(self):
+ """Clear all recorded CUDA graphs. This does not work."""
self.solver.reset_cuda_graph()
self.interpolate_rollout.reset_cuda_graph()
self.rollout_fn.reset_cuda_graph()
def reset_shape(self):
+ """Reset the shape of the rollout function and the solver."""
self.solver.reset_shape()
self.interpolate_rollout.reset_shape()
self.rollout_fn.reset_shape()
@property
def kinematics(self) -> CudaRobotModel:
+ """Get the kinematics instance of the robot."""
return self.rollout_fn.dynamics_model.robot_model
@property
- def retract_config(self):
+ def retract_config(self) -> torch.Tensor:
+ """Get the retract/home configuration of the robot.
+
+ Returns:
+ torch.Tensor: Retract configuration of the robot.
+ """
return self.rollout_fn.dynamics_model.retract_config.view(1, -1)
def fk(self, q: torch.Tensor) -> CudaRobotModelState:
+ """Compute forward kinematics for the robot.
+
+ Args:
+ q: Joint configuration of the robot.
+
+ Returns:
+ CudaRobotModelState: Forward kinematics of the robot.
+ """
return self.kinematics.get_state(q)
@property
- def solver_dt(self):
+ def solver_dt(self) -> torch.Tensor:
+ """Get the current trajectory dt for the solver.
+
+ Returns:
+ torch.Tensor: Trajectory dt for the solver.
+ """
return self.solver.safety_rollout.dynamics_model.traj_dt[0]
- # return self.solver.safety_rollout.dynamics_model.dt_traj_params.base_dt
@property
- def solver_dt_tensor(self):
+ def solver_dt_tensor(self) -> torch.Tensor:
+ """Get the current trajectory dt for the solver.
+
+ Returns:
+ torch.Tensor: Trajectory dt for the solver.
+ """
return self.solver.safety_rollout.dynamics_model.traj_dt[0]
+ @property
+ def minimum_trajectory_dt(self) -> float:
+ """Get the minimum trajectory dt that is allowed, smaller dt will be clamped to this value.
+
+ Returns:
+ float: Minimum trajectory dt that is allowed.
+ """
+ return self.traj_evaluator.min_dt
+
def update_solver_dt(
self,
dt: Union[float, torch.Tensor],
@@ -1250,21 +1810,70 @@ class TrajOptSolver(TrajOptSolverConfig):
max_dt: Optional[float] = None,
base_ratio: Optional[float] = None,
):
+ """Update the trajectory dt for the solver.
+
+ This dt is used to calculate the velocity, acceleration and jerk of the trajectory through
+ five point stencil (finite difference).
+
+ Args:
+ dt: New trajectory dt.
+ base_dt: Base dt for the trajectory. This is not supported.
+ max_dt: Maximum dt for the trajectory. This is not supported.
+ base_ratio: Ratio in trajectory length to scale from base_dt to max_dt. This is not
+ supported.
+ """
all_rollouts = self.get_all_rollout_instances()
for rollout in all_rollouts:
rollout.update_traj_dt(dt, base_dt, max_dt, base_ratio)
- def compute_metrics(self, opt_trajectory: bool, interpolated_trajectory: bool):
- self.solver.compute_metrics = opt_trajectory
- self.evaluate_interpolated_trajectory = interpolated_trajectory
-
def get_full_js(self, active_js: JointState) -> JointState:
+ """Get full joint state from controlled joint state, appending locked joints.
+
+ Args:
+ active_js: Controlled joint state
+
+ Returns:
+ JointState: Full joint state.
+ """
return self.rollout_fn.get_full_dof_from_solution(active_js)
+ def get_active_js(
+ self,
+ in_js: JointState,
+ ):
+ """Get controlled joint state from input joint state.
+
+ This is used to get the joint state for only joints that are optimization variables. This
+ also re-orders the joints to match the order of optimization variables.
+
+ Args:
+ in_js: Input joint state.
+
+ Returns:
+ JointState: Active joint state.
+ """
+ opt_jnames = self.rollout_fn.joint_names
+ opt_js = in_js.get_ordered_joint_state(opt_jnames)
+ return opt_js
+
def update_pose_cost_metric(
self,
metric: PoseCostMetric,
):
+ """Update the pose cost metric for :ref:`tut_constrained_planning`.
+
+ Only supports for the main end-effector. Does not support for multiple links that are
+ specified with `link_poses` in planning methods.
+
+ Args:
+ metric: Type and parameters for pose constraint to add.
+ start_state: Start joint state for the constraint.
+ goal_pose: Goal pose for the constraint.
+
+ Returns:
+ bool: True if the constraint can be added, False otherwise.
+ """
+
rollouts = self.get_all_rollout_instances()
[
rollout.update_pose_cost_metric(metric)
@@ -1274,8 +1883,35 @@ class TrajOptSolver(TrajOptSolverConfig):
@property
def newton_iters(self):
+ """Get the number of newton outer iterations during L-BFGS optimization.
+
+ Returns:
+ int: Number of newton outer iterations during L-BFGS optimization.
+ """
return self._og_newton_iters
+ @property
+ def dof(self) -> int:
+ """Get the number of degrees of freedom of the robot.
+
+ Returns:
+ int: Number of degrees of freedom of the robot.
+ """
+ return self.rollout_fn.d_action
+
+ @property
+ def action_horizon(self) -> int:
+ """Get the action horizon of the trajectory optimization problem.
+
+ Number of actions in trajectory optimization can be smaller than number of waypoints as
+ the first waypoint is the current state of the robot and the last two waypoints are
+ the same as T-2 waypoint to implicitly enforce zero acceleration and zero velocity at T.
+
+ Returns:
+ int: Action horizon of the trajectory optimization problem.
+ """
+ return self.rollout_fn.action_horizon
+
@get_torch_jit_decorator()
def jit_feasible_success(
@@ -1287,6 +1923,7 @@ def jit_feasible_success(
rotation_threshold: float,
cspace_threshold: float,
):
+ """JIT function to check if the optimization is successful."""
feasible = torch.all(feasible, dim=-1)
converge = feasible
if position_error is not None and rotation_error is not None:
@@ -1318,6 +1955,7 @@ def jit_trajopt_best_select(
col,
opt_dt,
):
+ """JIT function to select the best solution from optimized seeds."""
success[~smooth_label] = False
convergence_error = 0
# get the best solution:
diff --git a/src/curobo/wrap/wrap_base.py b/src/curobo/wrap/wrap_base.py
index 447d7e6..37ed51f 100644
--- a/src/curobo/wrap/wrap_base.py
+++ b/src/curobo/wrap/wrap_base.py
@@ -127,6 +127,10 @@ class WrapBase(WrapConfig):
def rollout_fn(self):
return self.safety_rollout
+ @property
+ def tensor_args(self):
+ return self.safety_rollout.tensor_args
+
def solve(self, goal: Goal, seed: Optional[torch.Tensor] = None):
metrics = None
diff --git a/src/curobo/wrap/wrap_mpc.py b/src/curobo/wrap/wrap_mpc.py
index a2f2229..51dbc71 100644
--- a/src/curobo/wrap/wrap_mpc.py
+++ b/src/curobo/wrap/wrap_mpc.py
@@ -53,13 +53,13 @@ class WrapMpc(WrapBase):
self.update_params(goal)
if self.sync_cuda_time:
- torch.cuda.synchronize()
+ torch.cuda.synchronize(device=self.tensor_args.device)
# print("In: ", seed[0,:,0])
start_time = time.time()
with profiler.record_function("mpc/opt"):
act_seq = self.optimize(seed, shift_steps=shift_steps)
if self.sync_cuda_time:
- torch.cuda.synchronize()
+ torch.cuda.synchronize(device=self.tensor_args.device)
self.opt_dt = time.time() - start_time
with profiler.record_function("mpc/filter"):
act = self.safety_rollout.get_robot_command(
diff --git a/tests/interpolation_test.py b/tests/interpolation_test.py
index 11fe94d..5cbbf50 100644
--- a/tests/interpolation_test.py
+++ b/tests/interpolation_test.py
@@ -43,16 +43,21 @@ def test_linear_interpolation():
# create max_velocity buffer:
out_traj_gpu, _, _ = get_batch_interpolated_trajectory(
- in_traj, int_dt, max_vel, max_acc=max_acc, max_jerk=max_jerk, raw_dt=raw_dt
+ in_traj,
+ raw_dt,
+ int_dt,
+ max_vel,
+ max_acc=max_acc,
+ max_jerk=max_jerk,
)
#
out_traj_gpu = out_traj_gpu.clone()
out_traj_cpu, _, _ = get_batch_interpolated_trajectory(
in_traj,
+ raw_dt,
int_dt,
max_vel,
- raw_dt=raw_dt,
kind=InterpolateType.LINEAR,
max_acc=max_acc,
max_jerk=max_jerk,
diff --git a/tests/motion_gen_cuda_graph_test.py b/tests/motion_gen_cuda_graph_test.py
new file mode 100644
index 0000000..da0ed9e
--- /dev/null
+++ b/tests/motion_gen_cuda_graph_test.py
@@ -0,0 +1,330 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Third Party
+import pytest
+import torch
+
+# CuRobo
+from curobo.geom.types import WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import JointState, RobotConfig
+from curobo.util.trajectory import InterpolateType
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
+
+
+@pytest.fixture(scope="function")
+def motion_gen():
+ tensor_args = TensorDeviceType()
+ world_file = "collision_table.yml"
+ robot_file = "franka.yml"
+ motion_gen_config = MotionGenConfig.load_from_robot_config(
+ robot_file,
+ world_file,
+ tensor_args,
+ use_cuda_graph=False,
+ )
+ motion_gen_instance = MotionGen(motion_gen_config)
+ return motion_gen_instance
+
+
+@pytest.fixture(scope="function")
+def motion_gen_batch_env():
+ tensor_args = TensorDeviceType()
+ world_files = ["collision_table.yml", "collision_test.yml"]
+ world_cfg = [
+ WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
+ for world_file in world_files
+ ]
+ robot_file = "franka.yml"
+ motion_gen_config = MotionGenConfig.load_from_robot_config(
+ robot_file,
+ world_cfg,
+ tensor_args,
+ use_cuda_graph=False,
+ )
+ motion_gen_instance = MotionGen(motion_gen_config)
+
+ return motion_gen_instance
+
+
+@pytest.mark.parametrize(
+ "motion_gen_str,interpolation",
+ [
+ ("motion_gen", InterpolateType.LINEAR),
+ ("motion_gen", InterpolateType.CUBIC),
+ # ("motion_gen", InterpolateType.KUNZ_STILMAN_OPTIMAL),
+ ("motion_gen", InterpolateType.LINEAR_CUDA),
+ ],
+)
+def test_motion_gen_single(motion_gen_str, interpolation, request):
+ motion_gen = request.getfixturevalue(motion_gen_str)
+ motion_gen.update_interpolation_type(interpolation)
+ motion_gen.warmup()
+
+ retract_cfg = motion_gen.get_retract_config()
+
+ state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+ goal_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+ m_config = MotionGenPlanConfig(False, True)
+
+ result = motion_gen.plan_single(start_state, goal_pose, m_config)
+
+ # get final solutions:
+ assert torch.count_nonzero(result.success) == 1
+ reached_state = motion_gen.compute_kinematics(result.optimized_plan[-1])
+ assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
+
+
+def test_motion_gen_goalset(motion_gen):
+ motion_gen.warmup(n_goalset=2)
+
+ retract_cfg = motion_gen.get_retract_config()
+
+ state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+ goal_pose = Pose(
+ state.ee_pos_seq.repeat(2, 1).view(1, -1, 3),
+ quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
+ )
+ goal_pose.position[0, 0, 0] -= 0.1
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+ m_config = MotionGenPlanConfig(False, True)
+
+ result = motion_gen.plan_goalset(start_state, goal_pose, m_config)
+
+ # get final solutions:
+ assert torch.count_nonzero(result.success) == 1
+
+ reached_state = motion_gen.compute_kinematics(result.optimized_plan[-1])
+
+ assert (
+ torch.min(
+ torch.norm(goal_pose.position[:, 0, :] - reached_state.ee_pos_seq),
+ torch.norm(goal_pose.position[:, 1, :] - reached_state.ee_pos_seq),
+ )
+ < 0.005
+ )
+
+ assert result.goalset_index is not None
+
+ assert (
+ torch.norm(goal_pose.position[:, result.goalset_index, :] - reached_state.ee_pos_seq)
+ < 0.005
+ )
+
+
+def test_motion_gen_batch_goalset(motion_gen):
+ motion_gen.warmup(n_goalset=3, batch=3, warmup_js_trajopt=False, enable_graph=False)
+ retract_cfg = motion_gen.get_retract_config()
+
+ state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+ goal_pose = Pose(
+ state.ee_pos_seq.repeat(6, 1).view(3, -1, 3).clone(),
+ quaternion=state.ee_quat_seq.repeat(6, 1).view(3, -1, 4).clone(),
+ )
+ goal_pose.position[0, 1, 1] = 0.2
+ goal_pose.position[1, 0, 1] = 0.2
+ goal_pose.position[2, 1, 1] = 0.2
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
+
+ m_config = MotionGenPlanConfig(False, True, max_attempts=1, enable_graph_attempt=None)
+ result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config)
+
+ # get final solutions:
+ assert torch.count_nonzero(result.success) == result.success.shape[0]
+
+ reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1))
+
+ #
+ goal_position = torch.cat(
+ [
+ goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
+ for x in range(len(result.goalset_index))
+ ]
+ )
+
+ assert result.goalset_index is not None
+
+ assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
+
+
+def test_motion_gen_batch(motion_gen):
+ motion_gen.warmup(batch=2)
+ retract_cfg = motion_gen.get_retract_config()
+ state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+ goal_pose = Pose(
+ state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
+ ).repeat_seeds(2)
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2)
+
+ goal_pose.position[1, 0] -= 0.1
+
+ m_config = MotionGenPlanConfig(False, True)
+
+ result = motion_gen.plan_batch(start_state, goal_pose.clone(), m_config)
+ assert torch.count_nonzero(result.success) == 2
+
+ # get final solutions:
+ q = result.optimized_plan.trim_trajectory(-1).squeeze(1)
+ reached_state = motion_gen.compute_kinematics(q)
+ assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
+
+
+@pytest.mark.parametrize(
+ "motion_gen_str,interpolation",
+ [
+ ("motion_gen", InterpolateType.LINEAR),
+ ("motion_gen", InterpolateType.CUBIC),
+ # ("motion_gen", InterpolateType.KUNZ_STILMAN_OPTIMAL),
+ ("motion_gen", InterpolateType.LINEAR_CUDA),
+ ],
+)
+def test_motion_gen_batch_graph(motion_gen_str: str, interpolation: InterpolateType, request):
+ motion_gen = request.getfixturevalue(motion_gen_str)
+
+ motion_gen.graph_planner.interpolation_type = interpolation
+ motion_gen.reset()
+ retract_cfg = motion_gen.get_retract_config()
+ state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+ goal_pose = Pose(
+ state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
+ ).repeat_seeds(5)
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(5)
+
+ goal_pose.position[1, 0] -= 0.05
+
+ m_config = MotionGenPlanConfig(True, False)
+
+ result = motion_gen.plan_batch(start_state, goal_pose, m_config)
+ assert torch.count_nonzero(result.success) > 0
+
+ # get final solutions:
+ q = result.interpolated_plan.trim_trajectory(-1).squeeze(1)
+ reached_state = motion_gen.compute_kinematics(q)
+ assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
+
+
+def test_motion_gen_batch_env(motion_gen_batch_env):
+ motion_gen_batch_env.warmup(batch=2, batch_env_mode=True, enable_graph=False)
+
+ # motion_gen_batch_env.reset()
+ retract_cfg = motion_gen_batch_env.get_retract_config()
+ state = motion_gen_batch_env.compute_kinematics(
+ JointState.from_position(retract_cfg.view(1, -1))
+ )
+
+ goal_pose = Pose(
+ state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
+ ).repeat_seeds(2)
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2)
+
+ goal_pose.position[1, 0] -= 0.1
+
+ m_config = MotionGenPlanConfig(False, True, max_attempts=1)
+
+ result = motion_gen_batch_env.plan_batch_env(start_state, goal_pose, m_config)
+ assert torch.count_nonzero(result.success) == 2
+
+ # get final solutions:
+ reached_state = motion_gen_batch_env.compute_kinematics(
+ result.optimized_plan.trim_trajectory(-1).squeeze(1)
+ )
+ assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
+
+
+def test_motion_gen_batch_env_goalset(motion_gen_batch_env):
+ motion_gen_batch_env.warmup(batch=2, batch_env_mode=True, n_goalset=2, enable_graph=False)
+ retract_cfg = motion_gen_batch_env.get_retract_config()
+ state = motion_gen_batch_env.compute_kinematics(
+ JointState.from_position(retract_cfg.view(1, -1))
+ )
+
+ goal_pose = Pose(
+ state.ee_pos_seq.repeat(4, 1).view(2, -1, 3),
+ quaternion=state.ee_quat_seq.repeat(4, 1).view(2, -1, 4),
+ )
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2)
+
+ goal_pose.position[1, 0] -= 0.2
+
+ m_config = MotionGenPlanConfig(False, True, enable_graph_attempt=None)
+
+ result = motion_gen_batch_env.plan_batch_env_goalset(start_state, goal_pose, m_config)
+ assert torch.count_nonzero(result.success) > 0
+
+ # get final solutions:
+ reached_state = motion_gen_batch_env.compute_kinematics(
+ result.optimized_plan.trim_trajectory(-1).squeeze(1)
+ )
+ assert (
+ torch.min(
+ torch.norm(goal_pose.position[:, 0, :] - reached_state.ee_pos_seq),
+ torch.norm(goal_pose.position[:, 1, :] - reached_state.ee_pos_seq),
+ )
+ < 0.005
+ )
+
+ goal_position = torch.cat(
+ [
+ goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
+ for x in range(len(result.goalset_index))
+ ]
+ )
+
+ assert result.goalset_index is not None
+
+ assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
+
+
+@pytest.mark.parametrize(
+ "motion_gen_str,enable_graph",
+ [
+ ("motion_gen", True),
+ ("motion_gen", False),
+ ],
+)
+def test_motion_gen_single_js(motion_gen_str, enable_graph, request):
+ motion_gen = request.getfixturevalue(motion_gen_str)
+
+ motion_gen.warmup(warmup_js_trajopt=True)
+
+ retract_cfg = motion_gen.get_retract_config()
+
+ start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+ m_config = MotionGenPlanConfig(enable_graph=enable_graph, max_attempts=2)
+ goal_state = start_state.clone()
+ goal_state.position -= 0.3
+
+ result = motion_gen.plan_single_js(start_state, goal_state, m_config)
+
+ assert torch.count_nonzero(result.success) == 1
+
+ reached_state = result.optimized_plan[-1]
+
+ assert torch.norm(goal_state.position - reached_state.position) < 0.05
diff --git a/tests/motion_gen_speed_test.py b/tests/motion_gen_speed_test.py
index 9a208ef..773fff1 100644
--- a/tests/motion_gen_speed_test.py
+++ b/tests/motion_gen_speed_test.py
@@ -40,6 +40,24 @@ def motion_gen(request):
return motion_gen_instance
+@pytest.fixture(scope="module")
+def motion_gen_ur5e():
+ tensor_args = TensorDeviceType()
+ world_file = "collision_table.yml"
+ robot_file = "ur5e.yml"
+ motion_gen_config = MotionGenConfig.load_from_robot_config(
+ robot_file,
+ world_file,
+ tensor_args,
+ interpolation_steps=10000,
+ interpolation_dt=0.05,
+ )
+ motion_gen_instance = MotionGen(motion_gen_config)
+ motion_gen_instance.warmup(warmup_js_trajopt=False, enable_graph=False)
+
+ return motion_gen_instance
+
+
@pytest.mark.parametrize(
"motion_gen",
[
@@ -66,3 +84,126 @@ def test_motion_gen_velocity_scale(motion_gen):
result = motion_gen.plan_single(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) == 1
+
+
+@pytest.mark.parametrize(
+ "velocity_scale, acceleration_scale",
+ [
+ (1.0, 1.0),
+ (0.75, 1.0),
+ (0.5, 1.0),
+ (0.25, 1.0),
+ (0.15, 1.0),
+ (0.1, 1.0),
+ (1.0, 0.1),
+ (0.75, 0.1),
+ (0.5, 0.1),
+ (0.25, 0.1),
+ (0.15, 0.1),
+ (0.1, 0.1),
+ ],
+)
+def test_pose_sequence_speed_ur5e_scale(velocity_scale, acceleration_scale):
+ # load ur5e motion gen:
+
+ world_file = "collision_table.yml"
+ robot_file = "ur5e.yml"
+ motion_gen_config = MotionGenConfig.load_from_robot_config(
+ robot_file,
+ world_file,
+ interpolation_dt=(1.0 / 5.0),
+ velocity_scale=velocity_scale,
+ acceleration_scale=acceleration_scale,
+ )
+ motion_gen = MotionGen(motion_gen_config)
+ motion_gen.warmup(warmup_js_trajopt=False, enable_graph=False)
+ retract_cfg = motion_gen.get_retract_config()
+ start_state = JointState.from_position(retract_cfg.view(1, -1))
+
+ # poses for ur5e:
+ home_pose = [-0.431, 0.172, 0.348, 0, 1, 0, 0]
+ pose_1 = [0.157, -0.443, 0.427, 0, 1, 0, 0]
+ pose_2 = [0.126, -0.443, 0.729, 0, 0, 1, 0]
+ pose_3 = [-0.449, 0.339, 0.414, -0.681, -0.000, 0.000, 0.732]
+ pose_4 = [-0.449, 0.339, 0.414, 0.288, 0.651, -0.626, -0.320]
+ pose_5 = [-0.218, 0.508, 0.670, 0.529, 0.169, 0.254, 0.792]
+ pose_6 = [-0.865, 0.001, 0.411, 0.286, 0.648, -0.628, -0.321]
+
+ pose_list = [home_pose, pose_1, pose_2, pose_3, pose_4, pose_5, pose_6, home_pose]
+ trajectory = start_state
+ motion_time = 0
+ fail = 0
+ for i, pose in enumerate(pose_list):
+ goal_pose = Pose.from_list(pose, q_xyzw=False)
+ start_state = trajectory[-1].unsqueeze(0).clone()
+ start_state.velocity[:] = 0.0
+ start_state.acceleration[:] = 0.0
+ result = motion_gen.plan_single(
+ start_state.clone(),
+ goal_pose,
+ plan_config=MotionGenPlanConfig(
+ max_attempts=5,
+ ),
+ )
+ if result.success.item():
+ plan = result.get_interpolated_plan()
+ trajectory = trajectory.stack(plan.clone())
+ motion_time += result.motion_time
+ else:
+ fail += 1
+ assert fail == 0
+
+
+@pytest.mark.parametrize(
+ "motion_gen_str, time_dilation_factor",
+ [
+ ("motion_gen_ur5e", 1.0),
+ ("motion_gen_ur5e", 0.75),
+ ("motion_gen_ur5e", 0.5),
+ ("motion_gen_ur5e", 0.25),
+ ("motion_gen_ur5e", 0.15),
+ ("motion_gen_ur5e", 0.1),
+ ("motion_gen_ur5e", 0.001),
+ ],
+)
+def test_pose_sequence_speed_ur5e_time_dilation(motion_gen_str, time_dilation_factor, request):
+ # load ur5e motion gen:
+ motion_gen = request.getfixturevalue(motion_gen_str)
+
+ retract_cfg = motion_gen.get_retract_config()
+ start_state = JointState.from_position(retract_cfg.view(1, -1))
+
+ # poses for ur5e:
+ home_pose = [-0.431, 0.172, 0.348, 0, 1, 0, 0]
+ pose_1 = [0.157, -0.443, 0.427, 0, 1, 0, 0]
+ pose_2 = [0.126, -0.443, 0.729, 0, 0, 1, 0]
+ pose_3 = [-0.449, 0.339, 0.414, -0.681, -0.000, 0.000, 0.732]
+ pose_4 = [-0.449, 0.339, 0.414, 0.288, 0.651, -0.626, -0.320]
+ pose_5 = [-0.218, 0.508, 0.670, 0.529, 0.169, 0.254, 0.792]
+ pose_6 = [-0.865, 0.001, 0.411, 0.286, 0.648, -0.628, -0.321]
+
+ pose_list = [home_pose, pose_1, pose_2, pose_3, pose_4, pose_5, pose_6, home_pose]
+ trajectory = start_state
+ motion_time = 0
+ fail = 0
+ for i, pose in enumerate(pose_list):
+ goal_pose = Pose.from_list(pose, q_xyzw=False)
+ start_state = trajectory[-1].unsqueeze(0).clone()
+ start_state.velocity[:] = 0.0
+ start_state.acceleration[:] = 0.0
+ result = motion_gen.plan_single(
+ start_state.clone(),
+ goal_pose,
+ plan_config=MotionGenPlanConfig(
+ max_attempts=5,
+ time_dilation_factor=time_dilation_factor,
+ ),
+ )
+ if result.success.item():
+ plan = result.get_interpolated_plan()
+ trajectory = trajectory.stack(plan.clone())
+ motion_time += result.motion_time
+ else:
+ fail += 1
+ assert fail == 0
+ assert motion_time < 15 * (1 / time_dilation_factor)
diff --git a/tests/trajopt_test.py b/tests/trajopt_test.py
index 2f7f226..0666faf 100644
--- a/tests/trajopt_test.py
+++ b/tests/trajopt_test.py
@@ -57,15 +57,15 @@ def trajopt_solver_batch_env():
robot_cfg = RobotConfig.from_dict(
load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
)
- # world_cfg = WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
trajopt_config = TrajOptSolverConfig.load_from_robot_config(
robot_cfg,
world_cfg,
tensor_args,
use_cuda_graph=False,
- num_seeds=10,
+ num_seeds=4,
evaluate_interpolated_trajectory=True,
+ grad_trajopt_iters=200,
)
trajopt_solver = TrajOptSolver(trajopt_config)
@@ -73,7 +73,7 @@ def trajopt_solver_batch_env():
def test_trajopt_single_js(trajopt_solver):
- q_start = trajopt_solver.retract_config
+ q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.2
goal_state = JointState.from_position(q_goal)
current_state = JointState.from_position(q_start)
@@ -88,7 +88,7 @@ def test_trajopt_single_js(trajopt_solver):
def test_trajopt_single_pose(trajopt_solver):
trajopt_solver.reset_seed()
- q_start = trajopt_solver.retract_config
+ q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.1
kin_state = trajopt_solver.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
@@ -102,7 +102,7 @@ def test_trajopt_single_pose(trajopt_solver):
def test_trajopt_single_pose_no_seed(trajopt_solver):
trajopt_solver.reset_seed()
- q_start = trajopt_solver.retract_config
+ q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.05
kin_state = trajopt_solver.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
@@ -116,7 +116,7 @@ def test_trajopt_single_pose_no_seed(trajopt_solver):
def test_trajopt_single_goalset(trajopt_solver):
# run goalset planning:
- q_start = trajopt_solver.retract_config
+ q_start = trajopt_solver.retract_config.clone()
q_goal = q_start.clone() + 0.1
kin_state = trajopt_solver.fk(q_goal)
goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
@@ -133,7 +133,7 @@ def test_trajopt_single_goalset(trajopt_solver):
def test_trajopt_batch(trajopt_solver):
# run goalset planning:
- q_start = trajopt_solver.retract_config.repeat(2, 1)
+ q_start = trajopt_solver.retract_config.clone().repeat(2, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -153,7 +153,7 @@ def test_trajopt_batch(trajopt_solver):
def test_trajopt_batch_js(trajopt_solver):
# run goalset planning:
- q_start = trajopt_solver.retract_config.repeat(2, 1)
+ q_start = trajopt_solver.retract_config.clone().repeat(2, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -173,7 +173,7 @@ def test_trajopt_batch_js(trajopt_solver):
def test_trajopt_batch_goalset(trajopt_solver):
# run goalset planning:
- q_start = trajopt_solver.retract_config.repeat(3, 1)
+ q_start = trajopt_solver.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -196,14 +196,12 @@ def test_trajopt_batch_goalset(trajopt_solver):
def test_trajopt_batch_env_js(trajopt_solver_batch_env):
# run goalset planning:
- q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
+ q_start = trajopt_solver_batch_env.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal += 0.1
- q_goal[2] += 0.1
+ q_goal[2][0] += 0.1
q_goal[1] -= 0.2
# q_goal[2, -1] += 0.1
- kin_state = trajopt_solver_batch_env.fk(q_goal)
- goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
goal_state = JointState.from_position(q_goal)
current_state = JointState.from_position(q_start)
@@ -213,14 +211,15 @@ def test_trajopt_batch_env_js(trajopt_solver_batch_env):
traj = result.solution.position
interpolated_traj = result.interpolated_solution.position
assert torch.count_nonzero(result.success) == 3
- assert torch.linalg.norm((goal_state.position - traj[:, -1, :])).item() < 0.005
- assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.005
+ error = torch.linalg.norm((goal_state.position - traj[:, -1, :]), dim=-1)
+ assert torch.max(error).item() < 0.05
+ assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.05
assert len(result) == 3
def test_trajopt_batch_env(trajopt_solver_batch_env):
# run goalset planning:
- q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
+ q_start = trajopt_solver_batch_env.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1
@@ -262,7 +261,7 @@ def test_trajopt_batch_env_goalset(trajopt_solver_batch_env):
def test_trajopt_batch_env(trajopt_solver):
# run goalset planning:
- q_start = trajopt_solver.retract_config.repeat(3, 1)
+ q_start = trajopt_solver.retract_config.clone().repeat(3, 1)
q_goal = q_start.clone()
q_goal[0] += 0.1
q_goal[1] -= 0.1