diff --git a/lib_run_single.py b/lib_run_single.py index e0e342e..ce9b215 100644 --- a/lib_run_single.py +++ b/lib_run_single.py @@ -461,3 +461,83 @@ def run_single_example_uipath(agent, env, example, max_steps, instruction, args, 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")) + +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_os_symphony(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") diff --git a/run_multienv_os_symphony.py b/run_multienv_os_symphony.py index 2c3f1f3..eb2802f 100644 --- a/run_multienv_os_symphony.py +++ b/run_multienv_os_symphony.py @@ -12,7 +12,7 @@ from multiprocessing import Process, Manager, current_process, Queue from mm_agents.os_symphony.agents.os_symphony import OSSymphony from mm_agents.os_symphony.agents.os_aci import OSACI -import lib_run_single_os_symphony as lib_run_single +import lib_run_single # Modify desktop_env, add a new function 'start' from desktop_env.desktop_env_os_symphony import DesktopEnv as OSWorldDesktopEnv @@ -230,7 +230,7 @@ def run_env_tasks( logger.info(f"[{current_process().name}][Example ID]: {example_id}") logger.info(f"[{current_process().name}][Instruction]: {instruction}") try: - lib_run_single.run_single_example( + lib_run_single.run_single_example_os_symphony( agent, env, example,