Accomplish the exp scripts v1; Add video recording and trajectory recording of desktop agent; Fix minor bugs

experiment.py

@@ -44,7 +44,7 @@ logger = logging.getLogger("desktopenv.experiment")
 PATH_TO_VM = r"C:\Users\tianbaox\Documents\Virtual Machines\Ubuntu\Ubuntu.vmx"
 
 
-def run_one_example(example, agent, max_steps=20, example_trajectory_dir="exp_trajectory"):
+def run_one_example(example, agent, max_steps=2, example_trajectory_dir="exp_trajectory", recording=True):
     trajectory_recording_path = os.path.join(example_trajectory_dir, "trajectory.json")
     env = DesktopEnv(
         path_to_vm=PATH_TO_VM,
@@ -57,25 +57,53 @@ def run_one_example(example, agent, max_steps=20, example_trajectory_dir="exp_tr
     done = False
     step_num = 0
 
-    # todo: save the screenshots and actions to a folder
+    if recording:
+        # send a request to the server to start recording
+        env.controller.start_recording()
+
     while not done and step_num < max_steps:
         actions = agent.predict(observation)
         for action in actions:
+            step_num += 1
+
+            # Capture the timestamp before executing the action
+            action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
+
             observation, reward, done, info = env.step(action)
             observation['instruction'] = example['instruction']
-        step_num += 1
-        logger.info("Step %d", step_num)
-        logger.info("Action: %s", actions)
-        observation.pop("accessibility_tree")
-        logger.info("Observation: %s", observation)
-        logger.info("Reward: %.2f", reward)
-        logger.info("Info: %s", info)
 
-        logger.info("================================\n")
+            # Logging
+            logger.info("Step %d: %s", step_num, action)
+            logger.info("Reward: %.2f", reward)
+            logger.info("Done: %s", done)
+            logger.info("Info: %s", info)
 
-        if done:
-            logger.info("The episode is done.")
-            break
+
+            # Save screenshot and trajectory information
+            with open(os.path.join(example_trajectory_dir, f"step_{step_num}_{action_timestamp}.png"), "wb") as _f:
+                with open(observation['screenshot'], "rb") as __f:
+                    screenshot = __f.read()
+                _f.write(screenshot)
+
+            with open(trajectory_recording_path, "a") as f:
+                f.write(json.dumps({
+                    "step_num": step_num,
+                    "action_timestamp": action_timestamp,
+                    "action": action,
+                    "reward": reward,
+                    "done": done,
+                    "info": info,
+                    "screenshot_file": f"step_{step_num}_{action_timestamp}.png"
+                }))
+                f.write("\n")
+
+            if done:
+                logger.info("The episode is done.")
+                break
+
+    if recording:
+        # send a request to the server to stop recording
+        env.controller.end_recording(os.path.join(example_trajectory_dir, "recording.mp4"))
 
     result = env.evaluate()
     logger.info("Result: %.2f", result)
@@ -91,7 +119,7 @@ if __name__ == "__main__":
 
     with open(f"evaluation_examples/examples/{example_class}/{example_id}.json", "r") as f:
         example = json.load(f)
-    example["snapshot"] = "chrome_setup"
+    example["snapshot"] = "exp_setup"
 
     api_key = os.environ.get("OPENAI_API_KEY")
     agent = GPT4v_Agent(api_key=api_key, action_space=action_space)
@@ -101,4 +129,4 @@ if __name__ == "__main__":
     example_trajectory_dir = os.path.join(root_trajectory_dir, example_class, example_id)
     os.makedirs(example_trajectory_dir, exist_ok=True)
 
-    run_one_example(example, agent, 20, example_trajectory_dir)
+    run_one_example(example, agent, 2, example_trajectory_dir)