[gym/envs] Automatically detect relevant contact points for Atlas.

duburcqa · Dec 4, 2021 · 7bd6196 · 7bd6196
1 parent 377e355
commit 7bd6196
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 9 deletions.
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -97,11 +97,11 @@ deployPythonPackage("${LIBRARY_NAME}_py")
 ################ gym_jiminy ################
 
 buildPythonWheel("python/gym_${LIBRARY_NAME}/common"
-                 "python/gym_${LIBRARY_NAME}/envs"
                  "python/gym_${LIBRARY_NAME}/toolbox"
+                 "python/gym_${LIBRARY_NAME}/envs"
                  "python/gym_${LIBRARY_NAME}/rllib"
                  "${CMAKE_BINARY_DIR}/pypi/dist/gym_${LIBRARY_NAME}")
 deployPythonPackageDevelop("python/gym_${LIBRARY_NAME}/common"
-                           "python/gym_${LIBRARY_NAME}/envs"
                            "python/gym_${LIBRARY_NAME}/toolbox"
+                           "python/gym_${LIBRARY_NAME}/envs"
                            "python/gym_${LIBRARY_NAME}/rllib")
diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py b/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py
@@ -12,6 +12,7 @@
     F_PROFILE_WAVELENGTH, F_PROFILE_PERIOD)
 from gym_jiminy.common.controllers import PDController
 from gym_jiminy.common.pipeline import build_pipeline
+from gym_jiminy.toolbox.math import ConvexHull
 
 from pinocchio import neutral
 
@@ -79,6 +80,23 @@
 }
 
 
+def _cleanup_contact_points(env: "AtlasJiminyEnv") -> None:
+    contact_frames_idx = env.robot.contact_frames_idx
+    contact_frames_names = env.robot.contact_frames_names
+    num_contacts = int(len(env.robot.contact_frames_idx) // 2)
+    for contact_slice in (slice(num_contacts), slice(num_contacts, None)):
+        contact_positions = np.stack([
+            env.robot.pinocchio_data.oMf[frame_idx].translation
+            for frame_idx in contact_frames_idx[contact_slice]], axis=0)
+        contact_bottom_idx = np.argsort(
+            contact_positions[:, 2])[:int(num_contacts//2)]
+        convex_hull = ConvexHull(contact_positions[contact_bottom_idx, :2])
+        env.robot.remove_contact_points([
+            contact_frames_names[contact_slice][i]
+            for i in set(range(num_contacts)).difference(
+                contact_bottom_idx[convex_hull._vertex_indices])])
+
+
 class AtlasJiminyEnv(WalkerJiminyEnv):
     def __init__(self, debug: bool = False, **kwargs: Any) -> None:
         # Get the urdf and mesh paths
@@ -98,9 +116,7 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None:
             debug=debug), **kwargs})
 
         # Remove unrelevant contact points
-        self.robot.remove_contact_points([
-            name for name in self.robot.contact_frames_names
-            if int(name.split("_")[-1]) in (1, 3, 5, 7)])
+        _cleanup_contact_points(self)
 
     def _neutral(self):
         def joint_position_idx(joint_name):
@@ -200,9 +216,7 @@ def joint_position_idx(joint_name):
             step_dt=STEP_DT, debug=debug), **kwargs})
 
         # Remove unrelevant contact points
-        self.robot.remove_contact_points([
-            name for name in self.robot.contact_frames_names
-            if int(name.split("_")[-1]) in (1, 3, 5, 7)])
+        _cleanup_contact_points(self)
 
 
 AtlasPDControlJiminyEnv = build_pipeline(**{

diff --git a/python/gym_jiminy/toolbox/setup.py b/python/gym_jiminy/toolbox/setup.py
@@ -29,7 +29,7 @@
     keywords="reinforcement-learning robotics gym jiminy",
     packages=find_namespace_packages(),
     install_requires=[
-        f"gym_jiminy=={version}"
+        f"gym_jiminy_toolbox=={version}"
     ],
     zip_safe=False
 )