sample demonstrations randomly

Author: BartekCupial

balrog/dataset.py

@@ -1,4 +1,5 @@
 import glob
+import logging
 import os
 import random
 import re
@@ -63,9 +64,15 @@ def load_episode(self, filename):
             episode = {k: data[k] for k in data.files}
         return episode
 
-    def load_in_context_learning_episode(self, i, task, agent):
+    def load_in_context_learning_episodes(self, num_episodes, task, agent):
         demo_task = self.demo_task(task)
-        demo_path = self.demo_path(i, demo_task)
+        demo_paths = [self.demo_path(i, demo_task) for i in range(num_episodes)]
+        random.shuffle(demo_paths)
+
+        for demo_path in demo_paths:
+            self.load_in_context_learning_episode(demo_path, agent)
+
+    def load_in_context_learning_episode(self, demo_path, agent):
         episode = self.load_episode(demo_path)
 
         actions = episode.pop("action").tolist()
@@ -89,6 +96,6 @@ def load_in_context_learning_episode(self, i, task, agent):
                 break
 
         if not done:
-            print("warning: icl trajectory ended without done")
+            logging.info("icl trajectory ended without done")
 
         agent.wrap_episode()

balrog/evaluator.py

@@ -290,8 +290,7 @@ def run_episode(self, task, agent, process_num=None, position=0, episode_idx=0):
 
             # If the agent is an FewShotAgent, load the in-context learning episode
             if isinstance(agent, FewShotAgent):
-                for icl_episode in range(self.config.eval.icl_episodes):
-                    self.dataset.load_in_context_learning_episode(icl_episode, task, agent)
+                self.dataset.load_in_context_learning_episodes(self.config.eval.icl_episodes, task, agent)
 
                 if self.config.agent.cache_icl and self.config.client.client_name == "gemini":
                     agent.cache_icl()