Edit pybullet coffee so that INTERPRET can run, and add render=True to.

Author: ashayathalye

predicators/cogman.py

@@ -215,9 +215,10 @@ def run_episode_and_get_observations(
     render_obs = cogman.get_approach_name == "oracle" and\
                  CFG.offline_data_method == "geo_and_demo_with_vlm_imgs"
     if isinstance(env, PyBulletEnv):
-        obs = env.get_observation(render=render_obs)
+        # obs = env.get_observation(render=render_obs)
+        obs = env.get_observation(render=True)
     else:
-        obs = env.get_observation()
+        obs = env.get_observation(render=True)
     observations = [obs]
     actions: List[Action] = []
     curr_option: Optional[_Option] = None

predicators/datasets/generate_atom_trajs_with_vlm.py

@@ -750,7 +750,6 @@ def _generate_ground_atoms_with_vlm_pure_visual_preds(
         vlm: VisionLanguageModel) -> List[List[Set[GroundAtom]]]:
     """Given a collection of ImageOptionTrajectories, query a VLM to convert
     these into ground atom trajectories."""
-    import pdb; pdb.set_trace()
     if not CFG.grammar_search_vlm_atom_proposal_use_debug:
         logging.info("Querying VLM for candidate atom proposals...")
         atom_strs_proposals_list = _sample_vlm_atom_proposals_from_trajectories(
@@ -1128,6 +1127,7 @@ def create_ground_atom_data_from_generated_demos(
         generate_func = _generate_ground_atoms_with_vlm_pure_visual_preds
     ground_atoms_trajs = generate_func(img_option_trajs, env, train_tasks,
                                        known_predicates, all_task_objs, vlm)
+    # return Dataset(option_segmented_trajs, None)
     return Dataset(option_segmented_trajs, ground_atoms_trajs)
 
 

predicators/envs/pybullet_coffee.py

@@ -48,7 +48,7 @@
     create_single_arm_pybullet_robot
 from predicators.settings import CFG
 from predicators.structs import Action, Array, EnvironmentTask, Object, \
-    Predicate, State, Observation, Task
+    Predicate, State, Observation, Task, Type
 
 class PyBulletCoffeeEnv(PyBulletEnv, CoffeeEnv):
     """PyBullet Coffee domain.
@@ -214,6 +214,14 @@ def pour_z_offset(cls) -> float:
     _camera_pitch: ClassVar[float]
     _camera_target: ClassVar[Pose3D]
 
+    # Types
+    _table_type = Type("table", [])
+    _robot_type = Type("robot", ["x", "y", "z", "tilt", "wrist", "fingers"])
+    _jug_type = Type("jug", ["x", "y", "z", "rot", "is_held", "is_filled"])
+    _machine_type = Type("coffee_machine", ["is_on"])
+    _cup_type = Type("cup", ["x", "y", "z", "capacity_liquid", "target_liquid", "current_liquid"])
+    _plug_type = Type("plug", ["x", "y", "z", "plugged_in"])
+
     def __init__(self, use_gui: bool = True) -> None:
         super().__init__(use_gui)
 
@@ -626,7 +634,6 @@ def _get_state(self, render_obs: bool = False) -> State:
         assert set(state) == set(self._current_state), \
             (f"Reconstructed state has objects {set(state)}, but "
              f"self._current_state has objects {set(self._current_state)}.")
-
         return state
 
     def step(self, action: Action, render_obs: bool = False) -> State: