improve unit test

CHANGELOG.md

@@ -10,6 +10,12 @@ its affiliates is strictly prohibited.
 -->
 # Changelog
 
+## Latest Commit
+
+### BugFixes & Misc.
+- Fix bug in evaluator to account for dof maximum acceleration and jerk.
+- Add unit test for different acceleration and jerk limits.
+
 ## Version 0.7.2
 
 ### New Features

tests/motion_gen_eval_test.py

@@ -16,45 +16,57 @@ import torch
 # CuRobo
 from curobo.types.base import TensorDeviceType
 from curobo.types.robot import JointState
-from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
+from curobo.util_file import get_robot_configs_path, join_path, load_yaml
 from curobo.wrap.reacher.evaluator import TrajEvaluatorConfig
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
 
-@pytest.fixture(scope="module", params=[True, False])
-def evaluate_interpolated_trajectory(request):
-    return request.param
 
-@pytest.fixture(scope="module")
-def motion_gen(evaluate_interpolated_trajectory):
+def run_motion_gen(robot_file, evaluate_interpolated_trajectory, max_acc, max_jerk):
     tensor_args = TensorDeviceType()
     world_file = "collision_test.yml"
-    robot_file = "franka.yml"
-    dof = 9
-    traj_evaluator_config = TrajEvaluatorConfig(
-         max_acc=torch.ones((dof), device=tensor_args.device, dtype=tensor_args.dtype),
-         max_jerk=torch.ones((dof), device=tensor_args.device, dtype=tensor_args.dtype),
-         min_dt=torch.tensor(0.01, device=tensor_args.device, dtype=tensor_args.dtype),
-         max_dt=torch.tensor(1.5, device=tensor_args.device, dtype=tensor_args.dtype)
-    )
+    robot_data = load_yaml(join_path(get_robot_configs_path(), robot_file))
+    dof = len(robot_data["robot_cfg"]["kinematics"]["cspace"]["joint_names"])
+
+    robot_data["robot_cfg"]["kinematics"]["cspace"]["max_acceleration"] = [
+        max_acc for i in range(9)
+    ]
+    robot_data["robot_cfg"]["kinematics"]["cspace"]["max_jerk"] = [max_jerk for i in range(9)]
+
     motion_gen_config = MotionGenConfig.load_from_robot_config(
-        robot_file,
+        robot_data,
         world_file,
         tensor_args,
-        trajopt_tsteps=26,
         use_cuda_graph=False,
-        num_trajopt_seeds=50,
-        fixed_iters_trajopt=True,
+        maximum_trajectory_dt=1.5,
         evaluate_interpolated_trajectory=evaluate_interpolated_trajectory,
-        traj_evaluator_config=traj_evaluator_config,
     )
     motion_gen = MotionGen(motion_gen_config)
-    motion_gen.warmup(warmup_js_trajopt=True)
 
     retract_cfg = motion_gen.get_retract_config()
 
     start_state = JointState.from_position(retract_cfg.view(1, -1).clone())
-    goal_state = JointState.from_position(retract_cfg.view(1, -1).clone())
+    goal_state = JointState.from_position(retract_cfg.view(1, -1).clone() + 0.2)
 
-    result = motion_gen.plan_single_js(start_state, goal_state, MotionGenPlanConfig(max_attempts=1))
+    result = motion_gen.plan_single_js(
+        start_state, goal_state, MotionGenPlanConfig(max_attempts=5, enable_graph_attempt=10)
+    )
+    return result
 
-def test_motion_gen(motion_gen):
-    return True
+
+@pytest.mark.parametrize(
+    "robot_file, evaluate_interpolated_traj, max_acc, max_jerk",
+    [
+        ("franka.yml", False, 1.0, 500.0),
+        ("franka.yml", True, 0.1, 500.0),
+        ("franka.yml", True, 1.0, 500.0),
+        ("ur5e.yml", False, 1.0, 500.0),
+        ("ur5e.yml", True, 0.1, 500.0),
+        ("ur5e.yml", True, 1.0, 500.0),
+    ],
+)
+def test_motion_gen_trajectory(robot_file, evaluate_interpolated_traj, max_acc, max_jerk):
+    result = run_motion_gen(robot_file, evaluate_interpolated_traj, max_acc, max_jerk)
+
+    assert result.success.item()
+    assert torch.max(torch.abs(result.optimized_plan.acceleration)) <= max_acc
+    assert torch.max(torch.abs(result.optimized_plan.jerk)) <= max_jerk