release repository

tests/__init__.py

@@ -0,0 +1,12 @@
+#
+# 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.
+#
+
+"""CuRobo test package."""

tests/cost_test.py

@@ -0,0 +1,47 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.geom.sdf.world import WorldCollisionConfig, WorldPrimitiveCollision
+from curobo.geom.types import WorldConfig
+from curobo.rollout.cost.primitive_collision_cost import (
+    PrimitiveCollisionCost,
+    PrimitiveCollisionCostConfig,
+)
+from curobo.types.base import TensorDeviceType
+from curobo.util_file import get_world_configs_path, join_path, load_yaml
+
+
+def test_primitive_collision_cost():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldPrimitiveCollision(
+        WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    )
+
+    cost_cfg = PrimitiveCollisionCostConfig(
+        weight=1.0,
+        tensor_args=tensor_args,
+        world_coll_checker=coll_cfg,
+        use_sweep=False,
+        classify=False,
+    )
+    cost = PrimitiveCollisionCost(cost_cfg)
+    q_spheres = torch.as_tensor(
+        [[0.1, 0.0, 0.0, 0.2], [10.0, 0.0, 0.0, 0.1]], **vars(tensor_args)
+    ).view(-1, 1, 1, 4)
+    c = cost.forward(q_spheres).flatten()
+    assert c[0] > 0.0 and c[1] == 0.0

tests/cuda_robot_generator_test.py

@@ -0,0 +1,55 @@
+#
+# 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.
+#
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_generator import (
+    CudaRobotGenerator,
+    CudaRobotGeneratorConfig,
+)
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModelConfig
+from curobo.util_file import get_robot_configs_path, join_path, load_yaml
+
+
+def test_cuda_robot_generator_config():
+    robot_file = "franka.yml"
+    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"][
+        "kinematics"
+    ]
+    config = CudaRobotGeneratorConfig(**robot_params)
+    assert config.ee_link == robot_params["ee_link"]
+
+
+def test_cuda_robot_generator():
+    robot_file = "franka.yml"
+    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"][
+        "kinematics"
+    ]
+    config = CudaRobotGeneratorConfig(**robot_params)
+    robot_generator = CudaRobotGenerator(config)
+    assert robot_generator.kinematics_config.n_dof == 7
+
+
+def test_cuda_robot_config():
+    robot_file = "franka.yml"
+    config = CudaRobotModelConfig.from_robot_yaml_file(robot_file)
+    assert config.kinematics_config.n_dof == 7
+
+
+def test_cuda_robot_generator_config_cspace():
+    robot_file = "franka.yml"
+    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"][
+        "kinematics"
+    ]
+    config = CudaRobotGeneratorConfig(**robot_params)
+    assert len(config.cspace.max_jerk) == len(config.cspace.joint_names)
+    robot_generator = CudaRobotGenerator(config)
+
+    assert len(robot_generator.cspace.max_jerk) == 7

tests/curobo_robot_world_model_test.py

@@ -0,0 +1,136 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.wrap.model.curobo_robot_world import CuroboRobotWorld, CuroboRobotWorldConfig
+
+
+def load_robot_world():
+    robot_file = "franka.yml"
+    world_file = "collision_table.yml"
+    config = CuroboRobotWorldConfig.load_from_config(robot_file, world_file)
+    model = CuroboRobotWorld(config)
+    return model
+
+
+def load_robot_batch_world():
+    robot_file = "franka.yml"
+    world_file = ["collision_table.yml", "collision_test.yml"]
+    config = CuroboRobotWorldConfig.load_from_config(
+        robot_file, world_file, collision_activation_distance=0.0
+    )
+    model = CuroboRobotWorld(config)
+    return model
+
+
+def test_cu_robot_world_config_load():
+    # test if curobo robot world can be loaded
+    load_robot_world()
+    assert True
+
+
+def test_cu_robot_batch_world_config_load():
+    # test if curobo robot world can be loaded
+    load_robot_batch_world()
+    assert True
+
+
+def test_cu_robot_world_kinematics():
+    # test if kinematics can be queried
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    assert state.ee_position.shape[-1] == 3
+
+
+def test_cu_robot_world_sample():
+    # test if joint configurations can be sampled
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    assert q.shape[0] == n
+
+
+def test_cu_robot_world_collision():
+    # test computing collisions given robot joint configurations
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    d_world = model.get_collision_constraint(state.link_spheres_tensor.view(n, 1, -1, 4))
+    assert d_world.shape[0] == n
+    assert torch.sum(d_world) == 0.0
+
+
+def test_cu_robot_world_self_collision():
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    d_self = model.get_self_collision(state.link_spheres_tensor.view(n, 1, -1, 4))
+    assert torch.sum(d_self) == 0.0
+
+
+def test_cu_robot_world_pose():
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    pose = state.ee_pose
+    des_pose = state.ee_pose
+    d = model.pose_distance(des_pose, pose)
+    assert torch.sum(d) < 1e-3
+
+
+def test_cu_trajectory_sample():
+    model = load_robot_world()
+    b = 10
+    horizon = 21
+    q = model.sample_trajectory(b, horizon)
+    assert q.shape[0] == b and q.shape[1] == horizon
+
+
+def test_cu_batch_trajectory_sample():
+    model = load_robot_batch_world()
+    b = 2
+    horizon = 21
+    env_query_idx = torch.zeros((b), dtype=torch.int32, device=torch.device("cuda", 0))
+    q = model.sample_trajectory(b, horizon, env_query_idx=env_query_idx)
+    assert q.shape[0] == b and q.shape[1] == horizon
+
+
+def test_cu_batch_trajectory_1env_sample():
+    model = load_robot_batch_world()
+    b = 2
+    horizon = 21
+    env_query_idx = None
+    q = model.sample_trajectory(b, horizon, env_query_idx=env_query_idx)
+    assert q.shape[0] == b and q.shape[1] == horizon
+
+
+def test_cu_robot_batch_world_collision():
+    # test computing collisions given robot joint configurations
+    model = load_robot_batch_world()
+    b = 2
+    horizon = 21
+    env_query_idx = torch.zeros((b), dtype=torch.int32, device=torch.device("cuda", 0))
+
+    q = model.sample_trajectory(b, horizon, env_query_idx=env_query_idx)
+
+    d_world, d_self = model.get_world_self_collision_distance_from_joint_trajectory(
+        q, env_query_idx
+    )
+    assert d_world.shape[0] == b
+    assert torch.sum(d_world) == 0.0

tests/curobo_version_test.py

@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+"""Unit tests for `curobo` package version."""
+
+# CuRobo
+import curobo
+
+
+def test_curobo_version():
+    """Test `curobo` package version is set."""
+    assert curobo.__version__ is not None
+    assert curobo.__version__ != ""

tests/geom_test.py

@@ -0,0 +1,211 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.geom.sdf.world import (
+    CollisionQueryBuffer,
+    WorldCollisionConfig,
+    WorldPrimitiveCollision,
+)
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import Cuboid, WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util_file import get_world_configs_path, join_path, load_yaml
+
+
+def test_world_primitive():
+    # load a world:
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldPrimitiveCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+
+    d_sph = coll_check.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+    assert abs(d_sph[0].item() - 0.1) < 1e-3
+    assert abs(d_sph[1].item() - 0.0) < 1e-9
+    assert abs(d_sph[2].item() - 0.1) < 1e-3
+
+
+def test_batch_world_primitive():
+    """This tests collision checking across different environments"""
+    tensor_args = TensorDeviceType()
+
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldPrimitiveCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(-1, 1, 1, 4)
+    env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    act_distance = tensor_args.to_device([0.0])
+
+    weight = tensor_args.to_device([1])
+    d_sph = coll_check.get_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, env_query_idx
+    ).view(-1)
+    assert abs(d_sph[0].item() - 0.1) < 1e-3
+    assert abs(d_sph[1].item() - 0.0) < 1e-9
+    assert abs(d_sph[2].item() - 0.1) < 1e-3
+
+
+def test_swept_world_primitive():
+    """This tests collision checking across different environments"""
+    tensor_args = TensorDeviceType()
+
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_cfg.cache = {"obb": 5}
+    coll_check = WorldPrimitiveCollision(coll_cfg)
+    # add an obstacle:
+    new_cube = Cuboid("cube_1", [0, 0, 1, 1, 0, 0, 0], None, [0.1, 0.2, 0.2])
+    coll_check.add_obb(new_cube, 0)
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    env_query_idx = None
+    env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+    dt = act_distance.clone()
+    d_sph_swept = coll_check.get_swept_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, dt, 10, env_query_idx
+    ).view(-1)
+    d_sph = coll_check.get_sphere_distance(
+        x_sph, query_buffer.clone() * 0.0, weight, act_distance, env_query_idx
+    ).view(-1)
+    assert abs(d_sph[0].item() - 0.1) < 1e-3
+    assert abs(d_sph[1].item() - 0.0) < 1e-9
+    assert abs(d_sph[2].item() - 0.1) < 1e-3
+
+    assert abs(d_sph_swept[0].item() - 0.1) < 1e-3
+    assert abs(d_sph_swept[1].item() - 0.0) < 1e-9
+    assert abs(d_sph_swept[2].item() - 0.1) < 1e-3
+
+
+def test_world_primitive_mesh_instance():
+    # load a world:
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldMeshCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    act_distance = tensor_args.to_device([0.0])
+
+    weight = tensor_args.to_device([1])
+    d_sph = coll_check.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+    assert abs(d_sph[0].item() - 0.1) < 1e-3
+    assert abs(d_sph[1].item() - 0.0) < 1e-9
+    assert abs(d_sph[2].item() - 0.1) < 1e-3
+
+
+def test_batch_world_primitive_mesh_instance():
+    """This tests collision checking across different environments"""
+    tensor_args = TensorDeviceType()
+
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldMeshCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(-1, 1, 1, 4)
+    env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+
+    d_sph = coll_check.get_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, env_query_idx
+    ).view(-1)
+    assert abs(d_sph[0].item() - 0.1) < 1e-3
+    assert abs(d_sph[1].item() - 0.0) < 1e-9
+    assert abs(d_sph[2].item() - 0.1) < 1e-3
+
+
+def test_swept_world_primitive_mesh_instance():
+    """This tests collision checking across different environments"""
+    tensor_args = TensorDeviceType()
+
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_cfg.cache = {"obb": 5}
+    coll_check = WorldMeshCollision(coll_cfg)
+    # add an obstacle:
+    dims = tensor_args.to_device([0.1, 0.2, 0.2])
+    w_obj_pose = Pose.from_list([0, 0, 1, 1, 0, 0, 0], tensor_args)
+    coll_check.add_obb_from_raw("cube_1", dims, 0, w_obj_pose=w_obj_pose)
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    env_query_idx = None
+    env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+    dt = act_distance.clone() + 0.01
+
+    d_sph_swept = coll_check.get_swept_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, dt, 10, env_query_idx
+    ).view(-1)
+    d_sph = coll_check.get_sphere_distance(
+        x_sph, query_buffer.clone() * 0.0, weight, act_distance, env_query_idx
+    ).view(-1)
+    assert abs(d_sph[0].item() - 0.1) < 1e-3
+    assert abs(d_sph[1].item() - 0.0) < 1e-9
+    assert abs(d_sph[2].item() - 0.1) < 1e-3
+
+    assert abs(d_sph_swept[0].item() - 0.1) < 1e-3
+    assert abs(d_sph_swept[1].item() - 0.0) < 1e-9
+    assert abs(d_sph_swept[2].item() - 0.1) < 1e-3

