245 lines
9.3 KiB
Python
245 lines
9.3 KiB
Python
import ctypes
|
|
import datetime
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import func_timeout
|
|
|
|
from desktop_env.envs.desktop_env import DesktopEnv
|
|
from mm_agents.gpt_4v_agent import GPT4v_Agent
|
|
|
|
# Logger Configs {{{ #
|
|
logger = logging.getLogger()
|
|
logger.setLevel(logging.DEBUG)
|
|
|
|
datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
|
|
|
|
file_handler = logging.FileHandler(os.path.join("logs", "normal-{:}.log".format(datetime_str)), encoding="utf-8")
|
|
debug_handler = logging.FileHandler(os.path.join("logs", "debug-{:}.log".format(datetime_str)), encoding="utf-8")
|
|
stdout_handler = logging.StreamHandler(sys.stdout)
|
|
sdebug_handler = logging.FileHandler(os.path.join("logs", "sdebug-{:}.log".format(datetime_str)), encoding="utf-8")
|
|
|
|
file_handler.setLevel(logging.INFO)
|
|
debug_handler.setLevel(logging.DEBUG)
|
|
stdout_handler.setLevel(logging.INFO)
|
|
sdebug_handler.setLevel(logging.DEBUG)
|
|
|
|
formatter = logging.Formatter(
|
|
fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s")
|
|
file_handler.setFormatter(formatter)
|
|
debug_handler.setFormatter(formatter)
|
|
stdout_handler.setFormatter(formatter)
|
|
sdebug_handler.setFormatter(formatter)
|
|
|
|
stdout_handler.addFilter(logging.Filter("desktopenv"))
|
|
sdebug_handler.addFilter(logging.Filter("desktopenv"))
|
|
|
|
logger.addHandler(file_handler)
|
|
logger.addHandler(debug_handler)
|
|
logger.addHandler(stdout_handler)
|
|
logger.addHandler(sdebug_handler)
|
|
# }}} Logger Configs #
|
|
|
|
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=10, 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,
|
|
action_space=agent.action_space,
|
|
task_config=example
|
|
)
|
|
# reset the environment to certain snapshot
|
|
observation = env.reset()
|
|
done = False
|
|
step_num = 0
|
|
|
|
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)
|
|
step_num += 1
|
|
for action in actions:
|
|
# Capture the timestamp before executing the action
|
|
action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
|
|
logger.info("Step %d: %s", step_num, action)
|
|
|
|
observation, reward, done, info = env.step(action)
|
|
|
|
logger.info("Reward: %.2f", reward)
|
|
logger.info("Done: %s", done)
|
|
logger.info("Info: %s", info)
|
|
|
|
# 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
|
|
|
|
def stop_recording():
|
|
try:
|
|
env.controller.end_recording(os.path.join(example_trajectory_dir, "recording.mp4"))
|
|
except Exception as e:
|
|
print(f"An error occurred while stopping the recording: {e}")
|
|
|
|
try:
|
|
func_timeout.func_timeout(30, stop_recording)
|
|
except func_timeout.exceptions.FunctionTimedOut:
|
|
logger.info("Recording timed out.")
|
|
|
|
result = env.evaluate()
|
|
logger.info("Result: %.2f", result)
|
|
|
|
with open(trajectory_recording_path, "a") as f:
|
|
f.write(json.dumps({
|
|
"result": result
|
|
}))
|
|
f.write("\n")
|
|
|
|
# env.close()
|
|
logger.info("Environment closed.")
|
|
|
|
|
|
def main(example_class, example_id):
|
|
action_space = "pyautogui"
|
|
gpt4_model = "gpt-4-0125-preview"
|
|
gemini_model = "gemini-pro-vision"
|
|
|
|
logger.info("Running example %s/%s", example_class, example_id)
|
|
logger.info("Using model %s", gpt4_model)
|
|
# logger.info("Using model %s", gemini_model)
|
|
|
|
with open(f"evaluation_examples/examples/{example_class}/{example_id}.json", "r", encoding="utf-8") as f:
|
|
example = json.load(f)
|
|
example["snapshot"] = "exp_v1"
|
|
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
agent = GPT4v_Agent(api_key=api_key, model=gpt4_model, instruction=example['instruction'], max_tokens=1000,
|
|
action_space=action_space, exp="a11y_tree")
|
|
|
|
# api_key = os.environ.get("GENAI_API_KEY")
|
|
# agent = GeminiPro_Agent(api_key=api_key, model=gemini_model, instruction=example['instruction'], action_space=action_space, exp="a11y_tree")
|
|
|
|
root_trajectory_dir = "exp_trajectory"
|
|
|
|
example_trajectory_dir = os.path.join(root_trajectory_dir, "a11y_tree", example_class, gpt4_model, example_id)
|
|
# example_trajectory_dir = os.path.join(root_trajectory_dir, "a11y_tree", example_class, gemini_model, example_id)
|
|
|
|
os.makedirs(example_trajectory_dir, exist_ok=True)
|
|
|
|
run_one_example(example, agent, 15, example_trajectory_dir)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
vlc_list = [
|
|
# "8ba5ae7a-5ae5-4eab-9fcc-5dd4fe3abf89",
|
|
# "8ba5ae7a-5ae5-4eab-9fcc-5dd4fe3abf89",
|
|
# "8f080098-ddb1-424c-b438-4e96e5e4786e",
|
|
# "bba3381f-b5eb-4439-bd9e-80c22218d5a7",
|
|
# "fba2c100-79e8-42df-ae74-b592418d54f4",
|
|
# "efcf0d81-0835-4880-b2fd-d866e8bc2294",
|
|
# "8d9fd4e2-6fdb-46b0-b9b9-02f06495c62f",
|
|
# "aa4b5023-aef6-4ed9-bdc9-705f59ab9ad6",
|
|
# "386dbd0e-0241-4a0a-b6a2-6704fba26b1c",
|
|
# "9195653c-f4aa-453d-aa95-787f6ccfaae9",
|
|
# "d06f0d4d-2cd5-4ede-8de9-598629438c6e",
|
|
# "a5bbbcd5-b398-4c91-83d4-55e1e31bbb81",
|
|
"f3977615-2b45-4ac5-8bba-80c17dbe2a37",
|
|
"215dfd39-f493-4bc3-a027-8a97d72c61bf"
|
|
]
|
|
|
|
chrome_list = [
|
|
# "bb5e4c0d-f964-439c-97b6-bdb9747de3f4",
|
|
"7b6c7e24-c58a-49fc-a5bb-d57b80e5b4c3",
|
|
"06fe7178-4491-4589-810f-2e2bc9502122",
|
|
"e1e75309-3ddb-4d09-92ec-de869c928143",
|
|
"35253b65-1c19-4304-8aa4-6884b8218fc0",
|
|
"2ad9387a-65d8-4e33-ad5b-7580065a27ca",
|
|
"7a5a7856-f1b6-42a4-ade9-1ca81ca0f263",
|
|
"44ee5668-ecd5-4366-a6ce-c1c9b8d4e938",
|
|
"2ae9ba84-3a0d-4d4c-8338-3a1478dc5fe3",
|
|
"480bcfea-d68f-4aaa-a0a9-2589ef319381",
|
|
"af630914-714e-4a24-a7bb-f9af687d3b91"
|
|
]
|
|
|
|
calc_list = [
|
|
"eb03d19a-b88d-4de4-8a64-ca0ac66f426b",
|
|
"0bf05a7d-b28b-44d2-955a-50b41e24012a",
|
|
"7b802dad-6e0f-4204-9815-d4e3f57627d8",
|
|
"7a4e4bc8-922c-4c84-865c-25ba34136be1",
|
|
"2bd59342-0664-4ccb-ba87-79379096cc08",
|
|
"a9f325aa-8c05-4e4f-8341-9e4358565f4f",
|
|
"ecb0df7a-4e8d-4a03-b162-053391d3afaf",
|
|
"7efeb4b1-3d19-4762-b163-63328d66303b",
|
|
"4e6fcf72-daf3-439f-a232-c434ce416af6",
|
|
"6054afcb-5bab-4702-90a0-b259b5d3217c",
|
|
"abed40dc-063f-4598-8ba5-9fe749c0615d",
|
|
"01b269ae-2111-4a07-81fd-3fcd711993b0",
|
|
"8b1ce5f2-59d2-4dcc-b0b0-666a714b9a14",
|
|
"af2b02f7-acee-4be4-8b66-499fab394915",
|
|
"da1d63b8-fa12-417b-ba18-f748e5f770f3",
|
|
"636380ea-d5f6-4474-b6ca-b2ed578a20f1",
|
|
"5ba77536-05c5-4aae-a9ff-6e298d094c3e",
|
|
"4bc4eaf4-ca5e-4db2-8138-8d4e65af7c0b",
|
|
"672a1b02-c62f-4ae2-acf0-37f5fb3052b0",
|
|
"648fe544-16ba-44af-a587-12ccbe280ea6",
|
|
"8985d1e4-5b99-4711-add4-88949ebb2308",
|
|
"9e606842-2e27-43bf-b1d1-b43289c9589b",
|
|
"fcb6e45b-25c4-4087-9483-03d714f473a9",
|
|
"68c0c5b7-96f3-4e87-92a7-6c1b967fd2d2",
|
|
"fff629ea-046e-4793-8eec-1a5a15c3eb35",
|
|
"5c9a206c-bb00-4fb6-bb46-ee675c187df5",
|
|
"e975ae74-79bd-4672-8d1c-dc841a85781d",
|
|
"34a6938a-58da-4897-8639-9b90d6db5391",
|
|
"b5a22759-b4eb-4bf2-aeed-ad14e8615f19",
|
|
"2f9913a1-51ed-4db6-bfe0-7e1c95b3139e",
|
|
"2558031e-401d-4579-8e00-3ecf540fb492",
|
|
"39aa4e37-dc91-482e-99af-132a612d40f3",
|
|
"0cecd4f3-74de-457b-ba94-29ad6b5dafb6",
|
|
"4188d3a4-077d-46b7-9c86-23e1a036f6c1",
|
|
"51b11269-2ca8-4b2a-9163-f21758420e78",
|
|
"7e429b8d-a3f0-4ed0-9b58-08957d00b127",
|
|
"f5a90742-3fa2-40fc-a564-f29b054e0337",
|
|
"22df9241-f8d7-4509-b7f1-37e501a823f7",
|
|
"1434ca3e-f9e3-4db8-9ca7-b4c653be7d17",
|
|
"347ef137-7eeb-4c80-a3bb-0951f26a8aff",
|
|
"6e99a1ad-07d2-4b66-a1ce-ece6d99c20a5",
|
|
"3aaa4e37-dc91-482e-99af-132a612d40f3",
|
|
"37608790-6147-45d0-9f20-1137bb35703d",
|
|
"f9584479-3d0d-4c79-affa-9ad7afdd8850",
|
|
"d681960f-7bc3-4286-9913-a8812ba3261a",
|
|
"21df9241-f8d7-4509-b7f1-37e501a823f7",
|
|
"f6a90742-3fa2-40fc-a564-f29b054e0337",
|
|
"1334ca3e-f9e3-4db8-9ca7-b4c653be7d17",
|
|
"357ef137-7eeb-4c80-a3bb-0951f26a8aff",
|
|
"aa3a8974-2e85-438b-b29e-a64df44deb4b",
|
|
"a01fbce3-2793-461f-ab86-43680ccbae25",
|
|
"4f07fbe9-70de-4927-a4d5-bb28bc12c52c",
|
|
]
|
|
|
|
for example_id in calc_list:
|
|
main("libreoffice_calc", example_id)
|