181 lines
6.7 KiB
Python
181 lines
6.7 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"
|
|
|
|
|
|
PATH_TO_VM = "../../../../大文件/镜像/Ubuntu-1218/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-vision-preview"
|
|
#gemini_model = "gemini-pro-vision"
|
|
|
|
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"
|
|
# example["snapshot"] = "exp_setup4"
|
|
example["snapshot"] = "Snapshot 30"
|
|
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
agent = GPT4v_Agent(api_key=api_key, model=gpt4_model, max_tokens=1000, instruction=example['instruction'],
|
|
action_space=action_space, exp="som")
|
|
|
|
# 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)
|
|
|
|
root_trajectory_dir = "exp_trajectory"
|
|
|
|
example_trajectory_dir = os.path.join(root_trajectory_dir, "som", example_class, gpt4_model, example_id)
|
|
# example_trajectory_dir = os.path.join(root_trajectory_dir, "som", 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__':
|
|
xx_list = [ "94d95f96-9699-4208-98ba-3c3119edf9c2"
|
|
, "bedcedc4-4d72-425e-ad62-21960b11fe0d"
|
|
, "43c2d64c-bab5-4dcb-a30c-b888321c319a"
|
|
, "7688b85f-87a4-4e4a-b2f8-f3d6c3f29b82"
|
|
, "ec4e3f68-9ea4-4c18-a5c9-69f89d1178b3"
|
|
, "a462a795-fdc7-4b23-b689-e8b6df786b78"
|
|
, "f9be0997-4b7c-45c5-b05c-4612b44a6118"
|
|
, "ae039631-2b12-4637-84f6-c67d51511be3"
|
|
, "e2eb4bf1-aa93-4192-b55d-03e2fb6dfd15"
|
|
, "28cc3b7e-b194-4bc9-8353-d04c0f4d56d2"
|
|
, "5ea617a3-0e86-4ba6-aab2-dac9aa2e8d57"
|
|
, "e0df059f-28a6-4169-924f-b9623e7184cc"
|
|
, "ddc75b62-7311-4af8-bfb3-859558542b36"
|
|
, "5c433d22-ed9a-4e31-91f5-54cf3e8acd63"
|
|
, "b6781586-6346-41cd-935a-a6b1487918fc"
|
|
, "b3d4a89c-53f2-4d6b-8b6a-541fb5d205fa"
|
|
, "3ce045a0-877b-42aa-8d2c-b4a863336ab8"
|
|
, "fe41f596-a71b-4c2f-9b2f-9dcd40b568c3"
|
|
, "a4d98375-215b-4a4d-aee9-3d4370fccc41"
|
|
, "765d2b74-88a7-4d50-bf51-34e4106fd24a"
|
|
, "13584542-872b-42d8-b299-866967b5c3ef"
|
|
, "23393935-50c7-4a86-aeea-2b78fd089c5c"
|
|
]
|
|
for example_id in xx_list:
|
|
main("os", example_id)
|