Merge branch 'main' into zdy

This commit is contained in:
David Chang
2023-12-25 20:29:20 +08:00
80 changed files with 835 additions and 90 deletions

View File

@@ -3,23 +3,24 @@ from __future__ import annotations
import os
import subprocess
import time
#import uuid
#import platform
# import uuid
# import platform
from typing import List, Dict
from typing import Callable, Any, Optional
import tempfile
import gymnasium as gym
#import requests
# import requests
from desktop_env.controllers.python import PythonController
from desktop_env.controllers.setup import SetupController
#from desktop_env.evaluators import eval_funcs
# from desktop_env.evaluators import eval_funcs
from desktop_env.evaluators import metrics, getters
Metric = Callable[[Any, Any], float]
Getter = Callable[[gym.Env, Dict[str, Any]], Any]
def _execute_command(command: List[str]) -> None:
if command[:4] == ["vmrun", "-T", "ws", "start"]:
p = subprocess.Popen(command)
@@ -104,7 +105,7 @@ class DesktopEnv(gym.Env):
output = subprocess.check_output("vmrun -T ws list", shell=True, stderr=subprocess.STDOUT)
output = output.decode()
output: List[str] = output.splitlines()
#if self.path_to_vm.lstrip("~/") in output:
# if self.path_to_vm.lstrip("~/") in output:
if self.path_to_vm in output:
print("VM is running.")
break
@@ -132,9 +133,9 @@ class DesktopEnv(gym.Env):
_execute_command(["vmrun", "-T", "ws" "snapshot", self.path_to_vm, self.snapshot_path])
def _get_screenshot(self):
#random_uuid = str(uuid.uuid4())
#os.makedirs(os.path.join("tmp", random_uuid), exist_ok=True)
#image_path = os.path.join("tmp", random_uuid, "screenshot.png")
# random_uuid = str(uuid.uuid4())
# os.makedirs(os.path.join("tmp", random_uuid), exist_ok=True)
# image_path = os.path.join("tmp", random_uuid, "screenshot.png")
image_path: str = os.path.join(self.tmp_dir, "screenshots", "{:d}.png".format(self._step_no))
# Get the screenshot and save to the image_path
@@ -173,9 +174,10 @@ class DesktopEnv(gym.Env):
self.action_history.clear()
print("Setup new temp dir...")
self.tmp_dir = tempfile.mkdtemp( prefix="{:d}.{:}.".format(self._traj_no, self.task_id)
, dir=self.tmp_dir_base
)
self.tmp_dir = tempfile.mkdtemp(
prefix="{:d}.{:}.".format(self._traj_no, self.task_id),
dir=self.tmp_dir_base
)
os.makedirs(os.path.join(self.tmp_dir, "screenshots"))
print("Reverting to snapshot to {}...".format(self.snapshot_path))
@@ -224,14 +226,14 @@ class DesktopEnv(gym.Env):
"""
# todo: make this more flexible by refactoring
#eval_func = eval_funcs[self.evaluator["func"]]
#eval_func_vars = {}
#
#for var_name, file_info in self.evaluator["paths"].items():
#path = copy_file_to_local(file_info)
#eval_func_vars[var_name] = path
#
#return eval_func(**eval_func_vars)
# eval_func = eval_funcs[self.evaluator["func"]]
# eval_func_vars = {}
#
# for var_name, file_info in self.evaluator["paths"].items():
# path = copy_file_to_local(file_info)
# eval_func_vars[var_name] = path
#
# return eval_func(**eval_func_vars)
result = self.result_getter(self, self.evaluator["result"])
expected = self.expected_getter(self, self.evaluator["expected"])