migrate klampt engine and visualiser to new style

Signed-off-by: Tin Lai <[email protected]>

collisionChecker.py

@@ -103,19 +103,22 @@ def get_coor_before_collision(self, pos1, pos2):
         return endPos
 
     def visible(self, pos1, pos2):
+        self.stats.visible_cnt += 1
         try:
             # get list of pixel between node A and B
             # pixels = lineGenerationAlgorithm(pos1, pos2)
             pixels = self.get_line(pos1, pos2)
             # check that all pixel are white (free space)
             for p in pixels:
-                if not self.feasible(p):
+                if not self.feasible(p, save_stats=False):
                     return False
         except ValueError:
             return False
         return True
 
-    def feasible(self, p):
+    def feasible(self, p, save_stats=True):
+        if save_stats:
+            self.stats.feasible_cnt += 1
         try:
             return self._img[tuple(map(int, p))] == 1
         except IndexError:
@@ -196,71 +199,73 @@ def __init__(self, xml: str, stats: Stats):
         import klampt
         from klampt.plan import robotplanning
 
-        self.cc_feasible = 0
-        self.cc_visible = 0
-
         world = klampt.WorldModel()
         world.readFile(xml)  # very cluttered
         robot = world.robot(0)
 
-        # this is the CSpace that will be used.  Standard collision and joint limit constraints
-        # will be checked
+        # this is the CSpace that will be used.
+        # Standard collision and joint limit constraints will be checked
         space = robotplanning.makeSpace(world, robot, edgeCheckResolution=0.1)
 
-        # fire up a visual editor to get some start and goal configurations
-        qstart = robot.getConfig()
-        qgoal = robot.getConfig()
-
         self.space = space
         self.robot = robot
         self.world = world
 
         import copy
 
-        self.template_pos = copy.copy(qstart)
+        self.template_pos = copy.copy(robot.getConfig())
         self.template_pos[1:7] = [0] * 6
 
-        self.qstart = self.translate_from_klampt(qstart)
-        self.qgoal = self.translate_from_klampt(qgoal)
-
     def get_dimension(self):
         return 6
 
     def get_dimension_limits(self):
-        return self.robot.getJointLimits()
+        return self.space.bound
 
     def translate_to_klampt(self, p):
         """Translate the given configuration to klampt's configuration
 
-        :param p: configuratino to translate
+        :param p: configuration to translate
 
         """
-        assert len(p) == 6, p
-
-        new_pos = list(self.template_pos)
-        new_pos[1:7] = p
-        return new_pos
+        assert len(p) == self.get_dimension(), p
+        # add a dummy end configuration that denotes the gripper's angle
+        return list(p) + [0]
 
     def translate_from_klampt(self, p):
         """Translate the given klampt's configuration to our protocol
 
-        :param p: configuratino to translate
+        :param p: configuration to translate
 
         """
-        assert len(p) == 12, len(p)
-        return p[1:7]
+        if len(p) == 12:
+            # this is what klampt's robot interface return, mixed with arm joints and
+            # gripper angles
+            return p[1:7]
+        elif len(p) == 7:
+            # this is what klampt's space interface return, mixed with arm joints and
+            # gripper rotation (single value)
+            return p[:6]
+        raise ValueError(f"Unknown length {p}")
+
+    def get_eef_world_pos(self, config):
+        _config = self.robot.getConfig()
+        _config[1:7] = list(config)[:6]
+        self.robot.setConfig(_config)
+        link = self.robot.link(11)  # the eef
+        pos = link.getWorldPosition([0, 0, 0])
+        return pos
 
     def visible(self, a, b):
         a = self.translate_to_klampt(a)
         b = self.translate_to_klampt(b)
-        # print(self.space.visible(a, b))
         self.stats.visible_cnt += 1
         return self.space.isVisible(a, b)
 
-    def feasible(self, p, stats=False):
+    def feasible(self, p, save_stats=True):
         p = self.translate_to_klampt(p)
         self.stats.feasible_cnt += 1
-        return self.space.feasible(p)
+        return self.space.isFeasible(p)
 
 
 class RobotArm4dCollisionChecker(CollisionChecker):
@@ -359,8 +364,9 @@ def _interpolate_configs(self, c1, c2):
     def visible(self, pos1, pos2):
         # get list of pixel between node A and B
         # pixels = lineGenerationAlgorithm(pos1, pos2)
+        self.stats.visible_cnt += 1
         for p in self._interpolate_configs(pos1, pos2):
-            if not self.feasible(p):
+            if not self.feasible(p, save_stats=False):
                 return False
         return True
 
@@ -375,7 +381,9 @@ def _pt_feasible(self, p):
         except IndexError:
             return False
 
-    def feasible(self, p):
+    def feasible(self, p, save_stats=True):
+        if save_stats:
+            self.stats.feasible_cnt += 1
         pt1 = p[:2]
         pt2 = self.get_pt_from_angle_and_length(
             pt1, p[2], self.stick_robot_length_config[0]

env.py

@@ -4,7 +4,6 @@
 import random
 
 import numpy as np
-from scipy import spatial
 from tqdm import tqdm
 
 import collisionChecker
@@ -111,6 +110,14 @@ def parse_input_pt(pt_as_str):
                     f">than 1.0). It might not work well as klampt uses "
                     f"radian for joints value"
                 )