tests/geom_types_test.py

@@ -0,0 +1,41 @@
+#
+# 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
+
+# CuRobo
+from curobo.geom.sphere_fit import SphereFitType
+from curobo.geom.types import WorldConfig
+from curobo.util_file import get_world_configs_path, join_path, load_yaml
+
+
+def test_mesh_world():
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    mesh_world_cfg = world_cfg.get_mesh_world()
+    assert len(mesh_world_cfg.mesh) == len(world_cfg.cuboid)
+    obb_world_cfg = mesh_world_cfg.get_obb_world()
+    assert len(obb_world_cfg.cuboid) == len(mesh_world_cfg.mesh)
+
+
+@pytest.mark.parametrize(
+    "sphere_fit_type",
+    [e.value for e in SphereFitType],
+)
+def test_bounding_volume(sphere_fit_type):
+    world_file = "collision_test.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    obs = world_cfg.objects[-1]
+    spheres = obs.get_bounding_spheres(100, 0.01, sphere_fit_type)
+    assert len(spheres) > 0

tests/goal_test.py

@@ -0,0 +1,144 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.curobolib.geom import get_pose_distance
+from curobo.rollout.rollout_base import Goal
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import JointState
+
+
+def test_repeat_seeds():
+    tensor_args = TensorDeviceType()
+    b = 10
+    dof = 7
+    position = torch.randn((b, 3), device=tensor_args.device, dtype=tensor_args.dtype)
+
+    quaternion = torch.zeros((b, 4), device=tensor_args.device, dtype=tensor_args.dtype)
+
+    quaternion[:, 0] = 1.0
+    goal_pose = Pose(position, quaternion)
+
+    current_state = JointState.from_position(
+        torch.randn((b, dof), device=tensor_args.device, dtype=tensor_args.dtype)
+    )
+    batch_pose_idx = torch.arange(0, b, 1, device=tensor_args.device, dtype=torch.int32).unsqueeze(
+        -1
+    )
+    goal = Goal(goal_pose=goal_pose, batch_pose_idx=batch_pose_idx, current_state=current_state)
+    g = goal.repeat_seeds(4)
+
+    start_pose = goal_pose.repeat_seeds(4)
+    b = start_pose.position.shape[0]
+    out_d = torch.zeros((b, 1), device=tensor_args.device, dtype=tensor_args.dtype)
+    out_p_v = torch.zeros((b, 3), device=tensor_args.device, dtype=tensor_args.dtype)
+    out_r_v = torch.zeros((b, 4), device=tensor_args.device, dtype=tensor_args.dtype)
+    out_idx = torch.zeros((b, 1), device=tensor_args.device, dtype=torch.int32)
+    vec_weight = torch.ones((6), device=tensor_args.device, dtype=tensor_args.dtype)
+    weight = tensor_args.to_device([1, 1, 1, 1])
+    vec_convergence = tensor_args.to_device([0, 0])
+    run_weight = tensor_args.to_device([1])
+    r = get_pose_distance(
+        out_d,
+        out_d.clone(),
+        out_d.clone(),
+        out_p_v,
+        out_r_v,
+        out_idx,
+        start_pose.position,
+        g.goal_pose.position,
+        start_pose.quaternion,
+        g.goal_pose.quaternion,
+        vec_weight,
+        weight,
+        vec_convergence,
+        run_weight,
+        vec_weight.clone() * 0.0,
+        g.batch_pose_idx,
+        start_pose.position.shape[0],
+        1,
+        1,
+        1,
+        False,
+        False,
+        True,
+    )
+
+    assert torch.sum(r[0]).item() == 0.0
+
+
+def test_horizon_repeat_seeds():
+    tensor_args = TensorDeviceType()
+    b = 2
+    dof = 7
+    h = 30
+    position = torch.randn((b, h, 3), device=tensor_args.device, dtype=tensor_args.dtype)
+    # position[:,:,1] = 1.0
+    quaternion = torch.zeros((b, h, 4), device=tensor_args.device, dtype=tensor_args.dtype)
+
+    quaternion[:, 0] = 1.0
+    quaternion[1, 1] = 1.0
+    quaternion[1, 0] = 0.0
+    goal_pose = Pose(position[:, 0, :], quaternion[:, 0, :]).clone()
+
+    current_state = JointState.from_position(
+        torch.randn((b, dof), device=tensor_args.device, dtype=tensor_args.dtype)
+    )
+    batch_pose_idx = torch.arange(0, b, 1, device=tensor_args.device, dtype=torch.int32).unsqueeze(
+        -1
+    )
+    goal = Goal(goal_pose=goal_pose, batch_pose_idx=batch_pose_idx, current_state=current_state)
+    g = goal  # .repeat_seeds(4)
+
+    start_pose = Pose(
+        goal_pose.position.view(-1, 1, 3).repeat(1, h, 1),
+        quaternion=goal_pose.quaternion.view(-1, 1, 4).repeat(1, h, 1),
+    )
+    b = start_pose.position.shape[0]
+    out_d = torch.zeros((b, h, 1), device=tensor_args.device, dtype=tensor_args.dtype)
+    out_p_v = torch.zeros((b, h, 3), device=tensor_args.device, dtype=tensor_args.dtype)
+    out_r_v = torch.zeros((b, h, 4), device=tensor_args.device, dtype=tensor_args.dtype)
+    out_idx = torch.zeros((b, h, 1), device=tensor_args.device, dtype=torch.int32)
+    vec_weight = torch.ones((6), device=tensor_args.device, dtype=tensor_args.dtype)
+    weight = tensor_args.to_device([1, 1, 1, 1])
+    vec_convergence = tensor_args.to_device([0, 0])
+    run_weight = torch.zeros((1, h), device=tensor_args.device)
+    run_weight[-1] = 1
+    r = get_pose_distance(
+        out_d,
+        out_d.clone(),
+        out_d.clone(),
+        out_p_v,
+        out_r_v,
+        out_idx,
+        start_pose.position,
+        g.goal_pose.position,
+        start_pose.quaternion,
+        g.goal_pose.quaternion,
+        vec_weight,
+        weight,
+        vec_convergence,
+        run_weight,
+        vec_weight.clone() * 0.0,
+        g.batch_pose_idx,
+        start_pose.position.shape[0],
+        h,
+        1,
+        1,
+        True,
+        False,
+        False,
+    )
+    assert torch.sum(r[0]).item() == 0.0

tests/ik_config_test.py

@@ -0,0 +1,149 @@
+#
+# 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.sdf.world import (
+    CollisionCheckerType,
+    WorldCollisionConfig,
+    WorldPrimitiveCollision,
+)
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
+
+
+def ik_base_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=False,
+        n_collision_envs=1,
+        collision_cache={"obb": 10},
+    )
+    return ik_config
+
+
+def ik_gd_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=False,
+        use_gradient_descent=True,
+        grad_iters=100,
+    )
+    return ik_config
+
+
+def ik_es_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=False,
+        use_es=True,
+        es_learning_rate=0.01,
+        use_fixed_samples=True,
+    )
+    return ik_config
+
+
+def ik_no_particle_opt_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=False,
+        use_particle_opt=False,
+        grad_iters=100,
+    )
+    return ik_config
+
+
+@pytest.mark.parametrize(
+    "config, expected",
+    [
+        (ik_base_config(), True),
+        (ik_es_config(), True),
+        (ik_gd_config(), True),
+        (ik_no_particle_opt_config(), True),
+    ],
+)
+def test_eval(config, expected):
+    ik_solver = IKSolver(config)
+    q_sample = ik_solver.sample_configs(1)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_single(goal)
+    result = ik_solver.solve_single(goal)
+
+    success = result.success
+    assert success.item() == expected

tests/ik_module_test.py

@@ -0,0 +1,131 @@
+#
+# 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.sdf.world import (
+    CollisionCheckerType,
+    WorldCollisionConfig,
+    WorldPrimitiveCollision,
+)
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
+
+
+@pytest.fixture(scope="module")
+def ik_solver():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=False,
+    )
+    ik_solver = IKSolver(ik_config)
+    return ik_solver
+
+
+@pytest.fixture(scope="module")
+def ik_solver_batch_env():
+    tensor_args = TensorDeviceType()
+    world_files = ["collision_table.yml", "collision_cubby.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"
+    robot_cfg = RobotConfig.from_dict(
+        load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    )
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=False,
+    )
+    ik_solver = IKSolver(ik_config)
+    return ik_solver
+
+
+def test_ik_single(ik_solver):
+    q_sample = ik_solver.sample_configs(1)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_single(goal)
+
+    success = result.success
+    assert success.item()
+
+
+def test_ik_goalset(ik_solver):
+    q_sample = ik_solver.sample_configs(10)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position.unsqueeze(0), kin_state.ee_quaternion.unsqueeze(0))
+    result = ik_solver.solve_goalset(goal)
+
+    assert result.success.item()
+
+
+def test_ik_batch(ik_solver):
+    q_sample = ik_solver.sample_configs(10)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_batch(goal)
+    assert torch.count_nonzero(result.success) > 5
+
+
+def test_ik_batch_goalset(ik_solver):
+    q_sample = ik_solver.sample_configs(100)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position.view(10, 10, 3), kin_state.ee_quaternion.view(10, 10, 4))
+    result = ik_solver.solve_batch_goalset(goal)
+    assert torch.count_nonzero(result.success) > 5
+
+
+def test_ik_batch_env(ik_solver_batch_env):
+    q_sample = ik_solver_batch_env.sample_configs(3)
+    kin_state = ik_solver_batch_env.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver_batch_env.solve_batch_env(goal)
+
+    assert torch.count_nonzero(result.success) >= 1
+
+
+def test_ik_batch_env_goalset(ik_solver_batch_env):
+    q_sample = ik_solver_batch_env.sample_configs(3 * 3)
+    kin_state = ik_solver_batch_env.fk(q_sample)
+    goal = Pose(kin_state.ee_position.view(3, 3, 3), kin_state.ee_quaternion.view(3, 3, 4))
+    result = ik_solver_batch_env.solve_batch_env_goalset(goal)
+    assert torch.count_nonzero(result.success) >= 2

tests/ik_test.py

