Minor fixes and geom module documentation

tests/curobo_robot_world_model_test.py

@@ -135,8 +135,9 @@ def test_cu_robot_batch_world_collision():
     assert d_world.shape[0] == b
     assert torch.sum(d_world) == 0.0
 
+
 def test_cu_robot_get_link_transform():
     model = load_robot_world()
     world_T_panda_hand = model.kinematics.get_link_transform("panda_hand")
     # It seems the panda hand is initialized at the origin.
-    assert torch.sum(world_T_panda_hand.position) == 0.0
+    assert torch.sum(world_T_panda_hand.position) == 0.0

tests/mpc_test.py

@@ -21,7 +21,7 @@ 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
+from curobo.wrap.reacher.mpc import MpcSolver, MpcSolverConfig, PoseCostMetric
 
 
 @pytest.fixture(scope="module")
@@ -233,3 +233,65 @@ def test_mpc_batch_env(mpc_batch_env):
         if tstep > 1000:
             break
     assert converged
+
+
+@pytest.mark.parametrize(
+    "mpc_str, expected",
+    [
+        ("mpc_single_env", True),
+        # ("mpc_single_env_lbfgs", True), unstable
+    ],
+)
+def test_mpc_single_pose_metric(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))
+
+    state = mpc.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+    goal_pose = state.ee_pose.clone()
+
+    start_state = JointState.from_position(
+        retract_cfg.view(1, -1) + 0.3, joint_names=mpc.joint_names
+    )
+
+    start_pose = mpc.compute_kinematics(start_state).ee_pose.clone()
+    goal_pose.position = start_pose.position.clone()
+    goal_pose.quaternion = start_pose.quaternion.clone()
+
+    if mpc.project_pose_to_goal_frame:
+        offset_pose = Pose.from_list([0, 0.1, 0, 1, 0, 0, 0])
+        goal_pose = goal_pose.multiply(offset_pose)
+    else:
+        goal_pose.position[0, 1] += 0.2
+
+    goal = Goal(
+        current_state=JointState.from_position(retract_cfg + 0.5),
+        goal_state=JointState.from_position(retract_cfg),
+        goal_pose=goal_pose,
+    )
+    goal_buffer = mpc.setup_solve_single(goal, 1)
+
+    converged = False
+    tstep = 0
+    mpc.update_goal(goal_buffer)
+    current_state = start_state.clone()
+
+    pose_cost_metric = PoseCostMetric(
+        hold_partial_pose=True,
+        hold_vec_weight=mpc.tensor_args.to_device([1, 1, 1, 1, 0, 1]),
+    )
+    mpc.update_pose_cost_metric(pose_cost_metric, start_state, goal_pose)
+    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 > 200:
+            break
+    assert converged == expected

tests/world_config_test.py

@@ -34,9 +34,6 @@ def test_world_modify():
         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(
@@ -170,3 +167,70 @@ def test_batch_collision():
     )
 
     assert d[0] == 0.2 and d[1] == 0.0
+
+
+def test_world_modify_cpu():
+    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],
+    )
+
+    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, update_cpu_reference=True
+    )
+
+    assert world_ccheck.world_model.get_obstacle("cylinder_1").pose[2] == 1