Store table IDs as _table_ids in PyBullet envs for motion planning

Add self._table_ids attribute in _store_pybullet_bodies for boil, grow,
coffee, fan, and domino environments. For boil and domino, also capture
the second table's ID in initialize_pybullet.

Author: yichao-liang

predicators/envs/pybullet_boil.py

@@ -358,14 +358,16 @@ def initialize_pybullet(
         )
         bodies["table_id"] = table_id
         # add another table for more space to place jugs and burners
-        create_object(asset_path="urdf/table.urdf",
-                      position=(cls.table_pos[0],
-                                cls.table_pos[1] + (cls.y_ub - cls.y_lb) / 2,
-                                cls.table_pos[2]),
-                      orientation=cls.table_orn,
-                      scale=1.0,
-                      use_fixed_base=True,
-                      physics_client_id=physics_client_id)
+        table_id2 = create_object(
+            asset_path="urdf/table.urdf",
+            position=(cls.table_pos[0],
+                      cls.table_pos[1] + (cls.y_ub - cls.y_lb) / 2,
+                      cls.table_pos[2]),
+            orientation=cls.table_orn,
+            scale=1.0,
+            use_fixed_base=True,
+            physics_client_id=physics_client_id)
+        bodies["table_id2"] = table_id2
 
         # 2) Create jugs
         jug_ids = []
@@ -436,6 +438,8 @@ def initialize_pybullet(
 
     def _store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None:
         """Store references to all PyBullet IDs in the environment objects."""
+        self._table_ids = [pybullet_bodies["table_id"],
+                           pybullet_bodies["table_id2"]]
         self._robot.id = self._pybullet_robot.robot_id
         # Jugs
         for i, jug_obj in enumerate(self._jugs):

predicators/envs/pybullet_coffee.py

@@ -302,6 +302,7 @@ def initialize_pybullet(
         return physics_client_id, pybullet_robot, bodies
 
     def _store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None:
+        self._table_ids = [pybullet_bodies["table_id"]]
         self._table.id = pybullet_bodies["table_id"]
         self._jug.id = pybullet_bodies["jug_id"]
         self._machine.id = pybullet_bodies["machine_id"]

predicators/envs/pybullet_domino/composed_env.py

@@ -238,20 +238,23 @@ def initialize_pybullet(
                                  physics_client_id=physics_client_id)
 
         # Add second table for more space
-        create_object(asset_path="urdf/table.urdf",
-                      position=(cls.table_pos[0],
-                                cls.table_pos[1] + cls.table_width / 2,
-                                cls.table_pos[2]),
-                      orientation=cls.table_orn,
-                      scale=1.0,
-                      use_fixed_base=True,
-                      physics_client_id=physics_client_id)
-
-        bodies = {"table_id": table_id}
+        table_id2 = create_object(
+            asset_path="urdf/table.urdf",
+            position=(cls.table_pos[0],
+                      cls.table_pos[1] + cls.table_width / 2,
+                      cls.table_pos[2]),
+            orientation=cls.table_orn,
+            scale=1.0,
+            use_fixed_base=True,
+            physics_client_id=physics_client_id)
+
+        bodies = {"table_id": table_id, "table_id2": table_id2}
         return physics_client_id, pybullet_robot, bodies
 
     def _store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None:
         """Initialize and store PyBullet bodies for all components."""
+        self._table_ids = [pybullet_bodies["table_id"],
+                           pybullet_bodies["table_id2"]]
         # Initialize each component
         for comp in self._components:
             comp.set_physics_client_id(self._physics_client_id)

predicators/envs/pybullet_fan.py

@@ -563,6 +563,7 @@ def _get_joint_id(obj_id: int,
 
     def _store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None:
         """Store references to all PyBullet object IDs and their joints."""
+        self._table_ids = [pybullet_bodies["table_id"]]
         # 0 = left, 1 = right, 2 = back, 3 = front
 
         # Store all fan IDs grouped by side

predicators/envs/pybullet_grow.py

@@ -235,6 +235,7 @@ def initialize_pybullet(
     def _store_pybullet_bodies(self, pybullet_bodies: Dict[str, Any]) -> None:
         """Store references (IDs) to cups and jugs inside self._cups,
         self._jugs."""
+        self._table_ids = [pybullet_bodies["table_id"]]
         for i, cup in enumerate(self._cups):
             cup.id = pybullet_bodies["cup_ids"][i]
         for i, jug in enumerate(self._jugs):