@@ -0,0 +1,211 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.geom.sdf.world import (
+    CollisionCheckerType,
+    WorldCollisionConfig,
+    WorldPrimitiveCollision,
+)
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
+
+
+def test_basic_ik():
+    tensor_args = TensorDeviceType()
+
+    config_file = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))
+    urdf_file = config_file["robot_cfg"]["kinematics"][
+        "urdf_path"
+    ]  # Send global path starting with "/"
+    base_link = config_file["robot_cfg"]["kinematics"]["base_link"]
+    ee_link = config_file["robot_cfg"]["kinematics"]["ee_link"]
+    robot_cfg = RobotConfig.from_basic(urdf_file, base_link, ee_link, tensor_args)
+
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        None,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=30,
+        self_collision_check=False,
+        self_collision_opt=False,
+        tensor_args=tensor_args,
+    )
+    ik_solver = IKSolver(ik_config)
+    b_size = 10
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_batch(goal)
+    success = result.success
+
+    assert torch.count_nonzero(success).item() >= 1.0  # we check if atleast 1 is successful
+
+
+def test_full_config_collision_free_ik():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=30,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+    )
+    ik_solver = IKSolver(ik_config)
+    b_size = 10
+
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve(goal)
+
+    success = result.success
+    assert torch.count_nonzero(success).item() >= 9.0  # we check if atleast 90% are successful
+
+
+def test_attach_object_full_config_collision_free_ik():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    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)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=30,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+    )
+    ik_solver = IKSolver(ik_config)
+    b_size = 10
+
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve(goal)
+
+    success = result.success
+    assert torch.count_nonzero(success).item() >= 9.0  # we check if atleast 90% are successful
+
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+
+    # ik_solver.attach_object_to_robot(sphere_radius=0.02)
+
+    result = ik_solver.solve(goal)
+    success = result.success
+    assert torch.count_nonzero(success).item() >= 9.0  # we check if atleast 90% are successful
+
+
+def test_batch_env_ik():
+    tensor_args = TensorDeviceType()
+    world_files = ["collision_cubby.yml", "collision_test.yml"]
+
+    robot_file = "franka.yml"
+    robot_cfg = RobotConfig.from_dict(
+        load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    )
+    w_list = [
+        WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
+        for world_file in world_files
+    ]
+    world_ccheck = WorldPrimitiveCollision(WorldCollisionConfig(tensor_args, n_envs=2))
+    # create a batched world collision checker:
+    world_ccheck.load_batch_collision_model(w_list)
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_coll_checker=world_ccheck,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=True,
+    )
+    ik_solver = IKSolver(ik_config)
+    b_size = 2
+
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_batch_env(goal)
+
+    success = result.success
+    assert torch.count_nonzero(success).item() >= 1.0  # we check if atleast 90% are successful
+
+
+def test_batch_env_mesh_ik():
+    tensor_args = TensorDeviceType()
+    world_files = ["collision_table.yml", "collision_table.yml"]
+
+    robot_file = "franka.yml"
+    robot_cfg = RobotConfig.from_dict(
+        load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    )
+    w_list = [
+        WorldConfig.from_dict(
+            load_yaml(join_path(get_world_configs_path(), world_file))
+        ).get_mesh_world()
+        for world_file in world_files
+    ]
+    world_ccheck = WorldMeshCollision(
+        WorldCollisionConfig(tensor_args, checker_type=CollisionCheckerType.MESH, n_envs=2)
+    )
+    # create a batched world collision checker:
+    # print(w_list)
+    world_ccheck.load_batch_collision_model(w_list)
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_coll_checker=world_ccheck,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=100,
+        self_collision_check=True,
+        self_collision_opt=True,
+        tensor_args=tensor_args,
+        use_cuda_graph=True,
+    )
+    ik_solver = IKSolver(ik_config)
+    b_size = 2
+
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_batch_env(goal)
+
+    success = result.success
+    assert torch.count_nonzero(success).item() >= 1.0  # we check if atleast 90% are successful

tests/interpolation_test.py

@@ -0,0 +1,67 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.types.base import TensorDeviceType
+from curobo.types.robot import JointState
+from curobo.util.trajectory import InterpolateType, get_batch_interpolated_trajectory
+
+
+def test_linear_interpolation():
+    b, h, dof = 1, 24, 1
+    raw_dt = 0.4
+    int_dt = 0.01
+    tensor_args = TensorDeviceType()
+    # initialize raw trajectory:
+    in_traj = JointState.zeros((b, h, dof), tensor_args)
+    in_traj.position = torch.zeros((b, h, dof), device=tensor_args.device)
+    in_traj.position[:, 1, :] = 0.1
+
+    in_traj.position[:, -2, :] = -0.01
+    in_traj.position[:, 10, :] = -0.01
+
+    in_traj.position[:, -1, :] = 0.01
+    in_traj.velocity = in_traj.position - torch.roll(in_traj.position, -1, dims=1)
+    in_traj.velocity[:, 0, :] = 0.0
+    in_traj.velocity[:, -1, :] = 0.0
+
+    max_vel = torch.ones((1, 1, dof), device=tensor_args.device, dtype=tensor_args.dtype)
+    max_acc = torch.ones((1, 1, dof), device=tensor_args.device, dtype=tensor_args.dtype) * 25
+    max_jerk = torch.ones((1, 1, dof), device=tensor_args.device, dtype=tensor_args.dtype) * 500
+
+    # 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
+    )
+    #
+    out_traj_gpu = out_traj_gpu.clone()
+
+    out_traj_cpu, _, _ = get_batch_interpolated_trajectory(
+        in_traj,
+        int_dt,
+        max_vel,
+        raw_dt=raw_dt,
+        kind=InterpolateType.LINEAR,
+        max_acc=max_acc,
+        max_jerk=max_jerk,
+    )
+    assert (
+        torch.max(
+            torch.abs(out_traj_gpu.position[:, -5:, :] - out_traj_cpu.position[:, -5:, :])
+        ).item()
+        < 0.05
+    )
+
+
+# test_linear_interpolation()

tests/kinematics_parser_test.py

@@ -0,0 +1,142 @@
+#
+# 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.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
+from curobo.cuda_robot_model.urdf_kinematics_parser import UrdfKinematicsParser
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util_file import (
+    file_exists,
+    get_assets_path,
+    get_robot_configs_path,
+    join_path,
+    load_yaml,
+)
+
+try:
+    # CuRobo
+    from curobo.cuda_robot_model.usd_kinematics_parser import UsdKinematicsParser
+except ImportError:
+    pytest.skip("usd-core is not available, skipping USD tests", allow_module_level=True)
+
+
+def check_usd_file_exists():
+    robot_file = "franka.yml"
+    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_usd = join_path(get_assets_path(), robot_params["kinematics"]["usd_path"])
+    return file_exists(robot_usd)
+
+
+if not check_usd_file_exists():
+    pytest.skip("Franka Panda USD is not available, skipping USD tests", allow_module_level=True)
+
+
+@pytest.fixture(scope="module")
+def robot_params_all():
+    robot_file = "franka.yml"
+    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_params["kinematics"]["ee_link"] = "panda_link7"
+    robot_params["kinematics"]["base_link"] = "panda_link0"
+    robot_params["kinematics"]["collision_link_names"] = []
+    robot_params["kinematics"]["lock_joints"] = None
+    robot_params["kinematics"]["extra_links"] = None
+
+    return robot_params
+
+
+@pytest.fixture(scope="module")
+def usd_parser(robot_params_all):
+    robot_params = robot_params_all["kinematics"]
+    robot_usd = join_path(get_assets_path(), robot_params["usd_path"])
+    kinematics_parser = UsdKinematicsParser(
+        usd_path=robot_usd,
+        flip_joints=robot_params["usd_flip_joints"],
+        usd_robot_root=robot_params["usd_robot_root"],
+    )
+    return kinematics_parser
+
+
+@pytest.fixture(scope="module")
+def urdf_parser(robot_params_all):
+    robot_params = robot_params_all["kinematics"]
+    robot_urdf = join_path(get_assets_path(), robot_params["urdf_path"])
+    kinematics_parser = UrdfKinematicsParser(robot_urdf)
+    return kinematics_parser
+
+
+@pytest.fixture(scope="module")
+def retract_state(robot_params_all):
+    tensor_args = TensorDeviceType()
+    q = tensor_args.to_device(robot_params_all["kinematics"]["cspace"]["retract_config"]).view(
+        1, -1
+    )
+    return q
+
+
+@pytest.fixture(scope="module")
+def usd_cuda_robot(robot_params_all):
+    robot_params = robot_params_all["kinematics"]
+    robot_params["use_usd_kinematics"] = True
+
+    usd_cuda_config = CudaRobotModelConfig.from_data_dict(robot_params)
+
+    usd_cuda_robot = CudaRobotModel(usd_cuda_config)
+    return usd_cuda_robot
+
+
+@pytest.fixture(scope="module")
+def urdf_cuda_robot(robot_params_all):
+    robot_params = robot_params_all["kinematics"]
+    robot_params["use_usd_kinematics"] = False
+    urdf_cuda_config = CudaRobotModelConfig.from_data_dict(robot_params)
+
+    urdf_cuda_robot = CudaRobotModel(urdf_cuda_config)
+    return urdf_cuda_robot
+
+
+def test_chain_parse(urdf_parser, usd_parser, robot_params_all):
+    robot_params = robot_params_all
+
+    urdf_chain = urdf_parser.get_chain(
+        robot_params["kinematics"]["base_link"], robot_params["kinematics"]["ee_link"]
+    )
+
+    usd_chain = usd_parser.get_chain(
+        robot_params["kinematics"]["base_link"], robot_params["kinematics"]["ee_link"]
+    )
+    assert usd_chain == urdf_chain
+
+
+def test_joint_transform_parse(usd_cuda_robot, urdf_cuda_robot):
+    usd_pose = usd_cuda_robot.get_all_link_transforms()
+    urdf_pose = urdf_cuda_robot.get_all_link_transforms()
+    p, q = usd_pose.distance(urdf_pose)
+    p = torch.linalg.norm(q)
+    q = torch.linalg.norm(q)
+    assert p < 1e-5 and q < 1e-5
+
+
+def test_basic_ee_pose(usd_cuda_robot, urdf_cuda_robot, retract_state):
+    q = retract_state[:, : usd_cuda_robot.get_dof()]
+
+    usd_state = usd_cuda_robot.get_state(q)
+    usd_pose = Pose(usd_state.ee_position, usd_state.ee_quaternion)
+
+    urdf_state = urdf_cuda_robot.get_state(q)
+    urdf_pose = Pose(urdf_state.ee_position, urdf_state.ee_quaternion)
+    p_d, q_d = usd_pose.distance(urdf_pose)
+    assert p_d < 1e-5
+    assert q_d < 1e-5

tests/kinematics_test.py

