85 lines
3.3 KiB
Python
Executable File
85 lines
3.3 KiB
Python
Executable File
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from wrapt_timeout_decorator import *
|
|
from mm_agents.os_symphony.utils.common_utils import draw_coordinates
|
|
from mm_agents.os_symphony.utils.process_context import set_current_result_dir
|
|
|
|
|
|
logger = logging.getLogger("desktopenv.experiment")
|
|
|
|
def run_single_example(agent, env, example, max_steps, instruction, args, example_result_dir, scores):
|
|
set_current_result_dir(example_result_dir)
|
|
|
|
agent.reset(result_dir=example_result_dir)
|
|
env.reset(task_config=example)
|
|
time.sleep(30) # Wait for the environment to be ready
|
|
obs = env._get_obs() # Get the initial observation
|
|
done = False
|
|
step_idx = 0
|
|
# env.controller.start_recording()
|
|
start_time = time.time()
|
|
|
|
while not done and step_idx < max_steps:
|
|
response, actions = agent.predict(
|
|
instruction,
|
|
obs,
|
|
step_idx == max_steps - 1
|
|
)
|
|
for action in actions:
|
|
# Save screenshot and trajectory information
|
|
if "reflection" in response and response["reflection"].get("is_milestone"):
|
|
img_name = f"step_{step_idx + 1}_milestone.png"
|
|
else:
|
|
img_name = f"step_{step_idx + 1}.png"
|
|
|
|
with open(os.path.join(example_result_dir, img_name),
|
|
"wb") as _f:
|
|
_f.write(obs['screenshot'])
|
|
if "coordinates" in response and response["coordinates"]:
|
|
draw_coordinates(
|
|
image_bytes=obs['screenshot'],
|
|
coordinates=response["coordinates"],
|
|
save_path=os.path.join(example_result_dir, img_name[:-4] + "_draw.png")
|
|
)
|
|
|
|
logger.info("Step %d: %s", step_idx + 1, action)
|
|
obs, reward, done, info = env.step(action, args.sleep_after_execution)
|
|
logger.info("Done: %s", done)
|
|
|
|
with open(os.path.join(example_result_dir, "traj.jsonl"), "a", encoding="utf-8") as f:
|
|
f.write(json.dumps({
|
|
"instruction": instruction,
|
|
"step_num": step_idx + 1,
|
|
"action": action,
|
|
"response": response,
|
|
"done": done,
|
|
"info": info,
|
|
"screenshot_file": img_name
|
|
}))
|
|
f.write("\n")
|
|
with open(os.path.join(example_result_dir, f"traj_{step_idx+1}.json"), "w", encoding="utf-8") as f:
|
|
json.dump({
|
|
"step_num": step_idx + 1,
|
|
"action": action,
|
|
"response": response,
|
|
"done": done,
|
|
"info": info,
|
|
"screenshot_file": img_name
|
|
}, f, indent=4, ensure_ascii=False)
|
|
if done:
|
|
logger.info("The episode is done.")
|
|
time.sleep(60)
|
|
break
|
|
step_idx += 1
|
|
end_time = time.time()
|
|
result = float(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")
|
|
|
|
with open(os.path.join(example_result_dir, "time.txt"), "w", encoding="utf-8") as f:
|
|
f.write(f"{end_time-start_time:.2f}\n")
|