Updated robot tutorial (#297)

* docs * work * work * more docs
haosulab · Apr 30, 2024 · 082f02e · 082f02e
1 parent f95f72a
commit 082f02e
Show file tree

Hide file tree

Showing 12 changed files with 256 additions and 50 deletions.
diff --git a/docs/source/_static/videos/panda_random_actions.mp4 b/docs/source/_static/videos/panda_random_actions.mp4
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -1,4 +1,4 @@
-__version__ = "3.0.0.dev14"
+__version__ = "3.0.0.b0"
 # Configuration file for the Sphinx documentation builder.
 #
 # For the full list of built-in configuration values, see the documentation:

diff --git a/docs/source/user_guide/tutorials/custom_robots.md b/docs/source/user_guide/tutorials/custom_robots.md
diff --git a/docs/source/user_guide/tutorials/images/anymal-visual-collision.png b/docs/source/user_guide/tutorials/images/anymal-visual-collision.png
diff --git a/docs/source/user_guide/tutorials/images/loaded_anymal_panda_0.png b/docs/source/user_guide/tutorials/images/loaded_anymal_panda_0.png
diff --git a/docs/source/user_guide/tutorials/images/loaded_anymal_panda_keyframe.png b/docs/source/user_guide/tutorials/images/loaded_anymal_panda_keyframe.png
diff --git a/docs/source/user_guide/tutorials/images/panda_robot_camera_lines.png b/docs/source/user_guide/tutorials/images/panda_robot_camera_lines.png
diff --git a/mani_skill/agents/robots/fetch/fetch.py b/mani_skill/agents/robots/fetch/fetch.py
@@ -366,18 +366,10 @@ def _after_init(self):
         self.base_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "base_link"
         )
-        self.l_wheel_link: Link = sapien_utils.get_obj_by_name(
-            self.robot.get_links(), "l_wheel_link"
-        )
-        self.r_wheel_link: Link = sapien_utils.get_obj_by_name(
-            self.robot.get_links(), "r_wheel_link"
-        )
+        self.l_wheel_link: Link = self.robot.links_map["l_wheel_link"]
+        self.r_wheel_link: Link = self.robot.links_map["r_wheel_link"]
         for link in [self.l_wheel_link, self.r_wheel_link]:
-            for body in link._bodies:
-                cs = body.get_collision_shapes()[0]
-                cg = cs.get_collision_groups()
-                cg[2] |= FETCH_UNIQUE_COLLISION_BIT
-                cs.set_collision_groups(cg)
+            link.set_collision_group_bit(group=2, bit_idx=30, bit=1)
 
         self.torso_lift_link: Link = sapien_utils.get_obj_by_name(
             self.robot.get_links(), "torso_lift_link"

diff --git a/mani_skill/envs/tasks/empty_env.py b/mani_skill/envs/tasks/empty_env.py
@@ -26,7 +26,7 @@ def __init__(self, *args, robot_uids="panda", **kwargs):
 
     @property
     def _default_sensor_configs(self):
-        pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])
+        pose = sapien_utils.look_at([1.25, -1.25, 1.5], [0.0, 0.0, 0.2])
         return [CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)]
 
     @property
@@ -35,7 +35,8 @@ def _default_human_render_camera_configs(self):
         return CameraConfig("render_camera", pose, 2048, 2048, 1, 0.01, 100)
 
     def _load_scene(self, options: dict):
-        build_ground(self._scene)
+        self.ground = build_ground(self._scene)
+        self.ground.set_collision_group_bit(group=2, bit_idx=30, bit=1)
 
     def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
         pass

diff --git a/mani_skill/examples/demo_robot.py b/mani_skill/examples/demo_robot.py
@@ -26,7 +26,8 @@ def parse_args(args=None):
     )
     args = parser.parse_args()
     return args
-if __name__ == "__main__":
+
+def main():
     args = parse_args()
     env = gym.make(
         "Empty-v1",
@@ -67,3 +68,6 @@ def parse_args(args=None):
             assert kf is not None, "this robot has no keyframes, cannot use it to set actions"
             env.step(kf.qpos)
         viewer = env.render()
+
+if __name__ == "__main__":
+    main()
diff --git a/mani_skill/utils/structs/actor.py b/mani_skill/utils/structs/actor.py
@@ -216,6 +216,14 @@ def is_static(self, lin_thresh=1e-2, ang_thresh=1e-1):
             torch.linalg.norm(self.angular_velocity, axis=1) <= ang_thresh,
         )
 
+    def set_collision_group_bit(self, group: int, bit_idx: int, bit: int):
+        """Set's a specific collision group bit for all collision shapes in all parallel actors"""
+        for body in self._bodies:
+            for cs in body.get_collision_shapes():
+                cg = cs.get_collision_groups()
+                cg[group] |= bit_idx << bit
+                cs.set_collision_groups(cg)
+
     # -------------------------------------------------------------------------- #
     # Exposed actor properties, getters/setters that automatically handle
     # CPU and GPU based actors

diff --git a/mani_skill/utils/structs/link.py b/mani_skill/utils/structs/link.py
@@ -181,7 +181,9 @@ def set_collision_group_bit(self, group: int, bit_idx: int, bit: int):
         """Set's a specific collision group bit for all collision shapes in all parallel actors"""
         for body in self._bodies:
             for cs in body.get_collision_shapes():
-                cs.collision_groups[group] |= bit << bit_idx
+                cg = cs.collision_groups
+                cg[group] |= bit << bit_idx
+                cs.set_collision_groups(cg)
 
     # -------------------------------------------------------------------------- #
     # Functions from sapien.Component