Clean code; Refactor environment to pass screenshot content instead of path

lib_run_single.py

@@ -2,7 +2,6 @@ import datetime
 import json
 import logging
 import os
-# import wandb
 
 from wrapt_timeout_decorator import *
 
@@ -14,6 +13,7 @@ with open("./settings.json", "r") as file:
     data = json.load(file)
 time_limit = data["time_limit"]
 
+
 @timeout(time_limit, use_signals=False)
 def run_single_example(agent, env, example, max_steps, instruction, args, example_result_dir, scores):
     agent.reset()
@@ -21,7 +21,6 @@ def run_single_example(agent, env, example, max_steps, instruction, args, exampl
     done = False
     step_idx = 0
     env.controller.start_recording()
-    # str_table = wandb.Table(columns=["Screenshot", "A11T", "Modle Response", "Action", "Action timestamp", "Done"])
     while not done and step_idx < max_steps:
         response, actions = agent.predict(
             instruction,
@@ -38,15 +37,7 @@ def run_single_example(agent, env, example, max_steps, instruction, args, exampl
             # Save screenshot and trajectory information
             with open(os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"),
                       "wb") as _f:
-                with open(obs['screenshot'], "rb") as __f:
-                    screenshot = __f.read()
-                _f.write(screenshot)
-            # get a11tree and save to wandb
-            # thisrun_a11tree = env.controller.get_accessibility_tree()
-            # str_table.add_data(wandb.Image(data_or_path=os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"), caption=f"step_{step_idx + 1}_{action_timestamp}"),
-            #                 thisrun_a11tree,
-            #                 response, action, action_timestamp, done)
-            # run.log({"Reward": reward})
+                _f.write(obs['screenshot'])
             with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                 f.write(json.dumps({
                     "step_num": step_idx + 1,
@@ -62,11 +53,9 @@ def run_single_example(agent, env, example, max_steps, instruction, args, exampl
                 logger.info("The episode is done.")
                 break
         step_idx += 1
-    # run.log({"str_trajectory": str_table})
     result = env.evaluate()
     logger.info("Result: %.2f", result)
     scores.append(result)
     with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
         f.write(f"{result}\n")
     env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))
-    # run.log({"Result": result})