Merge pull request #72 from haosulab/srdf-generation

Srdf generation

include/mplib/kinematics/pinocchio/pinocchio_model.h

@@ -1,8 +1,10 @@
 #pragma once
 
 #include <memory>
+#include <set>
 #include <string>
 #include <tuple>
+#include <utility>
 #include <vector>
 
 #include <pinocchio/multibody/data.hpp>
@@ -73,6 +75,13 @@ class PinocchioModelTpl {
    */
   const std::vector<std::string> &getLeafLinks() const { return leaf_links_; }
 
+  /**
+   * Get the all adjacent link names.
+   *
+   * @return: adjacent link names as a set of pairs of strings
+   */
+  std::set<std::pair<std::string, std::string>> getAdjacentLinks() const;
+
   /**
    * Pinocchio might have a different link order or it might add additional links.
    *

mplib/planner.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import os
+import xml.etree.ElementTree as ET
 from collections.abc import Callable
 from typing import Optional, Sequence
+from xml.dom import minidom
 
 import numpy as np
 import toppra as ta
@@ -11,7 +13,7 @@
 from transforms3d.quaternions import mat2quat, quat2mat
 
 from mplib.pymp import ArticulatedModel, PlanningWorld
-from mplib.pymp.collision_detection import WorldCollisionResult
+from mplib.pymp.collision_detection import AllowedCollisionMatrix, WorldCollisionResult
 from mplib.pymp.planning import ompl
 
 
@@ -77,25 +79,13 @@ def __init__(
         self.user_link_names = self.pinocchio_model.get_link_names()
         self.user_joint_names = self.pinocchio_model.get_joint_names()
 
-        self.planning_world = PlanningWorld(
-            [self.robot],
-            ["robot"],
-            kwargs.get("normal_objects", []),
-            kwargs.get("normal_object_names", []),
-        )
-        self.acm = self.planning_world.get_allowed_collision_matrix()
-
         self.joint_name_2_idx = {}
         for i, joint in enumerate(self.user_joint_names):
             self.joint_name_2_idx[joint] = i
         self.link_name_2_idx = {}
         for i, link in enumerate(self.user_link_names):
             self.link_name_2_idx[link] = i
 
-        if self.srdf == "":
-            self.generate_collision_pair()
-            self.robot.update_SRDF(self.srdf)
-
         assert (
             move_group in self.user_link_names
         ), f"end-effector not found as one of the links in {self.user_link_names}"
@@ -134,6 +124,18 @@ def __init__(
             t.startswith("JointModelR") for t in self.joint_types
         ] & (self.joint_limits[:, 1] - self.joint_limits[:, 0] > 2 * np.pi)
 
+        self.planning_world = PlanningWorld(
+            [self.robot],
+            ["robot"],
+            kwargs.get("normal_objects", []),
+            kwargs.get("normal_object_names", []),
+        )
+        self.acm = self.planning_world.get_allowed_collision_matrix()
+
+        if self.srdf == "":
+            self.generate_collision_pair()
+            self.robot.update_SRDF(self.srdf)
+
         self.planner = ompl.OMPLPlanner(world=self.planning_world)
 
     def replace_package_keyword(self, package_keyword_replacement):
@@ -155,56 +157,95 @@ def replace_package_keyword(self, package_keyword_replacement):
                         out_f.write(content)
         return rtn_urdf
 
-    def generate_collision_pair(self, sample_time=1000000, echo_freq=100000):
+    def generate_collision_pair(self, num_samples=100000):
         """
         We read the srdf file to get the link pairs that should not collide.
         If not provided, we need to randomly sample configurations
         to find the link pairs that will always collide.
+
+        :param num_samples: number of samples to find the link that will always collide
         """
         print(
             "Since no SRDF file is provided. We will first detect link pairs that will"
             " always collide. This may take several minutes."
         )
-        n_link = len(self.user_link_names)
-        cnt = np.zeros((n_link, n_link), dtype=np.int32)
-        for i in range(sample_time):
-            qpos = self.pinocchio_model.get_random_configuration()
-            self.robot.set_qpos(qpos, True)
-            collisions = self.planning_world.collide_full()
-            for collision in collisions:
-                u = self.link_name_2_idx[collision.link_name1]
-                v = self.link_name_2_idx[collision.link_name2]
-                cnt[u][v] += 1
-            if i % echo_freq == 0:
-                print("Finish %.1f%%!" % (i * 100 / sample_time))
-
-        import xml.etree.ElementTree as ET
-        from xml.dom import minidom
 
         root = ET.Element("robot")
         robot_name = self.urdf.split("/")[-1].split(".")[0]
         root.set("name", robot_name)
         self.srdf = self.urdf.replace(".urdf", ".srdf")
 
-        for i in range(n_link):
-            for j in range(n_link):
-                if cnt[i][j] == sample_time:
-                    link1 = self.user_link_names[i]
-                    link2 = self.user_link_names[j]
+        acm = AllowedCollisionMatrix()
+
+        # 1. disable adjacent link pairs
+        for link1, link2 in self.pinocchio_model.get_adjacent_links():
+            print(f"Ignore collision pair: ({link1}, {link2}), reason: Adjacent")
+            acm.set_entry(link1, link2, True)
+            _ = ET.SubElement(
+                root,
+                "disable_collisions",
+                attrib={"link1": link1, "link2": link2, "reason": "Adjacent"},
+            )
+
+        # 2. disable all-zeros qpos (default) collision
+        self.robot.set_qpos(np.zeros(len(self.user_joint_names)), True)
+        for collision in self.check_for_self_collision():
+            link1, link2 = collision.link_name1, collision.link_name2
+            if acm.get_entry(link1, link2) is not None:  # already ignored
+                continue
+            print(
+                f"Ignore collision pair: ({link1}, {link2}), "
+                "reason: Default (collides at all-zeros qpos)"
+            )
+            acm.set_entry(link1, link2, True)
+            _ = ET.SubElement(
+                root,
+                "disable_collisions",
+                attrib={"link1": link1, "link2": link2, "reason": "Default"},
+            )
+
+        # 3. disable collision pairs that always collide and never collide via sampling
+        n_links = len(self.user_link_names)
+        collision_cnt = np.zeros((n_links, n_links), dtype=int)
+        for _ in range(num_samples):
+            self.robot.set_qpos(self.pinocchio_model.get_random_configuration(), True)
+            for collision in self.check_for_self_collision():
+                u = self.link_name_2_idx[collision.link_name1]
+                v = self.link_name_2_idx[collision.link_name2]
+                collision_cnt[u][v] += 1
+
+        for i, link1 in enumerate(self.user_link_names):
+            for j in range(i + 1, n_links):
+                link2 = self.user_link_names[j]
+                if acm.get_entry(link1, link2) is not None:  # already ignored
+                    continue
+                if (cnt := (collision_cnt[i][j] + collision_cnt[j][i])) == num_samples:
                     print(
                         f"Ignore collision pair: ({link1}, {link2}), "
-                        "reason: always collide"
+                        "reason: Always collide"
+                    )
+                    _ = ET.SubElement(
+                        root,
+                        "disable_collisions",
+                        attrib={"link1": link1, "link2": link2, "reason": "Always"},
                     )
-                    collision = ET.SubElement(root, "disable_collisions")
-                    collision.set("link1", link1)
-                    collision.set("link2", link2)
-                    collision.set("reason", "Default")
+                elif cnt == 0:
+                    print(
+                        f"Ignore collision pair: ({link1}, {link2}), "
+                        "reason: Never collide"
+                    )
+                    _ = ET.SubElement(
+                        root,
+                        "disable_collisions",
+                        attrib={"link1": link1, "link2": link2, "reason": "Never"},
+                    )
+
+        # Save SRDF
         with open(self.srdf, "w") as srdf_file:
             srdf_file.write(
-                minidom.parseString(ET.tostring(root)).toprettyxml(indent="    ")
+                minidom.parseString(ET.tostring(root)).toprettyxml(indent="  ")
             )
-            srdf_file.close()
-        print("Saving the SRDF file to %s" % self.srdf)
+        print(f"Saving the SRDF file to {self.srdf}")
 
     def distance_6D(self, p1, q1, p2, q2):
         """

mplib/pymp/kinematics/pinocchio.pyi

@@ -161,6 +161,12 @@ class PinocchioModel:
             frame of the link
         :return: 6 x n jacobian of the link
         """
+    def get_adjacent_links(self) -> set[tuple[str, str]]:
+        """
+        Get the all adjacent link names.
+
+        :return: adjacent link names as a set of pairs of strings
+        """
     def get_chain_joint_index(self, end_effector: str) -> list[int]:
         """
         Get the joint indices of the joints in the chain from the root to the given

pybind/docstring/kinematics/pinocchio/pinocchio_model.h

@@ -118,6 +118,12 @@ Constructs a PinocchioModel from URDF string
 :param verbose: print debug information. Default: ``False``.
 :return: a unique_ptr to PinocchioModel)doc";
 
+static const char *__doc_mplib_kinematics_pinocchio_PinocchioModelTpl_getAdjacentLinks =
+R"doc(
+Get the all adjacent link names.
+
+:return: adjacent link names as a set of pairs of strings)doc";
+
 static const char *__doc_mplib_kinematics_pinocchio_PinocchioModelTpl_getChainJointIndex =
 R"doc(
 Get the joint indices of the joints in the chain from the root to the given

pybind/kinematics/pinocchio/pybind_pinocchio_model.cpp

@@ -40,6 +40,8 @@ void build_pypinocchio_model(py::module &m) {
 
       .def("get_leaf_links", &PinocchioModel::getLeafLinks,
            DOC(mplib, kinematics, pinocchio, PinocchioModelTpl, getLeafLinks))
+      .def("get_adjacent_links", &PinocchioModel::getAdjacentLinks,
+           DOC(mplib, kinematics, pinocchio, PinocchioModelTpl, getAdjacentLinks))
 
       .def("set_link_order", &PinocchioModel::setLinkOrder, py::arg("names"),
            DOC(mplib, kinematics, pinocchio, PinocchioModelTpl, setLinkOrder))

src/kinematics/pinocchio/pinocchio_model.cpp

@@ -236,6 +236,26 @@ void PinocchioModelTpl<S>::dfsParseTree(const urdf::LinkConstSharedPtr &link,
   if (link->child_links.empty()) leaf_links_.push_back(link->name);
 }
 
+template <typename S>
+std::set<std::pair<std::string, std::string>> PinocchioModelTpl<S>::getAdjacentLinks()
+    const {
+  std::set<std::pair<std::string, std::string>> ret;
+
+  const auto root_name = urdf_model_->getRoot()->name;
+  for (auto link_name : getLeafLinks())
+    while (link_name != root_name) {
+      const auto parent_link = urdf_model_->getLink(link_name)->getParent();
+      const auto parent_name = parent_link->name;
+      ret.insert({parent_name, link_name});
+      // if the parent has a fixed link to its parent, also disable collision
+      if (parent_name != root_name &&
+          parent_link->parent_joint->type == urdf::Joint::FIXED)
+        ret.insert({parent_link->getParent()->name, link_name});
+      link_name = parent_name;
+    }
+  return ret;
+}
+
 template <typename S>
 void PinocchioModelTpl<S>::setLinkOrder(const std::vector<std::string> &names) {
   user_link_names_ = names;