105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
import datetime
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
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=20, example_trajectory_dir="exp_trajectory"):
|
|
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()
|
|
observation['instruction'] = example['instruction']
|
|
done = False
|
|
step_num = 0
|
|
|
|
# todo: save the screenshots and actions to a folder
|
|
while not done and step_num < max_steps:
|
|
actions = agent.predict(observation)
|
|
for action in actions:
|
|
observation, reward, done, info = env.step(action)
|
|
observation['instruction'] = example['instruction']
|
|
step_num += 1
|
|
logger.info("Step %d", step_num)
|
|
logger.info("Action: %s", actions)
|
|
observation.pop("accessibility_tree")
|
|
logger.info("Observation: %s", observation)
|
|
logger.info("Reward: %.2f", reward)
|
|
logger.info("Info: %s", info)
|
|
|
|
logger.info("================================\n")
|
|
|
|
if done:
|
|
logger.info("The episode is done.")
|
|
break
|
|
|
|
result = env.evaluate()
|
|
logger.info("Result: %.2f", result)
|
|
|
|
# env.close()
|
|
logger.info("Environment closed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
action_space = "pyautogui"
|
|
example_class = "vlc"
|
|
example_id = "8f080098-ddb1-424c-b438-4e96e5e4786e"
|
|
|
|
with open(f"evaluation_examples/examples/{example_class}/{example_id}.json", "r") as f:
|
|
example = json.load(f)
|
|
example["snapshot"] = "chrome_setup"
|
|
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
agent = GPT4v_Agent(api_key=api_key, action_space=action_space)
|
|
|
|
root_trajectory_dir = "exp_trajectory"
|
|
|
|
example_trajectory_dir = os.path.join(root_trajectory_dir, example_class, example_id)
|
|
os.makedirs(example_trajectory_dir, exist_ok=True)
|
|
|
|
run_one_example(example, agent, 20, example_trajectory_dir)
|