Updated code to call o1

agents/recoma/README.md

@@ -5,7 +5,7 @@
 1. Install the Recoma library
 
 ```shell
-pip install git+https://github.com/allenai/recoma.git@b262ab5476841551c5d6874e8ddecef78c71a072
+pip install recoma==0.0.4
 ```
 
 2. Set the OPENAI_API_KEY env variable
@@ -190,3 +190,17 @@ for t in "Archaeology Dating" "Plant Nutrients" "Space Sick" "Combinatorial Chem
 ```
 
 
+### Running o1
+
+To run on o1-mini, you can run the following script. Make corresponding changes for the other datasets
+```shell
+export DIFF=Normal
+export TASK="Small Skills"
+export MAX_ENV_CALLS=10
+export SEED=123 # Used for GPT
+export MODEL=o1-mini-2024-09-12
+export OUTPUT_DIR=output_dir/react/${DIFF}_${MAX_ENV_CALLS}env_${MODEL}_s${SEED}/${TASK// /_}
+python agents/recoma/run_recoma.py \
+    --output_dir ${OUTPUT_DIR} \
+    --config agents/recoma/configs/react_o1.jsonnet
+```

agents/recoma/configs/default_o1_generator.libsonnet

@@ -0,0 +1,9 @@
+{
+    "type": "openai_chat",
+    "model": std.extVar("MODEL"),
+    "max_tokens": 10000,
+    "temperature": 0.0,
+    "use_cache": true,
+    "seed": std.parseInt(std.extVar("SEED")),
+    "stop": ["\n"],
+}

agents/recoma/configs/react_o1.jsonnet

@@ -0,0 +1,50 @@
+local output_dir = std.extVar('OUTPUT_DIR');
+local max_env_calls = std.parseInt(std.extVar('MAX_ENV_CALLS'));
+local task = std.extVar('TASK');
+local diff = std.extVar('DIFF');
+local generator_params = import "default_o1_generator.libsonnet";
+{
+    "models": {
+        "discoveryworld_init": {
+            "type": "discoveryworld_loader",
+            "threadid_offset": max_env_calls / 100,
+            "next_model": "react"
+        },
+        "react": {
+            "type": "discoveryworld_react_controller",
+            "action_model": "action",
+            "observation_model": "environment",
+            "add_roles": true,
+            "max_output_length": 10000,
+            "max_history": -1  // Maximum number of history steps to show
+        },
+        "action": {
+            "type": "discoveryworld_promptedlm",
+            "prompt_file": "agents/recoma/prompts/react_prompt.txt",
+            "generator_params": generator_params + {"max_tokens": 400, "stop": ["```\n"]},
+        },
+        "environment": {
+            "type": "discoveryworld_env",
+            "output_dir": output_dir,
+        },
+    },
+    "search": {
+        "type": "best_first",
+        "start_model": "discoveryworld_init",
+        "answerer": {
+            "type": "discoveryworld_answerer",
+            "output_dir": output_dir,
+        },
+        "stopping_conditions": [
+            {"type": "max_env_calls", "max_env_calls": max_env_calls},
+            {"type": "max_llm_calls", "max_llm_calls": 1000}, // Not necessary; mainly there to catch any rogue usage
+            {"type": "max_llm_cost", "max_llm_cost": 50.00}  // Not necessary; mainly there to catch any rogue usage
+        ]
+    },
+    "reader": {
+       "type": "discoveryworld_reader",
+       "limit_prefixes": [task],
+       "limit_difficulties": [diff],
+       "limit_seeds": [1, 2, 3, 4, 5],
+    }
+}

agents/recoma/discoveryworld_env_models.py

@@ -111,7 +111,10 @@ def generate_output(self, state: SearchState) -> GenerationOutputs:
                 "DiscoveryWorldEnvironmentModel called without any open node!!")
         env = SingletonEnvironment().env
         # execute the input against the DiscoveryWorld environment
-        action_json = json.loads(current_node.input_str)
+        if current_node.input_str:
+            action_json = json.loads(current_node.input_str)
+        else:
+            action_json = {}
         output = env.performAgentAction(agentIdx=0, actionJSON=action_json)
         if self.output_dir is not None:
             output_tracking_info(output_dir=self.output_dir, state=state, action=action_json,

agents/recoma/react_controller.py

@@ -122,8 +122,11 @@ def append_message_to_history(self, current_history: List[Any], last_child: Sear
         else:
             action_json = self.extract_json_output(last_child)
             try:
-                formatted_json = json.loads(action_json)
-                current_history.append(Action(action_str=last_child.output, action_json=formatted_json))
+                if action_json:
+                    formatted_json = json.loads(action_json)
+                    current_history.append(Action(action_str=last_child.output, action_json=formatted_json))
+                else:
+                    current_history.append(Action(action_str=last_child.output, action_json={}))
             except json.JSONDecodeError:
                 raise ValueError("Failed to decode JSON from action output: {}".format(last_child.output))