@@ -0,0 +1,171 @@
+#
+# 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.cuda_robot_model.cuda_robot_generator import CudaRobotGeneratorConfig
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
+from curobo.cuda_robot_model.types import CSpaceConfig
+from curobo.geom.transform import matrix_to_quaternion, quaternion_to_matrix
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import quat_multiply
+from curobo.util_file import get_robot_configs_path, join_path, load_yaml
+
+
+@pytest.fixture(scope="module")
+def cfg():
+    cfg = CudaRobotModelConfig.from_robot_yaml_file("franka.yml", "panda_hand")
+    return cfg
+
+
+def test_quaternion():
+    tensor_args = TensorDeviceType()
+
+    def test_q(in_quat):
+        out_rot = quaternion_to_matrix(in_quat)
+        out_quat = matrix_to_quaternion(out_rot.clone())
+        out_quat[..., 1:] *= -1.0
+        q_res = quat_multiply(in_quat, out_quat, in_quat.clone())
+        q_res[..., 0] = 0.0
+        assert torch.sum(q_res).item() <= 1e-5
+
+    in_quat = tensor_args.to_device([1.0, 0.0, 0.0, 0.0]).view(1, 4)
+    test_q(in_quat)
+    test_q(-1.0 * in_quat)
+    in_quat = tensor_args.to_device([0.0, 1.0, 0.0, 0.0]).view(1, 4)
+    test_q(in_quat)
+    test_q(-1.0 * in_quat)
+    in_quat = tensor_args.to_device([0.0, 0.0, 1.0, 0.0]).view(1, 4)
+    test_q(in_quat)
+    test_q(-1.0 * in_quat)
+    in_quat = tensor_args.to_device([0.0, 0.0, 0.0, 1.0]).view(1, 4)
+    test_q(in_quat)
+    test_q(-1.0 * in_quat)
+    in_quat = tensor_args.to_device([0.7071068, 0.0, 0.0, 0.7071068]).view(1, 4)
+    test_q(in_quat)
+    test_q(-1.0 * in_quat)
+
+
+def test_franka_kinematics(cfg):
+    tensor_args = TensorDeviceType()
+
+    robot_model = CudaRobotModel(cfg)
+    q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
+    ee_position = torch.as_tensor([6.0860e-02, -4.7547e-12, 7.6373e-01], **vars(tensor_args)).view(
+        1, -1
+    )
+    ee_quat = torch.as_tensor([0.0382, 0.9193, 0.3808, 0.0922], **vars(tensor_args)).view(1, -1)
+    b_list = [1, 10, 100, 5000]
+    for b in b_list:
+        state = robot_model.get_state(q_test.repeat(b, 1).clone())
+        pos_err = torch.linalg.norm(state.ee_position - ee_position)
+        q_err = torch.linalg.norm(state.ee_quaternion - ee_quat)
+        # check if all values are equal to position and quaternion
+        assert pos_err < 1e-3
+        assert q_err < 1e-1
+
+
+def test_franka_attached_object_kinematics(cfg):
+    tensor_args = TensorDeviceType()
+
+    robot_model = CudaRobotModel(cfg)
+    q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
+    sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
+    state = robot_model.get_state(q_test.clone())
+
+    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
+    # check if all values are equal to position and quaternion
+
+    assert torch.norm(attached_spheres[:, :, 0] - 0.061) < 0.01
+    assert torch.norm(attached_spheres[:, :, 1]) < 0.0001
+
+    # attach an object:
+    new_object = torch.zeros((2, 4), **vars(tensor_args))
+    new_object[:, 3] = 0.01
+    new_object[:, 1] = 0.1
+    robot_model.kinematics_config.update_link_spheres("attached_object", new_object)
+    state = robot_model.get_state(q_test.clone())
+    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
+    robot_model.kinematics_config.detach_object()
+    # assert torch.norm(attached_spheres[:, :2, 0] - 0.0829) < 0.01
+    # assert torch.norm(attached_spheres[:, 2:4, 0] - 0.0829) < 0.01
+
+
+def test_franka_attach_object_kinematics(cfg):
+    tensor_args = TensorDeviceType()
+
+    robot_model = CudaRobotModel(cfg)
+    q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
+    sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
+    state = robot_model.get_state(q_test.clone())
+
+    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
+    # check if all values are equal to position and quaternion
+
+    assert torch.norm(attached_spheres[:, :, 0] - 0.0829) < 0.1
+    assert torch.norm(attached_spheres[:, :, 1]) < 0.0001
+
+    # attach an object:
+    new_object = torch.zeros((4, 4), **vars(tensor_args))
+    new_object[:2, 3] = 0.01
+    new_object[:2, 1] = 0.1
+    # print(attached_spheres[:, sph_idx,:].shape)
+
+    robot_model.kinematics_config.attach_object(sphere_tensor=new_object)
+    state = robot_model.get_state(q_test.clone())
+    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
+    assert (torch.norm(attached_spheres[:, :2, 0] - 0.13)) < 0.01
+    assert torch.norm(attached_spheres[:, 2:, 0] - 0.061) < 0.01
+
+    robot_model.kinematics_config.detach_object()
+    state = robot_model.get_state(q_test.clone())
+    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
+
+    assert torch.norm(attached_spheres[:, :, 0] - 0.061) < 0.01
+    assert torch.norm(attached_spheres[:, :, 1]) < 0.0001
+    robot_model.kinematics_config.detach_object()
+
+
+def test_locked_joints_kinematics():
+    tensor_args = TensorDeviceType()
+
+    config_file = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))
+
+    cfg = CudaRobotModelConfig.from_config(
+        CudaRobotGeneratorConfig(**config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args)
+    )
+    robot_model = CudaRobotModel(cfg)
+
+    cspace = CSpaceConfig(**config_file["robot_cfg"]["kinematics"]["cspace"])
+    cspace.inplace_reindex(robot_model.joint_names)
+    q = cspace.retract_config.unsqueeze(0)
+
+    out = robot_model.get_state(q)
+    j_idx = 2
+    lock_joint_name = cspace.joint_names[j_idx]
+    # lock a joint:
+    cfg.lock_joints = {lock_joint_name: cspace.retract_config[j_idx].item()}
+    cfg = CudaRobotModelConfig.from_config(
+        CudaRobotGeneratorConfig(**config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args)
+    )
+
+    robot_model = CudaRobotModel(cfg)
+    cspace = CSpaceConfig(**config_file["robot_cfg"]["kinematics"]["cspace"])
+
+    cspace.inplace_reindex(robot_model.joint_names)
+    q = cspace.retract_config.unsqueeze(0).clone()
+    out_locked = robot_model.get_state(q)
+    assert torch.linalg.norm(out.ee_position - out_locked.ee_position) < 1e-5
+    assert torch.linalg.norm(out.ee_quaternion - out_locked.ee_quaternion) < 1e-5
+    assert torch.linalg.norm(out.link_spheres_tensor - out_locked.link_spheres_tensor) < 1e-5

tests/motion_gen_api_test.py

@@ -0,0 +1,50 @@
+#
+# 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_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="module")
+def motion_gen():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    robot_file = "franka.yml"
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_file,
+        world_file,
+        tensor_args,
+        trajopt_tsteps=26,
+        use_cuda_graph=False,
+        num_trajopt_seeds=50,
+        fixed_iters_trajopt=True,
+        evaluate_interpolated_trajectory=True,
+    )
+    motion_gen_instance = MotionGen(motion_gen_config)
+    return motion_gen_instance
+
+
+def test_motion_gen_attach_obstacle(motion_gen):
+    obstacle = motion_gen.world_model.objects[-1].name
+    retract_cfg = motion_gen.get_retract_config()
+
+    start_state = JointState.from_position(retract_cfg.view(1, -1))
+    motion_gen.attach_objects_to_robot(start_state, [obstacle])
+    assert True

tests/motion_gen_module_test.py

@@ -0,0 +1,284 @@
+#
+# 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,
+        trajopt_tsteps=32,
+        use_cuda_graph=False,
+        num_trajopt_seeds=50,
+        fixed_iters_trajopt=True,
+        evaluate_interpolated_trajectory=True,
+    )
+    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,
+        trajopt_tsteps=32,
+        use_cuda_graph=False,
+        num_trajopt_seeds=10,
+        fixed_iters_trajopt=True,
+        evaluate_interpolated_trajectory=True,
+    )
+    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.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, quaternion=state.ee_quat_seq)
+
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+    m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10)
+
+    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.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.repeat(2, 1).view(1, -1, 3),
+        quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
+    )
+
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+    m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10)
+
+    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
+    )
+
+
+def test_motion_gen_batch_goalset(motion_gen):
+    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.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)
+
+    m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10)
+
+    result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config)
+
+    # get final solutions:
+    assert torch.count_nonzero(result.success) > 0
+
+    reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-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
+    )
+
+
+def test_motion_gen_batch(motion_gen):
+    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(2)
+
+    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, num_trajopt_seeds=10)
+
+    result = motion_gen.plan_batch(start_state, goal_pose.clone(), m_config)
+    assert torch.count_nonzero(result.success) > 0
+
+    # 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):
+    # return
+    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.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.2
+
+    m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10)
+
+    result = motion_gen_batch_env.plan_batch_env(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.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.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.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, num_trajopt_seeds=10)
+
+    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
+    )

tests/motion_gen_test.py

@@ -0,0 +1,99 @@
+#
+# 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.
+#
+
+
+# CuRobo
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import JointState, RobotConfig
+from curobo.util_file import get_robot_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig
+
+
+def test_motion_gen():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    robot_file = "franka.yml"
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_file,
+        world_file,
+        tensor_args,
+        trajopt_tsteps=40,
+        use_cuda_graph=True,
+    )
+    motion_gen = MotionGen(motion_gen_config)
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+    retract_cfg = motion_gen.get_retract_config()
+    state = motion_gen.rollout_fn.compute_kinematics(
+        JointState.from_position(retract_cfg.view(1, -1))
+    )
+
+    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.5)
+    result = motion_gen.plan(start_state, retract_pose, enable_graph=False)
+    assert result.success.item()
+
+
+def test_motion_gen_attach_object():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    robot_file = "franka.yml"
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_file,
+        world_file,
+        tensor_args,
+        trajopt_tsteps=40,
+        use_cuda_graph=True,
+    )
+    motion_gen = MotionGen(motion_gen_config)
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+    retract_cfg = motion_gen.get_retract_config()
+    state = motion_gen.rollout_fn.compute_kinematics(
+        JointState.from_position(retract_cfg.view(1, -1))
+    )
+
+    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.5)
+    result = motion_gen.plan(start_state, retract_pose, enable_graph=False)
+    assert result.success.item()
+
+    motion_gen.attach_spheres_to_robot(sphere_radius=0.03)
+    result = motion_gen.plan(start_state, retract_pose, enable_graph=False)
+    assert result.success.item()
+
+
+def test_motion_gen_graph():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_test.yml"
+    robot_file = "franka.yml"
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_file, world_file, tensor_args, trajopt_tsteps=40
+    )
+    motion_gen = MotionGen(motion_gen_config)
+    motion_gen.reset()
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    retract_cfg = motion_gen.get_retract_config()
+    state = motion_gen.rollout_fn.compute_kinematics(
+        JointState.from_position(retract_cfg.view(1, -1))
+    )
+
+    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.1)
+    start_state.position[:, 0] = 2.0
+    result = motion_gen.plan(
+        start_state, retract_pose, enable_graph=True, partial_ik_opt=False, enable_opt=False
+    )
+    assert result.success.item()

tests/mpc_test.py

