Fix small bugs in max time limit setting

This commit is contained in:
Timothyxxx
2024-03-16 14:34:40 +08:00
parent 5a062c423f
commit 639f8c7db8
3 changed files with 47 additions and 13 deletions

View File

@@ -1,8 +1,10 @@
import os
import datetime
import json
import logging
import os
from wrapt_timeout_decorator import *
logger = logging.getLogger("desktopenv.experiment")
# Open the JSON file
@@ -11,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()
@@ -37,7 +40,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:
"wb") as _f:
with open(obs['screenshot'], "rb") as __f:
screenshot = __f.read()
_f.write(screenshot)