release repository

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