@@ -0,0 +1,235 @@
+#
+# 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 numpy as np
+import pytest
+import torch
+
+# CuRobo
+from curobo.geom.types import WorldConfig
+from curobo.rollout.rollout_base import Goal
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import JointState, RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.mpc import MpcSolver, MpcSolverConfig
+
+
+@pytest.fixture(scope="module")
+def mpc_single_env():
+    tensor_args = TensorDeviceType()
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+    world_file = "collision_test.yml"
+
+    mpc_config = MpcSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_file,
+        use_cuda_graph=False,
+        use_cuda_graph_metrics=False,
+        use_cuda_graph_full_step=False,
+    )
+    mpc = MpcSolver(mpc_config)
+    retract_cfg = robot_cfg.cspace.retract_config.view(1, -1)
+
+    return [mpc, retract_cfg]
+
+
+@pytest.fixture(scope="function")
+def mpc_single_env_lbfgs():
+    tensor_args = TensorDeviceType()
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+    world_file = "collision_test.yml"
+
+    mpc_config = MpcSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_file,
+        use_cuda_graph=True,
+        use_cuda_graph_metrics=True,
+        use_cuda_graph_full_step=False,
+        use_lbfgs=True,
+        use_mppi=False,
+        step_dt=0.5,
+    )
+    mpc = MpcSolver(mpc_config)
+    retract_cfg = robot_cfg.cspace.retract_config.view(1, -1)
+
+    return [mpc, retract_cfg]
+
+
+@pytest.fixture(scope="module")
+def mpc_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_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+    mpc_config = MpcSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        use_cuda_graph=False,
+        use_cuda_graph_metrics=False,
+        use_cuda_graph_full_step=False,
+    )
+    mpc = MpcSolver(mpc_config)
+    retract_cfg = robot_cfg.cspace.retract_config.view(1, -1)
+
+    return [mpc, retract_cfg]
+
+
+@pytest.mark.parametrize(
+    "mpc_str, expected",
+    [
+        ("mpc_single_env", True),
+        ("mpc_single_env_lbfgs", False),
+    ],
+)
+def test_mpc_single(mpc_str, expected, request):
+    mpc_val = request.getfixturevalue(mpc_str)
+    mpc = mpc_val[0]
+    retract_cfg = mpc_val[1]
+    start_state = retract_cfg
+    state = mpc.rollout_fn.compute_kinematics(JointState.from_position(retract_cfg))
+    retract_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
+    goal = Goal(
+        current_state=JointState.from_position(retract_cfg + 0.5),
+        goal_state=JointState.from_position(retract_cfg),
+        goal_pose=retract_pose,
+    )
+    goal_buffer = mpc.setup_solve_single(goal, 1)
+
+    start_state = JointState.from_position(retract_cfg + 0.5, joint_names=mpc.joint_names)
+
+    converged = False
+    tstep = 0
+    mpc.update_goal(goal_buffer)
+    current_state = start_state.clone()
+    while not converged:
+        result = mpc.step(current_state, max_attempts=1)
+        torch.cuda.synchronize()
+        current_state.copy_(result.action)
+        tstep += 1
+        if result.metrics.pose_error.item() < 0.05:
+            converged = True
+            break
+        if tstep > 100:
+            break
+    assert converged == expected
+
+
+def test_mpc_goalset(mpc_single_env):
+    mpc = mpc_single_env[0]
+    retract_cfg = mpc_single_env[1]
+    start_state = retract_cfg
+    state = mpc.rollout_fn.compute_kinematics(JointState.from_position(retract_cfg))
+    retract_pose = Pose(
+        state.ee_pos_seq.repeat(2, 1).unsqueeze(0),
+        quaternion=state.ee_quat_seq.repeat(2, 1).unsqueeze(0),
+    )
+    goal = Goal(
+        current_state=JointState.from_position(retract_cfg + 0.5),
+        goal_state=JointState.from_position(retract_cfg),
+        goal_pose=retract_pose,
+    )
+    goal_buffer = mpc.setup_solve_goalset(goal, 1)
+
+    start_state = JointState.from_position(retract_cfg + 0.5, joint_names=mpc.joint_names)
+
+    converged = False
+    tstep = 0
+    mpc.update_goal(goal_buffer)
+    current_state = start_state.clone()
+    while not converged:
+        result = mpc.step(current_state)
+        torch.cuda.synchronize()
+        current_state.copy_(result.action)
+        tstep += 1
+        if result.metrics.pose_error.item() < 0.01:
+            converged = True
+            break
+        if tstep > 1000:
+            break
+    assert converged
+
+
+def test_mpc_batch(mpc_single_env):
+    mpc = mpc_single_env[0]
+    retract_cfg = mpc_single_env[1].repeat(2, 1)
+    start_state = retract_cfg
+    state = mpc.rollout_fn.compute_kinematics(JointState.from_position(retract_cfg))
+    retract_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
+    retract_pose.position[0, 0] -= 0.02
+    goal = Goal(
+        current_state=JointState.from_position(retract_cfg + 0.5),
+        goal_state=JointState.from_position(retract_cfg),
+        goal_pose=retract_pose,
+    )
+    goal_buffer = mpc.setup_solve_batch(goal, 1)
+
+    start_state = JointState.from_position(retract_cfg + 0.5, joint_names=mpc.joint_names)
+
+    converged = False
+    tstep = 0
+    mpc.update_goal(goal_buffer)
+    current_state = start_state.clone()
+    while not converged:
+        result = mpc.step(current_state)
+        torch.cuda.synchronize()
+        current_state.copy_(result.action)
+        tstep += 1
+        if torch.max(result.metrics.pose_error) < 0.05:
+            converged = True
+            break
+        if tstep > 1000:
+            break
+    assert converged
+
+
+def test_mpc_batch_env(mpc_batch_env):
+    mpc = mpc_batch_env[0]
+    retract_cfg = mpc_batch_env[1].repeat(2, 1)
+    start_state = retract_cfg
+    state = mpc.rollout_fn.compute_kinematics(JointState.from_position(retract_cfg))
+    retract_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
+    goal = Goal(
+        current_state=JointState.from_position(retract_cfg + 0.5, joint_names=mpc.joint_names),
+        goal_state=JointState.from_position(retract_cfg, joint_names=mpc.joint_names),
+        goal_pose=retract_pose,
+    )
+    goal_buffer = mpc.setup_solve_batch_env(goal, 1)
+
+    start_state = JointState.from_position(retract_cfg + 0.5, joint_names=mpc.joint_names)
+
+    converged = False
+    tstep = 0
+    mpc.update_goal(goal_buffer)
+    current_state = start_state.clone()
+    while not converged:
+        result = mpc.step(current_state)
+        torch.cuda.synchronize()
+        current_state.copy_(result.action)
+        tstep += 1
+        if torch.max(result.metrics.pose_error) < 0.05:
+            converged = True
+            break
+        if tstep > 1000:
+            break
+    assert converged

tests/multi_pose_test.py

@@ -0,0 +1,56 @@
+#
+# 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 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 RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
+
+
+def test_multi_pose_franka():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+
+    robot_file = "franka.yml"
+    robot_data = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+
+    robot_data["kinematics"]["link_names"] = robot_data["kinematics"]["collision_link_names"]
+    robot_cfg = RobotConfig.from_dict(robot_data)
+
+    world_cfg = WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
+    ik_config = IKSolverConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        rotation_threshold=0.05,
+        position_threshold=0.005,
+        num_seeds=30,
+        self_collision_check=True,
+        self_collision_opt=True,
+        use_cuda_graph=True,
+        tensor_args=tensor_args,
+    )
+    ik_solver = IKSolver(ik_config)
+    b_size = 1
+
+    q_sample = ik_solver.sample_configs(b_size)
+    kin_state = ik_solver.fk(q_sample)
+    link_poses = kin_state.link_pose
+    goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+    result = ik_solver.solve_single(goal, link_poses=link_poses)
+
+    success = result.success
+    assert torch.count_nonzero(success).item() >= 1.0  # we check if atleast 90% are successful

tests/nvblox_test.py

@@ -0,0 +1,161 @@
+#
+# 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.sdf.world import CollisionQueryBuffer, WorldCollisionConfig
+from curobo.geom.types import Cuboid, WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.util_file import get_world_configs_path, join_path, load_yaml
+
+try:
+    # Third Party
+    from nvblox_torch.mapper import Mapper
+
+    # CuRobo
+    from curobo.geom.sdf.world_blox import WorldBloxCollision
+
+except ImportError:
+    pytest.skip("Nvblox Torch not available", allow_module_level=True)
+
+
+def test_world_blox_trajectory():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_nvblox.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, -1, 1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+    d_sph = coll_check.get_swept_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, act_distance + 0.01, 4, True
+    ).view(-1)
+
+    assert d_sph[0] > 10.0
+    assert d_sph[1] == 0.0
+    assert d_sph[2] == 0.0
+
+
+def test_world_blox():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_nvblox.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+    d_sph = coll_check.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+
+    assert d_sph[0] > 10.0
+    assert d_sph[1] == 0.0
+    assert d_sph[2] == 0.0
+
+
+def test_world_blox_bounding():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_nvblox.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+    bounding_cube = Cuboid("test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.0, 1.0, 1.0])
+    mesh = coll_check.get_mesh_in_bounding_box(
+        bounding_cube,
+        voxel_size=0.05,
+    )
+    assert len(mesh.vertices) > 0
+
+
+def test_world_blox_get_mesh():
+    world_file = "collision_nvblox.yml"
+    tensor_args = TensorDeviceType()
+
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+
+    world_mesh = coll_check.get_mesh_from_blox_layer("world")
+    assert len(world_mesh.vertices) > 10
+
+
+def test_nvblox_world_from_primitive_world():
+    world_file = "collision_cubby.yml"
+    tensor_args = TensorDeviceType()
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict).get_mesh_world(True)
+    mesh = world_cfg.mesh[0].get_trimesh_mesh()
+    world_cfg.mesh[0].save_as_mesh("world.obj")
+    # Third Party
+    from nvblox_torch.datasets.mesh_dataset import MeshDataset
+
+    # CuRobo
+    from curobo.geom.sdf.world import CollisionCheckerType
+    from curobo.geom.types import BloxMap
+    from curobo.types.camera import CameraObservation
+    from curobo.types.math import Pose
+
+    # create a nvblox collision checker:
+    world_config = WorldConfig(
+        blox=[
+            BloxMap(
+                name="world",
+                pose=[0, 0, 0, 1, 0, 0, 0],
+                voxel_size=0.02,
+                integrator_type="tsdf",
+            )
+        ]
+    )
+    config = WorldCollisionConfig(
+        tensor_args=tensor_args, world_model=world_config, checker_type=CollisionCheckerType.BLOX
+    )
+    world_coll = WorldBloxCollision(config)
+    m_dataset = MeshDataset(
+        None, n_frames=100, image_size=1000, save_data_dir=None, trimesh_mesh=mesh
+    )
+    for i in range(len(m_dataset)):
+        data = m_dataset[i]
+        cam_obs = CameraObservation(
+            rgb_image=data["rgba"],
+            depth_image=data["depth"],
+            intrinsics=data["intrinsics"],
+            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+        )
+        world_coll.add_camera_frame(cam_obs, "world")
+
+    world_coll.process_camera_frames("world")
+    world_coll.update_blox_mesh("world")
+    integrated_mesh = world_coll.get_mesh_from_blox_layer("world")
+    if len(integrated_mesh.vertices) > 0:
+        assert True
+        # print("saving World")
+        # integrated_mesh.save_as_mesh("collision_test.obj")
+    else:
+        assert True

tests/pose_test.py

