From 38226e09c75696bf22b76074ce05afdb3833a7de Mon Sep 17 00:00:00 2001
From: Nishanth Kumar <nishanth.kumar20@gmail.com>
Date: Wed, 9 Oct 2024 13:14:08 -0400
Subject: [PATCH] setup to get results we missed

---
 predicators/utils.py                    |  3 +-
 scripts/configs/pred_invention_vlm.yaml | 65 +++++++++++++------------
 2 files changed, 35 insertions(+), 33 deletions(-)

diff --git a/predicators/utils.py b/predicators/utils.py
index cedc41d23..8694a13a7 100644
--- a/predicators/utils.py
+++ b/predicators/utils.py
@@ -2485,7 +2485,8 @@ def query_vlm_for_atom_vals(
     all_vlm_responses = vlm_output_str.strip().split("\n")
     # NOTE: this assumption is likely too brittle; if this is breaking, feel
     # free to remove/adjust this and change the below parsing loop accordingly!
-    assert len(atom_queries_list) == len(all_vlm_responses)
+    if len(atom_queries_list) != len(all_vlm_responses):
+        return set()
     for i, (atom_query, curr_vlm_output_line) in enumerate(
             zip(atom_queries_list, all_vlm_responses)):
         assert atom_query + ":" in curr_vlm_output_line
diff --git a/scripts/configs/pred_invention_vlm.yaml b/scripts/configs/pred_invention_vlm.yaml
index c9d9167ec..5035f526a 100644
--- a/scripts/configs/pred_invention_vlm.yaml
+++ b/scripts/configs/pred_invention_vlm.yaml
@@ -1,41 +1,41 @@
 # Experiments to test predicate invention with VLMs
 ---
 APPROACHES:
-  ours:
-    NAME: "grammar_search_invention"
-    FLAGS:
-      grammar_search_vlm_atom_proposal_prompt_type: options_labels_whole_traj_diverse
-      offline_data_method: geo_and_demo_with_vlm_imgs
+  # ours:
+  #   NAME: "grammar_search_invention"
+  #   FLAGS:
+  #     grammar_search_vlm_atom_proposal_prompt_type: options_labels_whole_traj_diverse
+  #     offline_data_method: geo_and_demo_with_vlm_imgs
   ours-no-subselection:
     NAME: "grammar_search_invention"
     FLAGS:
       grammar_search_vlm_atom_proposal_prompt_type: options_labels_whole_traj_diverse
       offline_data_method: geo_and_demo_with_vlm_imgs
       grammar_search_pred_selection_approach: no_select
-  ours-no-invent:
-    NAME: "nsrt_learning"
-    FLAGS: {}
-  ours-no-visual:
-    NAME: "grammar_search_invention"
-    FLAGS: 
-      offline_data_method: demo
-  ours-no-geo:
-    NAME: "grammar_search_invention"
-    FLAGS:
-      grammar_search_vlm_atom_proposal_prompt_type: demo_with_vlm_imgs
-  interpret:
-    NAME: "grammar_search_invention"
-    FLAGS:
-      offline_data_method: demo_with_vlm_imgs
-      vlm_predicate_vision_api_generate_ground_atoms: True
-  vila-with-fewshot:
-    NAME: "vlm_open_loop"
-    FLAGS:
-      vlm_open_loop_use_training_demos: True
-  vila-pure:
-    NAME: "vlm_open_loop"
-    FLAGS:
-      vlm_open_loop_use_training_demos: False
+  # ours-no-invent:
+  #   NAME: "nsrt_learning"
+  #   FLAGS: {}
+  # ours-no-visual:
+  #   NAME: "grammar_search_invention"
+  #   FLAGS: 
+  #     offline_data_method: demo
+  # ours-no-geo:
+  #   NAME: "grammar_search_invention"
+  #   FLAGS:
+  #     grammar_search_vlm_atom_proposal_prompt_type: demo_with_vlm_imgs
+  # interpret:
+  #   NAME: "grammar_search_invention"
+  #   FLAGS:
+  #     offline_data_method: demo_with_vlm_imgs
+  #     vlm_predicate_vision_api_generate_ground_atoms: True
+  # vila-with-fewshot:
+  #   NAME: "vlm_open_loop"
+  #   FLAGS:
+  #     vlm_open_loop_use_training_demos: True
+  # vila-pure:
+  #   NAME: "vlm_open_loop"
+  #   FLAGS:
+  #     vlm_open_loop_use_training_demos: False
 
 ENVS:
   burger_no_move_more_stacks:
@@ -67,11 +67,12 @@ ENVS:
       cluster_and_intersect_prune_low_data_pnads: True
       cluster_and_intersect_min_datastore_fraction: 0.2
 
-ARGS: []
+ARGS:
+- "load_data"
 FLAGS: 
   vlm_model_name: gpt-4o
   num_train_tasks: 8
   num_test_tasks: 10
-START_SEED: 0
-NUM_SEEDS: 5
+START_SEED: 3
+NUM_SEEDS: 2
 ...