renamed function to make more sense and fixed failing test

mplib/planner.py

@@ -31,7 +31,7 @@ def __init__(
         joint_acc_limits: Optional[Sequence[float] | np.ndarray] = None,
         **kwargs,
     ):
-    # constructor ankor end
+        # constructor ankor end
         """Motion planner for robots.
 
         Args:
@@ -215,21 +215,22 @@ def distance_6D(self, p1, q1, p2, q2):
             np.linalg.norm(q1 - q2), np.linalg.norm(q1 + q2)
         )
 
-    def check_joint_limit(self, q):
+    def wrap_joint_limit(self, q) -> bool:
         """
-        check if the joint configuration is within the joint limits
+        Checks if the joint configuration is within the joint limits.
+        For revolute joints, the joint angle is wrapped to be within [q_min, q_min+2*pi)
 
         Args:
-            q: joint configuration
+            q: joint configuration, angles of revolute joints might be modified
 
         Returns:
-            True if the joint configuration is within the joint limits
+            True if q can be wrapped to be within the joint limits
         """
         n = len(q)
         flag = True
         for i in range(n):
             if self.joint_types[i].startswith("JointModelR"):
-                if np.abs(q[i] - self.joint_limits[i][0]) < 1e-3:
+                if -1e-3 < q[i] - self.joint_limits[i][0] < 0:
                     continue
                 q[i] -= (
                     2 * np.pi * np.floor((q[i] - self.joint_limits[i][0]) / (2 * np.pi))
@@ -368,7 +369,7 @@ def IK(self, goal_pose, start_qpos, mask=None, n_init_qpos=20, threshold=1e-3):
             ik_results = self.pinocchio_model.compute_IK_CLIK(
                 index, goal_pose, start_qpos, mask
             )
-            flag = self.check_joint_limit(ik_results[0])  # will clip qpos
+            flag = self.wrap_joint_limit(ik_results[0])  # will wrap revolute joints
 
             # check collision
             self.planning_world.set_qpos_all(ik_results[0][idx])
@@ -718,6 +719,7 @@ def plan_qpos_to_pose(
             constraint_tolerance,
             verbose=verbose,
         )
+
     # plan_screw ankor
     def plan_screw(
         self,
@@ -730,7 +732,7 @@ def plan_screw(
         wrt_world=True,
         verbose=False,
     ):
-    # plan_screw ankor end
+        # plan_screw ankor end
         """
         Plan from a start configuration to a goal pose of the end-effector using
         screw motion

tests/test_planner.py

@@ -75,21 +75,18 @@ def test_planning_screw(self):
     self.planner.robot.set_qpos(last_qpos)
     self.assertTrue(np.allclose(self.get_end_effector_pose(), self.target_pose, atol=1e-2))
 
-  # this test is a bit cursed
-  # def test_check_joint_limit(self, tolerance=2e-2):
-  #   for _ in range(100):
-  #     qpos = np.random.uniform(-np.pi, np.pi, size=7)
-  #     in_limit = True
-  #     for joint_angle, limit in zip(qpos, self.joint_limits):
-  #       if joint_angle < limit[0]:
-  #         joint_angle += 2 * np.pi - tolerance
-  #       elif joint_angle > limit[1]:
-  #         joint_angle -= 2 * np.pi + tolerance
-  #       if not (limit[0] <= joint_angle <= limit[1]):
-  #         in_limit = False
-  #         break
+  def test_wrap_joint_limit(self, tolerance=2e-2):
+    for _ in range(100):
+      qpos = np.random.uniform(-100, 100, size=7)
+      in_limit = True
+      for joint_angle, limit in zip(qpos, self.joint_limits):
+        k = np.ceil((limit[0]-joint_angle)/2/np.pi)
+        joint_angle += 2*np.pi*k
+        if not (limit[0]-tolerance <= joint_angle <= limit[1]+tolerance):
+          in_limit = False
+          break
 
-  #     self.assertEqual(self.planner.check_joint_limit(qpos), in_limit, f"Joint limit check failed for qpos: {qpos} which should be {'in' if in_limit else 'out'} of limit")
+      self.assertEqual(self.planner.wrap_joint_limit(qpos), in_limit, f"Joint limit check failed for qpos: {qpos} which should be {'in' if in_limit else 'out'} of limit")
 
   def test_pad_qpos(self):
     for _ in range(100):