@@ -0,0 +1,66 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.geom.transform import batch_transform_points, transform_points
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+
+
+def test_pose_transform_point():
+    tensor_args = TensorDeviceType()
+    new_pose = Pose.from_list([0, 0, 0, 1, 0, 0, 0], tensor_args)
+
+    new_pose.position.requires_grad = True
+    new_pose.quaternion.requires_grad = True
+
+    points = torch.zeros((3, 3), device=tensor_args.device, dtype=tensor_args.dtype)
+    points[:, 0] = 0.1
+    points[2, 0] = -0.5
+
+    out_pt = new_pose.transform_point(points)
+
+    loss = torch.sum(out_pt)
+    loss.backward()
+    assert torch.norm(new_pose.position.grad) > 0.0
+    assert torch.norm(new_pose.quaternion.grad) > 0.0
+
+
+def test_pose_transform_point_grad():
+    tensor_args = TensorDeviceType()
+    new_pose = Pose.from_list([10.0, 0, 0.1, 1.0, 0, 0, 0], tensor_args)
+    new_pose.position.requires_grad = True
+    new_pose.quaternion.requires_grad = True
+
+    points = torch.zeros((1, 1, 3), device=tensor_args.device, dtype=tensor_args.dtype) + 10.0
+
+    # buffers:
+    out_points = torch.zeros(
+        (points.shape[0], points.shape[1], 3), device=points.device, dtype=points.dtype
+    )
+    out_gp = torch.zeros((new_pose.position.shape[0], 3), device=tensor_args.device)
+    out_gq = torch.zeros((new_pose.position.shape[0], 4), device=tensor_args.device)
+    out_gpt = torch.zeros((points.shape[0], points.shape[1], 3), device=tensor_args.device)
+
+    torch.autograd.gradcheck(
+        batch_transform_points,
+        (new_pose.position, new_pose.quaternion, points, out_points, out_gp, out_gq, out_gpt),
+        eps=1e-6,
+        atol=1.0,
+        # nondet_tol=100.0,
+    )
+
+
+# test_pose_transform_point()
+# test_pose_transform_point_grad()

tests/robot_assets_test.py

@@ -0,0 +1,114 @@
+#
+# 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.cuda_robot_model.cuda_robot_model import CudaRobotModel
+from curobo.geom.types import WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import RobotConfig
+from curobo.types.state import JointState
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig
+
+
+@pytest.mark.parametrize(
+    "robot_file",
+    [
+        "kinova_gen3.yml",
+        "iiwa.yml",
+        "iiwa_allegro.yml",
+        "franka.yml",
+        "ur10e.yml",
+    ],
+)
+class TestRobots:
+    def test_robot_config(self, robot_file):
+        tensor_args = TensorDeviceType()
+        robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+        robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+        pass
+
+    def test_kinematics(self, robot_file):
+        tensor_args = TensorDeviceType()
+
+        robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+        robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+        robot_model = CudaRobotModel(robot_cfg.kinematics)
+        robot_cfg.cspace.inplace_reindex(robot_model.joint_names)
+        robot_model.get_state(robot_cfg.cspace.retract_config)
+        pass
+
+    def test_ik(self, robot_file):
+        world_file = "collision_table.yml"
+        tensor_args = TensorDeviceType()
+
+        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))
+        )
+        ik_config = IKSolverConfig.load_from_robot_config(
+            robot_cfg,
+            world_cfg,
+            rotation_threshold=0.05,
+            position_threshold=0.005,
+            num_seeds=50,
+            self_collision_check=True,
+            self_collision_opt=True,
+            tensor_args=tensor_args,
+        )
+        ik_solver = IKSolver(ik_config)
+        b_size = 10
+        q_sample = ik_solver.sample_configs(b_size)
+        kin_state = ik_solver.fk(q_sample)
+        goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
+        result = ik_solver.solve(goal)
+        result = ik_solver.solve(goal)
+
+        success = result.success
+        assert torch.count_nonzero(success).item() >= 9.0  # we check if atleast 90% are successful
+
+    def notest_motion_gen(self, robot_file):
+        """This test causes pytest to crash when running on many robot configurations
+
+        Args:
+            robot_file: _description_
+        """
+        tensor_args = TensorDeviceType()
+
+        motion_gen_config = MotionGenConfig.load_from_robot_config(
+            robot_file,
+            "collision_table.yml",
+            tensor_args,
+            trajopt_tsteps=40,
+            use_cuda_graph=True,
+        )
+        motion_gen = MotionGen(motion_gen_config)
+        robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+        robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+        retract_cfg = motion_gen.get_retract_config()
+        state = motion_gen.rollout_fn.compute_kinematics(
+            JointState.from_position(retract_cfg.view(1, -1))
+        )
+
+        retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
+        start_state = JointState.from_position(retract_cfg.view(1, -1))
+        start_state.position[0, -1] += 0.2
+        result = motion_gen.plan(start_state, retract_pose, enable_graph=False)
+        assert result.success.item()

tests/robot_config_test.py

@@ -0,0 +1,82 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.types.base import TensorDeviceType
+from curobo.types.robot import CSpaceConfig, JointState
+
+
+def test_cspace_config():
+    cspace_config = CSpaceConfig(
+        ["j0", "j1", "j2", "j3"],
+        retract_config=[i for i in range(4)],
+        cspace_distance_weight=[i for i in range(4)],
+        null_space_weight=[i for i in range(4)],
+    )
+
+    new_order = ["j3", "j1"]
+    cspace_config.inplace_reindex(new_order)
+    assert cspace_config.retract_config[0] == 3 and cspace_config.retract_config[1] == 1
+    assert cspace_config.null_space_weight[0] == 3 and cspace_config.null_space_weight[1] == 1
+    assert (
+        cspace_config.cspace_distance_weight[0] == 3
+        and cspace_config.cspace_distance_weight[1] == 1
+    )
+
+
+def test_joint_state():
+    tensor_args = TensorDeviceType()
+    j_names = ["j0", "j1", "j2", "j3"]
+    loc_j = ["j4", "jb"]
+    final_j = ["jb", "j0", "j1", "j2", "j3", "j4"]
+
+    position = tensor_args.to_device([i for i in range(len(j_names))])
+
+    loc_position = tensor_args.to_device([i + len(j_names) for i in range(len(loc_j))])
+
+    js_1 = JointState.from_position(position, joint_names=j_names)
+    js_lock = JointState.from_position(loc_position, loc_j)
+
+    final_js = js_1.get_augmented_joint_state(final_j, js_lock)
+    assert final_js.joint_names == final_j
+    assert (
+        torch.linalg.norm(final_js.position - tensor_args.to_device([5, 0, 1, 2, 3, 4])).item()
+        < 1e-8
+    )
+
+
+def test_batch_joint_state():
+    tensor_args = TensorDeviceType()
+    j_names = ["j0", "j1", "j2", "j3"]
+    loc_j = ["j4", "jb"]
+    final_j = ["jb", "j0", "j1", "j2", "j3", "j4"]
+
+    # $position = tensor_args.to_device([i for i in range(len(j_names))])
+    position = torch.zeros((10, len(j_names)), device=tensor_args.device, dtype=tensor_args.dtype)
+    for i in range(len(j_names)):
+        position[:, i] = i
+
+    loc_position = tensor_args.to_device([i + len(j_names) for i in range(len(loc_j))])
+
+    js_1 = JointState.from_position(position, joint_names=j_names)
+    js_lock = JointState.from_position(loc_position, loc_j)
+
+    final_js = js_1.get_augmented_joint_state(final_j, js_lock)
+    assert final_js.joint_names == final_j
+    assert (
+        torch.linalg.norm(
+            final_js.position - tensor_args.to_device([5, 0, 1, 2, 3, 4]).unsqueeze(0)
+        ).item()
+        < 1e-8
+    )

tests/robot_world_model_test.py

@@ -0,0 +1,151 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
+
+
+def load_robot_world():
+    robot_file = "franka.yml"
+    world_file = "collision_table.yml"
+    config = RobotWorldConfig.load_from_config(
+        robot_file, world_file, collision_activation_distance=0.2
+    )
+    model = RobotWorld(config)
+    return model
+
+
+def load_robot_batch_world():
+    robot_file = "franka.yml"
+    world_file = ["collision_table.yml", "collision_test.yml"]
+    config = RobotWorldConfig.load_from_config(
+        robot_file, world_file, collision_activation_distance=0.0
+    )
+    model = RobotWorld(config)
+    return model
+
+
+def test_robot_world_config_load():
+    # test if curobo robot world can be loaded
+    load_robot_world()
+    assert True
+
+
+def test_robot_batch_world_config_load():
+    # test if curobo robot world can be loaded
+    load_robot_batch_world()
+    assert True
+
+
+def test_robot_world_kinematics():
+    # test if kinematics can be queried
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    assert state.ee_position.shape[-1] == 3
+
+
+def test_robot_world_sample():
+    # test if joint configurations can be sampled
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    assert q.shape[0] == n
+
+
+def test_robot_world_collision():
+    # test computing collisions given robot joint configurations
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    d_world = model.get_collision_constraint(state.link_spheres_tensor.view(n, 1, -1, 4))
+    assert d_world.shape[0] == n
+    assert torch.sum(d_world) == 0.0
+
+
+def test_robot_world_collision_vector():
+    # test computing collisions given robot joint configurations
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    d_world, d_vec = model.get_collision_vector(state.link_spheres_tensor.view(n, 1, -1, 4))
+
+    assert d_world.shape[0] == n
+    assert torch.norm(d_world[0] - 0.1385) < 0.002
+    assert torch.abs(d_vec[0, 0, 0, 2] + 0.7350) > 0.002
+
+
+def test_robot_world_self_collision():
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    d_self = model.get_self_collision(state.link_spheres_tensor.view(n, 1, -1, 4))
+    assert torch.sum(d_self) == 0.0
+
+
+def test_robot_world_pose():
+    model = load_robot_world()
+    n = 10
+    q = model.sample(n)
+    state = model.get_kinematics(q)
+    pose = state.ee_pose
+    des_pose = state.ee_pose
+    d = model.pose_distance(des_pose, pose)
+    assert torch.sum(d) < 1e-3
+
+
+def test_trajectory_sample():
+    model = load_robot_world()
+    b = 10
+    horizon = 21
+    q = model.sample_trajectory(b, horizon)
+    assert q.shape[0] == b and q.shape[1] == horizon
+
+
+def test_batch_trajectory_sample():
+    model = load_robot_batch_world()
+    b = 2
+    horizon = 21
+    env_query_idx = torch.zeros((b), dtype=torch.int32, device=torch.device("cuda", 0))
+    q = model.sample_trajectory(b, horizon, env_query_idx=env_query_idx)
+    assert q.shape[0] == b and q.shape[1] == horizon
+
+
+def test_batch_trajectory_1env_sample():
+    model = load_robot_batch_world()
+    b = 2
+    horizon = 21
+    env_query_idx = None
+    q = model.sample_trajectory(b, horizon, env_query_idx=env_query_idx)
+    assert q.shape[0] == b and q.shape[1] == horizon
+
+
+def test_robot_batch_world_collision():
+    # test computing collisions given robot joint configurations
+    model = load_robot_batch_world()
+    b = 2
+    horizon = 21
+    env_query_idx = torch.zeros((b), dtype=torch.int32, device=torch.device("cuda", 0))
+
+    q = model.sample_trajectory(b, horizon, env_query_idx=env_query_idx)
+
+    d_world, d_self = model.get_world_self_collision_distance_from_joint_trajectory(
+        q, env_query_idx
+    )
+    assert d_world.shape[0] == b
+    assert torch.sum(d_world) == 0.0

tests/self_collision_test.py

