Minor fixes and geom module documentation

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