+            if self.args["radius"] > 2:
+                import warnings
+
+                warnings.warn(
+                    f"Radius value is very high at {self.args['radius']} ("
+                    f">than 2.0). It might not work well as klampt uses "
+                    f"radian for joints value"
+                )
 
     def __getattr__(self, attr):
         """This is called what self.attr doesn't exist.
@@ -120,18 +127,17 @@ def __getattr__(self, attr):
 
     @staticmethod
     def radian_dist(p1: np.ndarray, p2: np.ndarray):
-        """Return the cosine similarity for radian
+        """Return the (possibly wrapped) distance between two vector of angles in
+        radians.
 
         :param p1: first configuration :math:`q_1`
         :param p2: second configuration :math:`q_2`
 
         """
-        # distance metric for angles
-        # https://en.wikipedia.org/wiki/Cosine_similarity#Angular_distance_and_similarity
-
-        cosine_similarity = 1 - spatial.distance.cosine(p1, p2)
-        return np.arccos(cosine_similarity) / np.pi
-        # return spatial.distance.cosine(p1, p2)
+        # https://stackoverflow.com/questions/28036652/finding-the-shortest-distance-between-two-angles/28037434
+        diff = (p2 - p1 + np.pi) % (2 * np.pi) - np.pi
+        diff = np.where(diff < -np.pi, diff + (2 * np.pi), diff)
+        return np.linalg.norm(diff)
 
     @staticmethod
     def euclidean_dist(p1: np.ndarray, p2: np.ndarray):
@@ -156,7 +162,7 @@ def step_from_to(self, p1: np.ndarray, p2: np.ndarray):
         """
         if self.args.ignore_step_size:
             return p2
-        if np.all(p1 == p2):
+        if np.isclose(p1, p2).all():
             # p1 is at the same point as p2
             return p2
         unit_vector = p2 - p1
@@ -194,4 +200,4 @@ def run(self):
                 )
                 pbar.refresh()
 
-        self.planner.terminates_hook()
+        self.visualiser.terminates_hook()

main.py

@@ -36,7 +36,6 @@
 
 General Planner Options:
   --radius=RADIUS        Set radius for node connections.
-                         [default: 12.0]
   --epsilon=EPSILON      Set epsilon value.
                          [default: 10.0]
   --max-number-nodes=MAX_NODES
@@ -200,7 +199,6 @@ def generate_args(
         len(planner_canidates) == 1
     ), f"Planner to use '{planner_canidates}' has length {len(planner_canidates)}"
     planner_to_use = planner_canidates[0]
-    print(planner_to_use)
 
     planner_data_pack = planner_registry.PLANNERS[planner_to_use]
 
@@ -211,6 +209,14 @@ def generate_args(
         **{re.sub(r"^--", "", k).replace("-", "_"): v for (k, v) in args.items()},
     )
 
+    # setup defaults
+    if args["--engine"] == "klampt":
+        # setup defaults for klampt
+        if args["--epsilon"] == "10.0":
+            args["--epsilon"] = 1.0
+    if args["--radius"] is None:
+        args["--radius"] = float(args["--epsilon"]) * 1.1
+
     planning_option = MagicDict(
         planner_data_pack=planner_data_pack,
         skip_optimality=args["--skip-optimality"],

planners/birrtPlanner.py

@@ -146,11 +146,28 @@ def pygame_birrt_planner_paint(planner):
         planner.visualiser.draw_solution_path()
 
 
+from planners.rrtPlanner import klampt_draw_nodes_paint_func
+
+
+def klampt_birrt_paint(planner):
+    """Visualiser paint function for BiRRT
+
+    :param planner: the planner to be visualised
+
+    """
+    for c, nodes in (
+        ((1, 0, 0, 1), planner.nodes),
+        ((0, 0, 1, 1), planner.goal_tree_nodes),
+    ):
+        klampt_draw_nodes_paint_func(planner, nodes, c)
+
+
 # start register
 planner_registry.register_planner(
     "birrt",
     planner_class=BiRRTPlanner,
     visualise_pygame_paint=pygame_birrt_planner_paint,
+    visualise_klampt_paint=klampt_birrt_paint,
     sampler_id="birrt_sampler",
 )
 # finish register

planners/prmPlanner.py

@@ -189,12 +189,53 @@ def pygame_prm_planner_paint_when_terminate(planner):
     input("\nPress Enter to quit...")
 
 
+def klampt_prm_paint(planner) -> None:
+    """Visualiser paint function for PRM
+
+    :param planner: the planner to be visualised
+
+    """
+
+    colour = (1, 0, 0, 1)
+    for n in planner.nodes:
+        planner.args.env.draw_node(
+            planner.args.env.cc.get_eef_world_pos(n.pos), colour=colour
+        )
+    for edge in planner.graph.edges:
+        edge = np.array(edge).transpose()
+        planner.args.env.draw_path(
+            planner.args.env.cc.get_eef_world_pos(n.pos),
+            planner.args.env.cc.get_eef_world_pos(n.parent.pos),
+            colour=colour,
+        )
+
+
+def klampt_prm_planner_paint_when_terminate(planner):
+    """Visualisation function to paint for planner when termiante
+
+    :param planner: the planner to visualise
+
+    """
+    from utils.common import Colour
+
+    planner.build_graph()
+    # draw all edges
+    for n1, n2 in planner.tree.edges():
+        planner.args.env.draw_path(n1, n2, Colour.path_blue)
+    planner.get_solution()
+    planner.args.env.update_screen()
+
+    input("\nPress Enter to quit...")
+
+
 # start register
 planner_registry.register_planner(
     "prm",
     planner_class=PRMPlanner,
     visualise_pygame_paint=pygame_prm_planner_paint,
     visualise_pygame_paint_terminate=pygame_prm_planner_paint_when_terminate,
+    visualise_klampt_paint=klampt_prm_paint,
+    visualise_klampt_paint_terminate=klampt_prm_planner_paint_when_terminate,
     sampler_id=prmSampler.sampler_id,
 )
 # finish register