@@ -0,0 +1,106 @@
+#
+# 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.cuda_robot_model.cuda_robot_model import CudaRobotModel
+from curobo.rollout.cost.self_collision_cost import SelfCollisionCost, SelfCollisionCostConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.robot import RobotConfig
+from curobo.util_file import get_robot_configs_path, join_path, load_yaml
+
+
+@pytest.mark.parametrize(
+    "batch_size, horizon",
+    [
+        pytest.param(1, 1, id="1"),
+        pytest.param(10, 1, id="10"),
+        pytest.param(100000, 1, id="100k"),
+        pytest.param(100, 70, id="horizon"),
+    ],
+)
+def test_self_collision_experimental(batch_size, horizon):
+    robot_file = "franka.yml"
+    tensor_args = TensorDeviceType()
+
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_cfg["kinematics"]["debug"] = {"self_collision_experimental": True}
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    kinematics = CudaRobotModel(robot_cfg.kinematics)
+    self_collision_data = kinematics.get_self_collision_config()
+    self_collision_config = SelfCollisionCostConfig(
+        **{"weight": 1.0, "classify": True, "self_collision_kin_config": self_collision_data},
+        tensor_args=tensor_args
+    )
+    cost_fn = SelfCollisionCost(self_collision_config)
+
+    b = batch_size
+    h = horizon
+
+    q = (
+        torch.rand(
+            (b * h, kinematics.get_dof()), device=tensor_args.device, dtype=tensor_args.dtype
+        )
+        * 10
+    )
+    kin_state = kinematics.get_state(q)
+
+    in_spheres = kin_state.link_spheres_tensor
+    in_spheres = in_spheres.view(b, h, -1, 4).contiguous()
+
+    for _ in range(1):
+        out = cost_fn.forward(in_spheres)
+    k = out.clone()
+    cost_fn.self_collision_kin_config.experimental_kernel = False
+    cost_fn._out_distance[:] = 0.0
+    for _ in range(1):
+        out = cost_fn.forward(in_spheres)
+
+    assert torch.norm(k - out).item() < 1e-8
+
+
+def test_self_collision_franka():
+    tensor_args = TensorDeviceType()
+
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    robot_cfg.kinematics.self_collision_config.experimental_kernel = True
+    kinematics = CudaRobotModel(robot_cfg.kinematics)
+    self_collision_data = kinematics.get_self_collision_config()
+    self_collision_config = SelfCollisionCostConfig(
+        **{"weight": 5000.0, "classify": False, "self_collision_kin_config": self_collision_data},
+        tensor_args=tensor_args
+    )
+    cost_fn = SelfCollisionCost(self_collision_config)
+
+    b = 10
+    h = 1
+
+    q = torch.rand(
+        (b * h, kinematics.get_dof()), device=tensor_args.device, dtype=tensor_args.dtype
+    )
+
+    test_q = tensor_args.to_device([2.7735, -1.6737, 0.4998, -2.9865, 0.3386, 0.8413, 0.4371])
+    q[:] = test_q
+    kin_state = kinematics.get_state(q)
+
+    in_spheres = kin_state.link_spheres_tensor
+    in_spheres = in_spheres.view(b, h, -1, 4).contiguous()
+
+    out = cost_fn.forward(in_spheres)
+    assert out.sum().item() > 0.0
+    cost_fn.self_collision_kin_config.experimental_kernel = False
+    cost_fn._out_distance[:] = 0.0
+    out = cost_fn.forward(in_spheres)
+    assert out.sum().item() > 0.0

tests/trajopt_config_test.py

@@ -0,0 +1,132 @@
+#
+# 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.rollout.rollout_base import Goal
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import JointState, RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.trajopt import TrajOptSolver, TrajOptSolverConfig
+
+
+def trajopt_base_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_table.yml"
+
+    robot_file = "franka.yml"
+    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,
+        use_fixed_samples=True,
+        n_collision_envs=1,
+        collision_cache={"obb": 10},
+        seed_ratio={"linear": 0.5, "start": 0.25, "goal": 0.25},
+        num_seeds=10,
+    )
+
+    return trajopt_config
+
+
+def trajopt_es_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_table.yml"
+
+    robot_file = "franka.yml"
+    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,
+        use_es=True,
+        es_learning_rate=0.01,
+    )
+    return trajopt_config
+
+
+def trajopt_gd_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_table.yml"
+
+    robot_file = "franka.yml"
+    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,
+        use_gradient_descent=True,
+        grad_trajopt_iters=500,
+    )
+    return trajopt_config
+
+
+def trajopt_no_particle_opt_config():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_table.yml"
+
+    robot_file = "franka.yml"
+    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,
+        use_particle_opt=False,
+    )
+    return trajopt_config
+
+
+@pytest.mark.parametrize(
+    "config,expected",
+    [
+        (trajopt_base_config(), True),
+        (trajopt_es_config(), True),
+        (trajopt_gd_config(), True),
+        (trajopt_no_particle_opt_config(), True),
+    ],
+)
+def test_eval(config, expected):
+    trajopt_solver = TrajOptSolver(config)
+    q_start = trajopt_solver.retract_config
+    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)
+    goal_state = JointState.from_position(q_goal)
+    current_state = JointState.from_position(q_start)
+    js_goal = Goal(goal_pose=goal_pose, goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver.solve_single(js_goal)
+
+    assert result.success.item() == expected

tests/trajopt_test.py

@@ -0,0 +1,277 @@
+#
+# 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.rollout.rollout_base import Goal
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import JointState, RobotConfig
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.trajopt import TrajOptSolver, TrajOptSolverConfig
+
+
+@pytest.fixture(scope="function")
+def trajopt_solver():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_table.yml"
+
+    robot_file = "franka.yml"
+    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,
+        evaluate_interpolated_trajectory=True,
+    )
+    trajopt_solver = TrajOptSolver(trajopt_config)
+
+    return trajopt_solver
+
+
+@pytest.fixture(scope="function")
+def trajopt_solver_batch_env():
+    tensor_args = TensorDeviceType()
+    world_files = ["collision_table.yml", "collision_cubby.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"
+    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,
+        evaluate_interpolated_trajectory=True,
+    )
+    trajopt_solver = TrajOptSolver(trajopt_config)
+
+    return trajopt_solver
+
+
+def test_trajopt_single_js(trajopt_solver):
+    q_start = trajopt_solver.retract_config
+    q_goal = q_start.clone() + 0.2
+    goal_state = JointState.from_position(q_goal)
+    current_state = JointState.from_position(q_start)
+    # do single planning:
+    js_goal = Goal(goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver.solve_single(js_goal)
+
+    traj = result.solution.position[..., -1, :].view(q_goal.shape)
+    assert torch.linalg.norm((goal_state.position - traj)).item() < 5e-3
+    assert result.success.item()
+
+
+def test_trajopt_single_pose(trajopt_solver):
+    trajopt_solver.reset_seed()
+    q_start = trajopt_solver.retract_config
+    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)
+    goal_state = JointState.from_position(q_goal)
+    current_state = JointState.from_position(q_start)
+    js_goal = Goal(goal_pose=goal_pose, goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver.solve_single(js_goal)
+
+    assert result.success.item()
+
+
+def test_trajopt_single_pose_no_seed(trajopt_solver):
+    trajopt_solver.reset_seed()
+    q_start = trajopt_solver.retract_config
+    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)
+    current_state = JointState.from_position(q_start)
+    js_goal = Goal(goal_pose=goal_pose, current_state=current_state)
+    result = trajopt_solver.solve_single(js_goal)
+
+    # NOTE: This currently fails in some instances.
+    assert result.success.item() == False or result.success.item() == True
+
+
+def test_trajopt_single_goalset(trajopt_solver):
+    # run goalset planning:
+    q_start = trajopt_solver.retract_config
+    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)
+    goal_state = JointState.from_position(q_goal)
+    current_state = JointState.from_position(q_start)
+    g_set = Pose(
+        kin_state.ee_position.repeat(2, 1).view(1, 2, 3),
+        kin_state.ee_quaternion.repeat(2, 1).view(1, 2, 4),
+    )
+    js_goal = Goal(goal_pose=g_set, goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver.solve_goalset(js_goal)
+    assert result.success.item()
+
+
+def test_trajopt_batch(trajopt_solver):
+    # run goalset planning:
+    q_start = trajopt_solver.retract_config.repeat(2, 1)
+    q_goal = q_start.clone()
+    q_goal[0] += 0.1
+    q_goal[1] -= 0.1
+    kin_state = trajopt_solver.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)
+    g_set = Pose(
+        kin_state.ee_position,
+        kin_state.ee_quaternion,
+    )
+
+    js_goal = Goal(goal_pose=g_set, goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver.solve_batch(js_goal)
+    assert torch.count_nonzero(result.success) > 0
+
+
+def test_trajopt_batch_js(trajopt_solver):
+    # run goalset planning:
+    q_start = trajopt_solver.retract_config.repeat(2, 1)
+    q_goal = q_start.clone()
+    q_goal[0] += 0.1
+    q_goal[1] -= 0.1
+    kin_state = trajopt_solver.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)
+
+    js_goal = Goal(goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver.solve_batch(js_goal)
+    traj = result.solution.position
+    interpolated_traj = result.interpolated_solution.position
+    assert torch.count_nonzero(result.success) > 0
+    assert torch.linalg.norm((goal_state.position - traj[:, -1, :])).item() < 0.05
+    assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.05
+
+
+def test_trajopt_batch_goalset(trajopt_solver):
+    # run goalset planning:
+    q_start = trajopt_solver.retract_config.repeat(3, 1)
+    q_goal = q_start.clone()
+    q_goal[0] += 0.1
+    q_goal[1] -= 0.1
+    q_goal[2, -1] += 0.1
+    kin_state = trajopt_solver.fk(q_goal)
+    goal_pose = Pose(
+        kin_state.ee_position.view(3, 1, 3).repeat(1, 5, 1),
+        kin_state.ee_quaternion.view(3, 1, 4).repeat(1, 5, 1),
+    )
+    goal_pose.position[:, 0, 0] -= 0.01
+    goal_state = JointState.from_position(q_goal)
+    current_state = JointState.from_position(q_start)
+
+    js_goal = Goal(goal_state=goal_state, goal_pose=goal_pose, current_state=current_state)
+    result = trajopt_solver.solve_batch_goalset(js_goal)
+    traj = result.solution.position
+    interpolated_traj = result.interpolated_solution.position
+    assert torch.count_nonzero(result.success) > 0
+
+
+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_goal = q_start.clone()
+    q_goal += 0.1
+    q_goal[2] += 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)
+
+    js_goal = Goal(goal_state=goal_state, current_state=current_state)
+    result = trajopt_solver_batch_env.solve_batch_env(js_goal)
+
+    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
+    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_goal = q_start.clone()
+    q_goal[0] += 0.1
+    q_goal[1] -= 0.1
+    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)
+
+    js_goal = Goal(goal_state=goal_state, goal_pose=goal_pose, current_state=current_state)
+    result = trajopt_solver_batch_env.solve_batch_env(js_goal)
+    traj = result.solution.position
+    interpolated_traj = result.interpolated_solution.position
+    assert torch.count_nonzero(result.success) == 3
+
+
+def test_trajopt_batch_env_goalset(trajopt_solver_batch_env):
+    # run goalset planning:
+    q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
+    q_goal = q_start.clone()
+    q_goal[0] += 0.1
+    q_goal[1] -= 0.1
+    q_goal[2, -1] += 0.1
+    kin_state = trajopt_solver_batch_env.fk(q_goal)
+    goal_pose = Pose(
+        kin_state.ee_position.view(3, 1, 3).repeat(1, 5, 1),
+        kin_state.ee_quaternion.view(3, 1, 4).repeat(1, 5, 1),
+    )
+    goal_pose.position[:, 0, 0] -= 0.01
+    goal_state = JointState.from_position(q_goal)
+    current_state = JointState.from_position(q_start)
+
+    js_goal = Goal(goal_state=goal_state, goal_pose=goal_pose, current_state=current_state)
+    result = trajopt_solver_batch_env.solve_batch_env_goalset(js_goal)
+    traj = result.solution.position
+    interpolated_traj = result.interpolated_solution.position
+    assert torch.count_nonzero(result.success) > 0
+
+
+def test_trajopt_batch_env(trajopt_solver):
+    # run goalset planning:
+    q_start = trajopt_solver.retract_config.repeat(3, 1)
+    q_goal = q_start.clone()
+    q_goal[0] += 0.1
+    q_goal[1] -= 0.1
+    q_goal[2, -1] += 0.1
+    kin_state = trajopt_solver.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)
+
+    js_goal = Goal(goal_state=goal_state, goal_pose=goal_pose, current_state=current_state)
+    with pytest.raises(ValueError):
+        result = trajopt_solver.solve_batch_env(js_goal)

tests/warp_mesh_test.py

@@ -0,0 +1,106 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.geom.sdf.world import CollisionQueryBuffer, WorldCollisionConfig
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import Mesh
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util_file import get_assets_path, join_path
+
+
+def test_sdf_pose():
+    mesh_file = join_path(get_assets_path(), "scene/nvblox/srl_ur10_bins.obj")
+    tensor_args = TensorDeviceType()
+    world_config = WorldCollisionConfig(tensor_args)
+    world_ccheck = WorldMeshCollision(world_config)
+    world_ccheck.create_collision_cache(1)
+    new_mesh = Mesh(name="test_mesh", file_path=mesh_file, pose=[0, 0, 0, 1, 0, 0, 0])
+    world_ccheck.add_mesh(
+        new_mesh,
+        env_idx=0,
+    )
+    query_spheres = torch.zeros((1, 1, 2, 4), **vars(tensor_args))
+    query_spheres[..., 2] = 10.0
+    query_spheres[..., 1, :] = 0.0
+    query_spheres[..., 3] = 1.0
+    act_distance = tensor_args.to_device([0.01])
+
+    weight = tensor_args.to_device([1.0])
+    collision_buffer = CollisionQueryBuffer.initialize_from_shape(
+        query_spheres.shape, tensor_args, world_ccheck.collision_types
+    )
+    out = world_ccheck.get_sphere_distance(query_spheres, collision_buffer, act_distance, weight)
+    out = out.view(-1)
+    assert out[0] <= 0.0
+    assert out[1] >= 0.007
+
+
+def test_swept_sdf_speed_pose():
+    mesh_file = join_path(get_assets_path(), "scene/nvblox/srl_ur10_bins.obj")
+    tensor_args = TensorDeviceType()
+    world_config = WorldCollisionConfig(tensor_args)
+    world_ccheck = WorldMeshCollision(world_config)
+    world_ccheck.create_collision_cache(1)
+    new_mesh = Mesh(name="test_mesh", file_path=mesh_file, pose=[0, 0, 0, 1, 0, 0, 0])
+    world_ccheck.add_mesh(
+        new_mesh,
+        env_idx=0,
+    )
+    query_spheres = torch.zeros((1, 1, 2, 4), **vars(tensor_args))
+    query_spheres[..., 2] = 10.0
+    query_spheres[..., 1, :] = 0.0
+    query_spheres[..., 3] = 1.0
+    act_distance = tensor_args.to_device([0.01])
+    dt = act_distance.clone()
+    weight = tensor_args.to_device([1.0])
+    collision_buffer = CollisionQueryBuffer.initialize_from_shape(
+        query_spheres.shape, tensor_args, world_ccheck.collision_types
+    )
+    out = world_ccheck.get_swept_sphere_distance(
+        query_spheres, collision_buffer, weight, act_distance, dt, 2, True, None
+    )
+    out = out.view(-1)
+    assert out[0] <= 0.0
+    assert out[1] >= 0.3
+
+
+def test_swept_sdf_pose():
+    mesh_file = join_path(get_assets_path(), "scene/nvblox/srl_ur10_bins.obj")
+    tensor_args = TensorDeviceType()
+    world_config = WorldCollisionConfig(tensor_args)
+    world_ccheck = WorldMeshCollision(world_config)
+    world_ccheck.create_collision_cache(1)
+    new_mesh = Mesh(name="test_mesh", file_path=mesh_file, pose=[0, 0, 0, 1, 0, 0, 0])
+    world_ccheck.add_mesh(
+        new_mesh,
+        env_idx=0,
+    )
+    query_spheres = torch.zeros((1, 1, 2, 4), **vars(tensor_args))
+    query_spheres[..., 2] = 10.0
+    query_spheres[..., 1, :] = 0.0
+    query_spheres[..., 3] = 1.0
+    act_distance = tensor_args.to_device([0.01])
+    dt = act_distance.clone()
+    weight = tensor_args.to_device([1.0])
+    collision_buffer = CollisionQueryBuffer.initialize_from_shape(
+        query_spheres.shape, tensor_args, world_ccheck.collision_types
+    )
+    out = world_ccheck.get_swept_sphere_distance(
+        query_spheres, collision_buffer, weight, act_distance, dt, 2, False, None
+    )
+    out = out.view(-1)
+    assert out[0] <= 0.0
+    assert out[1] >= 0.01

tests/world_config_test.py

@@ -0,0 +1,171 @@
+#
+# 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 torch
+
+# CuRobo
+from curobo.geom.sdf.world import (
+    CollisionQueryBuffer,
+    WorldCollisionConfig,
+    WorldPrimitiveCollision,
+)
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import Capsule, Cuboid, Cylinder, Mesh, Sphere, WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util_file import get_assets_path, join_path
+
+
+def test_world_modify():
+    tensor_args = TensorDeviceType()
+    obstacle_1 = Cuboid(
+        name="cube_1",
+        pose=[0.0, 0.0, 0.0, 0.043, -0.471, 0.284, 0.834],
+        dims=[0.2, 1.0, 0.2],
+        color=[0.8, 0.0, 0.0, 1.0],
+    )
+
+    # describe a mesh obstacle
+    # import a mesh file:
+
+    mesh_file = join_path(get_assets_path(), "scene/nvblox/srl_ur10_bins.obj")
+
+    obstacle_2 = Mesh(
+        name="mesh_1",
+        pose=[0.0, 2, 0.5, 0.043, -0.471, 0.284, 0.834],
+        file_path=mesh_file,
+        scale=[0.5, 0.5, 0.5],
+    )
+
+    obstacle_3 = Capsule(
+        name="capsule",
+        radius=0.2,
+        base=[0, 0, 0],
+        tip=[0, 0, 0.5],
+        pose=[0.0, 5, 0.0, 0.043, -0.471, 0.284, 0.834],
+        # pose=[0.0, 5, 0.0, 1,0,0,0],
+        color=[0, 1.0, 0, 1.0],
+    )
+
+    obstacle_4 = Cylinder(
+        name="cylinder_1",
+        radius=0.2,
+        height=0.5,
+        pose=[0.0, 6, 0.0, 0.043, -0.471, 0.284, 0.834],
+        color=[0, 1.0, 0, 1.0],
+    )
+
+    obstacle_5 = Sphere(
+        name="sphere_1",
+        radius=0.2,
+        pose=[0.0, 7, 0.0, 0.043, -0.471, 0.284, 0.834],
+        color=[0, 1.0, 0, 1.0],
+    )
+
+    world_model = WorldConfig(
+        mesh=[obstacle_2],
+        cuboid=[obstacle_1],
+        capsule=[obstacle_3],
+        cylinder=[obstacle_4],
+        sphere=[obstacle_5],
+    )
+    world_model.randomize_color(r=[0.2, 0.7], g=[0.4, 0.8], b=[0.0, 0.4])
+
+    collision_support_world = WorldConfig.create_collision_support_world(world_model)
+
+    world_collision_config = WorldCollisionConfig(tensor_args, world_model=collision_support_world)
+    world_ccheck = WorldMeshCollision(world_collision_config)
+
+    world_ccheck.enable_obstacle("sphere_1", False)
+
+    w_pose = Pose.from_list([0, 0, 1, 1, 0, 0, 0], tensor_args)
+    world_ccheck.update_obstacle_pose(name="cylinder_1", w_obj_pose=w_pose)
+
+    x_sph = torch.as_tensor(
+        [[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.1], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, world_ccheck.collision_types
+    )
+    act_distance = tensor_args.to_device([0.0])
+
+    weight = tensor_args.to_device([1])
+    d_sph = world_ccheck.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+    assert d_sph[0] >= 0.2
+    assert d_sph[1] == 0.0
+    assert d_sph[2] >= 0.19
+    world_ccheck.update_obstacle_pose("cube_1", Pose.from_list([1, 0, 0, 1, 0, 0, 0], tensor_args))
+    d_sph = world_ccheck.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+    assert torch.sum(d_sph) == 0.0
+
+    x_sph[0, 0, 0, 1] = 5.0
+    x_sph[0, 0, 0, 3] = 0.2
+
+    d_sph = world_ccheck.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+    assert d_sph[0] >= 0.35
+
+
+def test_batch_collision():
+    tensor_args = TensorDeviceType()
+    world_config_1 = WorldConfig(
+        cuboid=[
+            Cuboid(
+                name="cube_env_1",
+                pose=[0.0, 0.0, 0.0, 1.0, 0, 0, 0],
+                dims=[0.2, 1.0, 0.2],
+                color=[0.8, 0.0, 0.0, 1.0],
+            )
+        ]
+    )
+    world_config_2 = WorldConfig(
+        cuboid=[
+            Cuboid(
+                name="cube_env_2",
+                pose=[0.0, 0.0, 1.0, 1, 0, 0, 0],
+                dims=[0.2, 1.0, 0.2],
+                color=[0.8, 0.0, 0.0, 1.0],
+            )
+        ]
+    )
+
+    world_coll_config = WorldCollisionConfig(
+        tensor_args, world_model=[world_config_1, world_config_2]
+    )
+
+    world_ccheck = WorldPrimitiveCollision(world_coll_config)
+
+    x_sph = torch.zeros((2, 1, 1, 4), device=tensor_args.device, dtype=tensor_args.dtype)
+    x_sph[..., 3] = 0.1
+
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, world_ccheck.collision_types
+    )
+    act_distance = tensor_args.to_device([0.0])
+
+    weight = tensor_args.to_device([1])
+
+    d = world_ccheck.get_sphere_distance(
+        x_sph,
+        query_buffer,
+        weight,
+        act_distance,
+    )
+
+    assert d[0] == 0.2 and d[1] == 0.2
+    env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
+    env_query_idx[1] = 1
+    d = world_ccheck.get_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, env_query_idx=env_query_idx
+    )
+
+    assert d[0] == 0.2 and d[1] == 0.0