diff --git a/lib_run_single.py b/lib_run_single.py index 0d3fefb..4bc37e1 100644 --- a/lib_run_single.py +++ b/lib_run_single.py @@ -267,3 +267,63 @@ def run_single_example_autoglm(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")) + +def run_single_example_mano(agent, env, example, max_steps, instruction, args, example_result_dir, scores): + runtime_logger = setup_logger(example, example_result_dir) + agent.reset(runtime_logger) + env.reset(task_config=example) + time.sleep(60) # Wait for the environment to be ready + obs = env._get_obs() # Get the initial observation + done = False + step_idx = 0 + env.controller.start_recording() + + with open(os.path.join(example_result_dir, f"step_0.png"), + "wb") as _f: + _f.write(obs['screenshot']) + while not done and step_idx < max_steps: + response, actions = agent.predict( + instruction, + obs + ) + if len(actions) > 1: + if (("pyautogui.hotkey('shift')" in actions[0] or "pyautogui.hotkey('ctrl')" in actions[0]) + and "pyautogui.click" in actions[1]): + hotkey_type = 'shift' if "shift" in actions[0] else 'ctrl' + action = f"pyautogui.keyDown('{hotkey_type}')\n{actions[1]}\npyautogui.keyUp('{hotkey_type}')" + actions = [action] + + 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_idx + 1, action) + obs, reward, done, info = env.step(action, args.sleep_after_execution) + + logger.info("Reward: %.2f", reward) + logger.info("Done: %s", done) + # Save screenshot and trajectory information + with open(os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"), + "wb") as _f: + _f.write(obs['screenshot']) + with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f: + f.write(json.dumps({ + "step_num": step_idx + 1, + "action_timestamp": action_timestamp, + "action": action, + "reward": reward, + "done": done, + "info": info, + "screenshot_file": f"step_{step_idx + 1}_{action_timestamp}.png", + "response":response + })) + f.write("\n") + if done: + logger.info("The episode is done.") + break + step_idx += 1 + result = 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") + env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4")) \ No newline at end of file diff --git a/mm_agents/mano_agent.py b/mm_agents/mano_agent.py new file mode 100644 index 0000000..d85a6a8 --- /dev/null +++ b/mm_agents/mano_agent.py @@ -0,0 +1,1068 @@ +import ast +import base64 +import logging +import math +import re +import os +import xml.etree.ElementTree as ET +from io import BytesIO +from typing import Dict, List + +import backoff +import numpy as np +from PIL import Image, ImageDraw +from requests.exceptions import SSLError +import openai +from openai import OpenAI +from google.api_core.exceptions import ( + BadRequest, + InternalServerError, + InvalidArgument, + ResourceExhausted, +) + +from mm_agents.accessibility_tree_wrap.heuristic_retrieve import ( + filter_nodes, +) + +# from mm_agents.prompts import MANO_PROMPT_THOUGHT + + +logger = logging.getLogger("desktopenv.agent") + +FINISH_WORD = "finished" +WAIT_WORD = "wait" +ENV_FAIL_WORD = "error_env" +CALL_USER = "call_user" + +IMAGE_FACTOR = 28 +MIN_PIXELS = 100 * 28 * 28 +MAX_PIXELS = 16384 * 28 * 28 +MAX_RATIO = 200 +MANO_PROMPT_THOUGHT = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. + +## Output Format +``` +Thought: ... +Action desp: ... +Action: ... +``` + +## Action Space +click(start_box='<|box_start|>(x1,y1)<|box_end|>') +left_double(start_box='<|box_start|>(x1,y1)<|box_end|>') +right_single(start_box='<|box_start|>(x1,y1)<|box_end|>') +drag(start_box='<|box_start|>(x1,y1)<|box_end|>', end_box='<|box_start|>(x3,y3)<|box_end|>') +hotkey(key='') +type(content='') #If you want to submit your input, use "\\n" at the end of `content`. +scroll(start_box='<|box_start|>(x1,y1)<|box_end|>', direction='down or up or right or left') +wait() #Sleep for 5s and take a screenshot to check for any changes. +finished(content='xxx') # Use escape characters \\', \\", and \\n in content part to ensure we can parse the content in normal python string format. + +## Note +- Use {language} in `Thought` part. +- Write a small plan and finally summarize your next action (with its target element) in one sentence in `Action desp` part. + +## User Instruction +{instruction} +""" + + +pure_text_settings = ["a11y_tree"] + +attributes_ns_ubuntu = "https://accessibility.windows.example.org/ns/attributes" +attributes_ns_windows = "https://accessibility.windows.example.org/ns/attributes" +state_ns_ubuntu = "https://accessibility.ubuntu.example.org/ns/state" +state_ns_windows = "https://accessibility.windows.example.org/ns/state" +component_ns_ubuntu = "https://accessibility.ubuntu.example.org/ns/component" +component_ns_windows = "https://accessibility.windows.example.org/ns/component" +value_ns_ubuntu = "https://accessibility.ubuntu.example.org/ns/value" +value_ns_windows = "https://accessibility.windows.example.org/ns/value" +class_ns_windows = "https://accessibility.windows.example.org/ns/class" +# More namespaces defined in OSWorld, please check desktop_env/server/main.py + +# 定义一个函数来解析每个 action +def parse_action(action_str): + try: + # 解析字符串为 AST 节点 + node = ast.parse(action_str, mode='eval') + + # 确保节点是一个表达式 + if not isinstance(node, ast.Expression): + raise ValueError("Not an expression") + + # 获取表达式的主体 + call = node.body + + # 确保主体是一个函数调用 + if not isinstance(call, ast.Call): + raise ValueError("Not a function call") + + # 获取函数名 + if isinstance(call.func, ast.Name): + func_name = call.func.id + elif isinstance(call.func, ast.Attribute): + func_name = call.func.attr + else: + func_name = None + + # 获取关键字参数 + kwargs = {} + for kw in call.keywords: + key = kw.arg + # 处理不同类型的值,这里假设都是常量 + if isinstance(kw.value, ast.Constant): + value = kw.value.value + elif isinstance(kw.value, ast.Str): # 兼容旧版本 Python + value = kw.value.s + else: + value = None + value = str(value) + kwargs[key] = value + + return { + 'function': func_name, + 'args': kwargs + } + + except Exception as e: + print(f"Failed to parse action '{action_str}': {e}") + return None + +def escape_single_quotes(text): + # 匹配未转义的单引号(不匹配 \\') + pattern = r"(? int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor + +def linear_resize( + height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS +) -> tuple[int, int]: + if width * height > max_pixels: + """ + 如果图片超过/低于像素限制,则计算一个缩放因子resize_factor,使图片的像素数缩小到等于或小于max_pixels。这个缩放因子是通过开平方根计算的,确保纵横比保持不变,这样原始的相对坐标可以不经转换直接复用 + """ + resize_factor = math.sqrt(max_pixels / (width * height)) + width, height = int(width * resize_factor), int(height * resize_factor) + if width * height < min_pixels: + resize_factor = math.sqrt(min_pixels / (width * height)) + width, height = math.ceil(width * resize_factor), math.ceil(height * resize_factor) + + return height, width + +def smart_resize( + height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS +) -> tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + + 3. The aspect ratio of the image is maintained as closely as possible. + """ + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + +def parse_action_to_structure_output(text, factor, origin_resized_height, origin_resized_width, model_type, max_pixels=16384*28*28, min_pixels=100*28*28): + text = text.strip() + if model_type == "qwen25vl": + smart_resize_height, smart_resize_width = smart_resize(origin_resized_height, origin_resized_width, factor=IMAGE_FACTOR, min_pixels=min_pixels, max_pixels=max_pixels) + + # 正则表达式匹配 Action 字符串 + if text.startswith("Thought:") and "Action desp" not in text: + thought_pattern = r"Thought: (.+?)(?=\s*Action:|$)" + thought_hint = "Thought: " + elif text.startswith("Thought:") and "Action desp" in text: + thought_pattern = r"Thought: (.+?)(?=\s*Action desp:|$)" + thought_hint = "Thought: " + elif text.startswith("Reflection:"): + thought_pattern = r"Reflection: (.+?)Action_Summary: (.+?)(?=\s*Action:|$)" + thought_hint = "Reflection: " + elif text.startswith("Action_Summary:"): + thought_pattern = r"Action_Summary: (.+?)(?=\s*Action:|$)" + thought_hint = "Action_Summary: " + else: + thought_pattern = r"Thought: (.+?)(?=\s*Action:|$)" + thought_hint = "Thought: " + reflection, thought = None, None + thought_match = re.search(thought_pattern, text, re.DOTALL) + if thought_match: + if len(thought_match.groups()) == 1: + thought = thought_match.group(1).strip() + elif len(thought_match.groups()) == 2: + thought = thought_match.group(2).strip() + reflection = thought_match.group(1).strip() + assert "Action:" in text + action_str = text.split("Action:")[-1] + + tmp_all_action = action_str.split("\n\n") + all_action = [] + for action_str in tmp_all_action: + if "type(content" in action_str: + # 正则表达式匹配 content 中的字符串并转义单引号 + def escape_quotes(match): + content = match.group(1) # 获取 content 的值 + return content + + # 使用正则表达式进行替换 + pattern = r"type\(content='(.*?)'\)" # 匹配 type(content='...') + content = re.sub(pattern, escape_quotes, action_str) + + # 处理字符串 + action_str = escape_single_quotes(content) + action_str = "type(content='" + action_str + "')" + all_action.append(action_str) + + parsed_actions = [parse_action(action.replace("\n","\\n").lstrip()) for action in all_action] + actions = [] + for action_instance, raw_str in zip(parsed_actions, all_action): + if action_instance == None: + print(f"Action can't parse: {raw_str}") + raise ValueError(f"Action can't parse: {raw_str}") + action_type = action_instance["function"] + params = action_instance["args"] + + # import pdb; pdb.set_trace() + action_inputs = {} + for param_name, param in params.items(): + if param == "": continue + param = param.lstrip() # 去掉引号和多余的空格 + # 处理start_box或者end_box参数格式 'x1 y1 x2 y2' + action_inputs[param_name.strip()] = param + + if "start_box" in param_name or "end_box" in param_name: + ori_box = param + # Remove parentheses and split the string by commas + numbers = ori_box.replace("(", "").replace(")", "").split(",") + + # Convert to float and scale by 1000 + # Qwen2.5vl output absolute coordinates, qwen2vl output relative coordinates + if model_type == "qwen25vl": + float_numbers = [] + for num_idx, num in enumerate(numbers): + num = float(num) + if (num_idx + 1) % 2 == 0: + float_numbers.append(float(num/smart_resize_height)) + else: + float_numbers.append(float(num/smart_resize_width)) + else: + float_numbers = [float(num) / factor for num in numbers] + + if len(float_numbers) == 2: + float_numbers = [float_numbers[0], float_numbers[1], float_numbers[0], float_numbers[1]] + action_inputs[param_name.strip()] = str(float_numbers) + + # import pdb; pdb.set_trace() + actions.append({ + "reflection": reflection, + "thought": thought, + "action_type": action_type, + "action_inputs": action_inputs, + "text": text + }) + return actions + +def parsing_response_to_pyautogui_code(responses, image_height: int, image_width:int, input_swap:bool=True) -> str: + ''' + 将M模型的输出解析为OSWorld中的action,生成pyautogui代码字符串 + 参数: + response: 包含模型输出的字典,结构类似于: + { + "action_type": "hotkey", + "action_inputs": { + "hotkey": "v ctrl", + "start_box": None, + "end_box": None + } + } + 返回: + 生成的pyautogui代码字符串 + ''' + + pyautogui_code = f"import pyautogui\nimport time\n" + if isinstance(responses, dict): + responses = [responses] + for response_id, response in enumerate(responses): + if "observation" in response: + observation = response["observation"] + else: + observation = "" + + if "thought" in response: + thought = response["thought"] + else: + thought = "" + + if response_id == 0: + pyautogui_code += f"'''\nObservation:\n{observation}\n\nThought:\n{thought}\n'''\n" + else: + pyautogui_code += f"\ntime.sleep(1)\n" + + action_dict = response + action_type = action_dict.get("action_type") + action_inputs = action_dict.get("action_inputs", {}) + + if action_type == "hotkey": + # Parsing hotkey action + if "key" in action_inputs: + hotkey = action_inputs.get("key", "") + else: + hotkey = action_inputs.get("hotkey", "") + + if hotkey == "arrowleft": + hotkey = "left" + + elif hotkey == "arrowright": + hotkey = "right" + + elif hotkey == "arrowup": + hotkey = "up" + + elif hotkey == "arrowdown": + hotkey = "down" + + if hotkey: + # Handle other hotkeys + keys = hotkey.split() # Split the keys by space + convert_keys = [] + for key in keys: + if key == "space": + key = ' ' + convert_keys.append(key) + pyautogui_code += f"\npyautogui.hotkey({', '.join([repr(k) for k in convert_keys])})" + + elif action_type == "press": + # Parsing press action + if "key" in action_inputs: + key_to_press = action_inputs.get("key", "") + else: + key_to_press = action_inputs.get("press", "") + + if hotkey == "arrowleft": + hotkey = "left" + + elif hotkey == "arrowright": + hotkey = "right" + + elif hotkey == "arrowup": + hotkey = "up" + + elif hotkey == "arrowdown": + hotkey = "down" + + elif hotkey == "space": + hotkey = " " + + if key_to_press: + # Simulate pressing a single key + pyautogui_code += f"\npyautogui.press({repr(key_to_press)})" + + elif action_type == "keyup": + key_to_up = action_inputs.get("key", "") + pyautogui_code += f"\npyautogui.keyUp({repr(key_to_up)})" + + elif action_type == "keydown": + key_to_down = action_inputs.get("key", "") + pyautogui_code += f"\npyautogui.keyDown({repr(key_to_down)})" + + elif action_type == "type": + # Parsing typing action using clipboard + content = action_inputs.get("content", "") + content = escape_single_quotes(content) + stripped_content = content + if content.endswith("\n") or content.endswith("\\n"): + stripped_content = stripped_content.rstrip("\\n").rstrip("\n") + if content: + if input_swap: + pyautogui_code += f"\nimport pyperclip" + pyautogui_code += f"\npyperclip.copy('{stripped_content}')" + pyautogui_code += f"\npyautogui.hotkey('ctrl', 'v')" + pyautogui_code += f"\ntime.sleep(0.5)\n" + if content.endswith("\n") or content.endswith("\\n"): + pyautogui_code += f"\npyautogui.press('enter')" + else: + pyautogui_code += f"\npyautogui.write('{stripped_content}', interval=0.1)" + pyautogui_code += f"\ntime.sleep(0.5)\n" + if content.endswith("\n") or content.endswith("\\n"): + pyautogui_code += f"\npyautogui.press('enter')" + + + elif action_type in ["drag", "select"]: + # Parsing drag or select action based on start and end_boxes + start_box = action_inputs.get("start_box") + end_box = action_inputs.get("end_box") + if start_box and end_box: + x1, y1, x2, y2 = eval(start_box) # Assuming box is in [x1, y1, x2, y2] + sx = round(float((x1 + x2) / 2) * image_width, 3) + sy = round(float((y1 + y2) / 2) * image_height, 3) + x1, y1, x2, y2 = eval(end_box) # Assuming box is in [x1, y1, x2, y2] + ex = round(float((x1 + x2) / 2) * image_width, 3) + ey = round(float((y1 + y2) / 2) * image_height, 3) + pyautogui_code += ( + f"\npyautogui.moveTo({sx}, {sy})\n" + f"\npyautogui.dragTo({ex}, {ey}, duration=1.0)\n" + ) + + + elif action_type == "scroll": + # Parsing scroll action + start_box = action_inputs.get("start_box") + if start_box: + x1, y1, x2, y2 = eval(start_box) # Assuming box is in [x1, y1, x2, y2] + x = round(float((x1 + x2) / 2) * image_width, 3) + y = round(float((y1 + y2) / 2) * image_height, 3) + else: + x = None + y = None + + direction = action_inputs.get("direction", "") + + # 获取自定义滑动距离,如果没有则使用默认值 + scroll_amount = action_inputs.get("scroll_amount", None) + + if scroll_amount is not None: + try: + scroll_amount = int(scroll_amount) + except (ValueError, TypeError): + scroll_amount = None # 转换失败使用默认值 + + # 确定滚动值 + if scroll_amount is not None: + # 使用自定义滑动距离 + if "up" in direction.lower(): + scroll_value = abs(scroll_amount) # 向上滚动:正值 + elif "down" in direction.lower(): + scroll_value = -abs(scroll_amount) # 向下滚动:负值 + else: + scroll_value = -abs(scroll_amount) # 默认向下 + else: + # 使用默认滑动距离 + if "up" in direction.lower(): + scroll_value = 5 # 向上滚动:正值 + elif "down" in direction.lower(): + scroll_value = -5 # 向下滚动:负值 + else: + scroll_value = -5 # 默认向下 + + # 生成滚动代码 + if x is None: + pyautogui_code += f"\npyautogui.scroll({scroll_value})" + else: + pyautogui_code += f"\npyautogui.scroll({scroll_value}, x={x}, y={y})" + + elif action_type in ["click", "left_single", "left_double", "right_single", "hover"]: + # Parsing mouse click actions + start_box = action_inputs.get("start_box") + start_box = str(start_box) + if start_box: + start_box = eval(start_box) + if len(start_box) == 4: + x1, y1, x2, y2 = start_box # Assuming box is in [x1, y1, x2, y2] + elif len(start_box) == 2: + x1, y1 = start_box + x2 = x1 + y2 = y1 + x = round(float((x1 + x2) / 2) * image_width, 3) + y = round(float((y1 + y2) / 2) * image_height, 3) + if action_type == "left_single" or action_type == "click": + pyautogui_code += f"\npyautogui.click({x}, {y}, button='left')" + elif action_type == "left_double": + pyautogui_code += f"\npyautogui.doubleClick({x}, {y}, button='left')" + elif action_type == "right_single": + pyautogui_code += f"\npyautogui.click({x}, {y}, button='right')" + elif action_type == "hover": + pyautogui_code += f"\npyautogui.moveTo({x}, {y})" + + elif action_type in ["finished"]: + pyautogui_code = f"DONE" + + else: + pyautogui_code += f"\n# Unrecognized action type: {action_type}" + + return pyautogui_code + +def add_box_token(input_string): + # Step 1: Split the string into individual actions + if "Action: " in input_string and "start_box=" in input_string: + suffix = input_string.split("Action: ")[0] + "Action: " + actions = input_string.split("Action: ")[1:] + processed_actions = [] + for action in actions: + action = action.strip() + # Step 2: Extract coordinates (start_box or end_box) using regex + coordinates = re.findall(r"(start_box|end_box)='\((\d+),\s*(\d+)\)'", action) + + updated_action = action # Start with the original action + for coord_type, x, y in coordinates: + # Convert x and y to integers + updated_action = updated_action.replace(f"{coord_type}='({x},{y})'", f"{coord_type}='<|box_start|>({x},{y})<|box_end|>'") + processed_actions.append(updated_action) + + # Step 5: Reconstruct the final string + final_string = suffix + "\n\n".join(processed_actions) + else: + final_string = input_string + return final_string + +def pil_to_base64(image): + buffer = BytesIO() + image.save(buffer, format="PNG") # 你可以改成 "JPEG" 等格式 + return base64.b64encode(buffer.getvalue()).decode("utf-8") + +def linearize_accessibility_tree(accessibility_tree, platform="ubuntu"): + + if platform == "ubuntu": + _attributes_ns = attributes_ns_ubuntu + _state_ns = state_ns_ubuntu + _component_ns = component_ns_ubuntu + _value_ns = value_ns_ubuntu + elif platform == "windows": + _attributes_ns = attributes_ns_windows + _state_ns = state_ns_windows + _component_ns = component_ns_windows + _value_ns = value_ns_windows + else: + raise ValueError("Invalid platform, must be 'ubuntu' or 'windows'") + + filtered_nodes = filter_nodes(ET.fromstring(accessibility_tree), platform) + linearized_accessibility_tree = [ + "tag\tname\ttext\tclass\tdescription\tposition (top-left x&y)\tsize (w&h)" + ] + + # Linearize the accessibility tree nodes into a table format + for node in filtered_nodes: + if node.text: + text = ( + node.text + if '"' not in node.text + else '"{:}"'.format(node.text.replace('"', '""')) + ) + + elif node.get("{{{:}}}class".format(class_ns_windows), "").endswith( + "EditWrapper" + ) and node.get("{{{:}}}value".format(_value_ns)): + node_text = node.get("{{{:}}}value".format(_value_ns), "") + text = ( + node_text + if '"' not in node_text + else '"{:}"'.format(node_text.replace('"', '""')) + ) + else: + text = '""' + + linearized_accessibility_tree.append( + "{:}\t{:}\t{:}\t{:}\t{:}\t{:}\t{:}".format( + node.tag, + node.get("name", ""), + text, + ( + node.get("{{{:}}}class".format(_attributes_ns), "") + if platform == "ubuntu" + else node.get("{{{:}}}class".format(class_ns_windows), "") + ), + node.get("{{{:}}}description".format(_attributes_ns), ""), + node.get("{{{:}}}screencoord".format(_component_ns), ""), + node.get("{{{:}}}size".format(_component_ns), ""), + ) + ) + + return "\n".join(linearized_accessibility_tree) + +def trim_accessibility_tree(linearized_accessibility_tree, max_tokens): + # enc = tiktoken.encoding_for_model("gpt-4") + # tokens = enc.encode(linearized_accessibility_tree) + # if len(tokens) > max_tokens: + # linearized_accessibility_tree = enc.decode(tokens[:max_tokens]) + # linearized_accessibility_tree += "[...]\n" + return linearized_accessibility_tree + +def extract_action_desp(raw_response: str) -> str: + """ + 从原始响应中提取 "Action desp:" 及其描述部分,不包括后续的Action部分。 + 如果未找到开始或结束标记,则抛出 ValueError。 + """ + start_marker = "Action desp:" + start_index = raw_response.find(start_marker) + + # 如果没找到开始标记,直接抛出错误 + if start_index == -1: + raise ValueError("未在响应中找到 'Action desp:' 标记。") + + # 找到下一个"\nAction:"的位置 + end_marker = "\nAction:" + end_index = raw_response.find(end_marker, start_index) + + # 如果没有找到结束标记,也抛出错误 + if end_index == -1: + raise ValueError("未在响应中找到 'Action desp:' 对应的结束标记 '\\nAction:'。") + + # 返回从开始标记到结束标记之前的内容 + return raw_response[start_index:end_index] + +def print_content_without_long_base64(content): + """Print content while truncating long base64 image data""" + import copy + printable_content = copy.deepcopy(content) + + for item in printable_content: + if item.get("type") == "image_url" and "image_url" in item: + url = item["image_url"]["url"] + if url.startswith("data:image"): + # Replace base64 data with a short indication + prefix = url[:30] # Keep the beginning + item["image_url"]["url"] = f"{prefix}...[BASE64_DATA_TRUNCATED]" + + print(printable_content) + +class ManoAgent: + def __init__( + self, + platform="ubuntu", + action_space="pyautogui", + observation_type="screenshot", + # observation_type can be in ["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"] + max_trajectory_length=100, + a11y_tree_max_tokens=10000, + model="mano", + model_type="qwen25vl", + runtime_conf: dict = { + "infer_mode": "qwen25vl_normal", + "prompt_style": "qwen25vl_normal", + "input_swap": True, + "language": "English", + "history_n": 3, + "max_pixels": 16384*28*28, + "min_pixels": 100*28*28, + "callusr_tolerance": 100, + "temperature": 0.0, + "top_k": -1, + "top_p": 0.9, + "max_tokens": 500 + + } + ): + self.platform = platform + self.action_space = action_space + self.observation_type = observation_type + self.max_trajectory_length = max_trajectory_length + self.a11y_tree_max_tokens = a11y_tree_max_tokens + self.model = model + self.model_type = model_type + self.runtime_conf = runtime_conf + self.vlm = OpenAI( + base_url=os.environ['MANO_API_URL'], + api_key=os.environ['MANO_API_KEY'], + ) + self.temperature = self.runtime_conf["temperature"] + self.top_k = self.runtime_conf["top_k"] + self.top_p = self.runtime_conf["top_p"] + self.max_tokens = self.runtime_conf["max_tokens"] + self.infer_mode = self.runtime_conf["infer_mode"] + self.prompt_style = self.runtime_conf["prompt_style"] + self.input_swap = self.runtime_conf["input_swap"] + self.language = self.runtime_conf["language"] + self.max_pixels = self.runtime_conf["max_pixels"] + self.min_pixels = self.runtime_conf["min_pixels"] + self.callusr_tolerance = self.runtime_conf["callusr_tolerance"] + + self.thoughts = [] + self.actions = [] + self.observations = [] + self.history_images = [] + self.history_responses = [] + self.action_parse_res_factor = 1000 + self.prompt_template = MANO_PROMPT_THOUGHT + + if "history_n" in self.runtime_conf: + self.history_n = self.runtime_conf["history_n"] + else: + self.history_n = 5 + + self.cur_callusr_count = 0 + + def predict( + self, instruction: str, obs: Dict, last_action_after_obs: Dict = None + ) -> List: + """ + Predict the next action(s) based on the current observation. + """ + + # Append trajectory + # print(len(self.observations), len(self.actions), len(self.actions)) + assert len(self.observations) == len(self.actions) and len(self.actions) == len( + self.thoughts + ), "The number of observations and actions should be the same." + + if len(self.observations) > self.max_trajectory_length: + if self.max_trajectory_length == 0: + _observations = [] + _actions = [] + _thoughts = [] + else: + _observations = self.observations[-self.max_trajectory_length :] + _actions = self.actions[-self.max_trajectory_length :] + _thoughts = self.thoughts[-self.max_trajectory_length :] + else: + _observations = self.observations + _actions = self.actions + _thoughts = self.thoughts + + for previous_obs, previous_action, previous_thought in zip( + _observations, _actions, _thoughts + ): + # {{{1 + if self.observation_type == "screenshot_a11y_tree": + _screenshot = previous_obs["screenshot"] + _linearized_accessibility_tree = previous_obs["accessibility_tree"] + elif self.observation_type == "screenshot": + _screenshot = previous_obs["screenshot"] + else: + raise ValueError( + "Invalid observation_type type: " + self.observation_type + ) # 1}}} + + self.history_images.append(obs["screenshot"]) + + if self.observation_type in ["screenshot", "screenshot_a11y_tree"]: + base64_image = obs["screenshot"] + try: + linearized_accessibility_tree = ( + linearize_accessibility_tree( + accessibility_tree=obs["accessibility_tree"], + platform=self.platform, + ) + if self.observation_type == "screenshot_a11y_tree" + else None + ) + except: + linearized_accessibility_tree = None + # logger.debug("LINEAR AT: %s", linearized_accessibility_tree) + + if linearized_accessibility_tree: + linearized_accessibility_tree = trim_accessibility_tree( + linearized_accessibility_tree, self.a11y_tree_max_tokens + ) + + if self.observation_type == "screenshot_a11y_tree": + self.observations.append( + { + "screenshot": base64_image, + "accessibility_tree": linearized_accessibility_tree, + } + ) + else: + self.observations.append( + {"screenshot": base64_image, "accessibility_tree": None} + ) + + else: + raise ValueError( + "Invalid observation_type type: " + self.observation_type + ) # 1}}} + + if self.infer_mode == "qwen25vl_normal": + user_prompt = self.prompt_template.format( + instruction=instruction, + language=self.language + ) + # print(user_prompt) + # exit() + + if len(self.history_images) > self.history_n: + self.history_images = self.history_images[-self.history_n:] + + messages, images = [], [] + if isinstance(self.history_images, bytes): + self.history_images = [self.history_images] + elif isinstance(self.history_images, np.ndarray): + self.history_images = list(self.history_images) + elif isinstance(self.history_images, list): + pass + else: + raise TypeError(f"Unidentified images type: {type(self.history_images)}") + + for turn, image in enumerate(self.history_images): + if len(images) >= self.history_n: + break + try: + image = Image.open(BytesIO(image)) + except Exception as e: + raise RuntimeError(f"Error opening image: {e}") + + if image.width * image.height > self.max_pixels: + """ + 如果图片超过/低于像素限制,则计算一个缩放因子resize_factor,使图片的像素数缩小到等于或小于max_pixels。这个缩放因子是通过开平方根计算的,确保纵横比保持不变,这样原始的相对坐标可以不经转换直接复用 + """ + resize_factor = math.sqrt(self.max_pixels / (image.width * image.height)) + width, height = int(image.width * resize_factor), int(image.height * resize_factor) + image = image.resize((width, height)) + if image.width * image.height < self.min_pixels: + resize_factor = math.sqrt(self.min_pixels / (image.width * image.height)) + width, height = math.ceil(image.width * resize_factor), math.ceil(image.height * resize_factor) + image = image.resize((width, height)) + + if image.mode != "RGB": + image = image.convert("RGB") + + images.append(image) + + # 修改为单轮对话形式 + messages = [ + { + "role": "system", + "content": [{"type": "text", "text": "You are a helpful assistant."}] + } + ] + + # 构建用户消息内容 + user_message_content = [] + + user_message_content.append({ + "type": "text", + "text": user_prompt + }) + #print(user_prompt) + + # 添加历史操作信息 - 包含所有历史步骤,不限制数量 + if len(self.history_responses) > 0: + # 确定哪些步骤需要显示图片 + steps_with_images = {} + + # 计算可用的历史图片数量(除去当前步骤的图片) + available_images = len(images) - 1 + + # 为最近的步骤分配图片 + for i in range(min(available_images, len(self.history_responses))): + # 最近的步骤优先获得图片 + step_idx = len(self.history_responses) - 1 - i + img_idx = available_images - 1 - i + steps_with_images[step_idx] = img_idx + + # 构建历史步骤文本 + for i, history_response in enumerate(self.history_responses): + step_num = i + 1 + + # 尝试提取Action描述 + try: + action_desp = extract_action_desp(history_response) + except ValueError as e: + print("获取Action desp失败") + raise ValueError(f"Action extraction error: {str(e)}") from e + + # 添加步骤描述文本 + if i > 0: # 添加换行,除了第一步 + user_message_content.append({"type": "text", "text": "\n"}) + + user_message_content.append({ + "type": "text", + "text": f"第{step_num}步:{action_desp}" + }) + + if i in steps_with_images: + img_idx = steps_with_images[i] + encoded_string = pil_to_base64(images[img_idx]) + user_message_content.append({ + "type": "text", + "text": "对应截图为:" + }) + user_message_content.append({ + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{encoded_string}"} + }) + + # 添加当前步骤描述 + user_message_content.append({ + "type": "text", + "text": "\n当前步骤的截图:" + }) + + # 添加当前步骤的图片 + encoded_string = pil_to_base64(images[-1]) + user_message_content.append({ + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{encoded_string}"} + }) + + # 构建最终的用户消息 + messages.append({ + "role": "user", + "content": user_message_content + }) + + #print_content_without_long_base64(user_message_content) + + + try_times = 3 + origin_resized_height = images[-1].height + origin_resized_width = images[-1].width + temperature = self.temperature + top_k = self.top_k + #print("[DEBUG] 模型调用中") + while True: + if try_times <= 0: + print(f"Reach max retry times to fetch response from client, as error flag.") + return "client error", ["DONE"], [] + try: + # print("in try") + response = self.vlm.chat.completions.create( + model="mano", + messages=messages, + #frequency_penalty=1, + max_tokens=self.max_tokens, + temperature=temperature + ) + # print("[DEBUG] 模型调用成功,收到 response") + #print(response.choices[0].message.content) + prediction = response.choices[0].message.content.strip() + #prediction = response[0]["prediction"].strip() + except Exception as e: + print(f"Error when fetching response from client, with response: {response}") + prediction = None + try_times -= 1 + + try: + parsed_responses = parse_action_to_structure_output( + prediction, + self.action_parse_res_factor, + origin_resized_height, + origin_resized_width, + self.model_type, + self.max_pixels, + self.min_pixels + ) + #print(f"[DEBUG] Step prediction: {prediction}") + #print(f"[DEBUG] Parsed responses: {parsed_responses}") + break + except Exception as e: + print(f"Error 111 when parsing response from client, with response: {response}") + # If fail to parse the model response, we use sampling parameters to avoid it + prediction = None + try_times -= 1 + temperature = 1 + top_k = -1 + + if prediction is None: + return "client error", ["DONE"] + + self.history_responses.append(prediction) + self.thoughts.append(prediction) + + try: + parsed_responses = parse_action_to_structure_output( + prediction, + self.action_parse_res_factor, + origin_resized_height, + origin_resized_width, + self.model_type, + self.max_pixels, + self.min_pixels + ) + except Exception as e: + print(f"Parsing action error: {prediction}, with error:\n{e}") + return f"Parsing action error: {prediction}, with error:\n{e}", ["DONE"] + + actions = [] + last_image = Image.open(BytesIO(self.history_images[-1])) + obs_image_height = last_image.height + obs_image_width = last_image.width + + for parsed_response in parsed_responses: + if "action_type" in parsed_response: + + if parsed_response["action_type"] == FINISH_WORD: + self.actions.append(actions) + + return prediction, ["DONE"] + + elif parsed_response["action_type"] == WAIT_WORD: + self.actions.append(actions) + return prediction, ["WAIT"] + + elif parsed_response["action_type"] == ENV_FAIL_WORD: + self.actions.append(actions) + return prediction, ["FAIL"] + + elif parsed_response["action_type"] == CALL_USER: + if self.callusr_tolerance > self.cur_callusr_count: + self.actions.append(actions) + self.cur_callusr_count += 1 + return prediction, ["WAIT"] + else: + self.actions.append(actions) + return prediction, ["FAIL"] + + pyautogui_code = parsing_response_to_pyautogui_code( + parsed_response, + obs_image_height, + obs_image_width, + self.input_swap + ) + actions.append(pyautogui_code) + + self.actions.append(actions) + + if len(self.history_responses) >= self.max_trajectory_length: + # Default to FAIL if exceed max steps + actions = ["FAIL"] + + return prediction, actions + + + @backoff.on_exception( + backoff.constant, + # here you should add more model exceptions as you want, + # but you are forbidden to add "Exception", that is, a common type of exception + # because we want to catch this kind of Exception in the outside to ensure each example won't exceed the time limit + ( + # General exceptions + SSLError, + # OpenAI exceptions + openai.RateLimitError, + openai.BadRequestError, + openai.InternalServerError, + # Google exceptions + InvalidArgument, + ResourceExhausted, + InternalServerError, + BadRequest, + # Groq exceptions + # todo: check + ), + interval=30, + max_tries=10, + ) + + def reset(self, runtime_logger): + self.thoughts = [] + self.actions = [] + self.observations = [] + self.history_images = [] + self.history_responses = [] diff --git a/run_multienv_mano.py b/run_multienv_mano.py new file mode 100644 index 0000000..3006efc --- /dev/null +++ b/run_multienv_mano.py @@ -0,0 +1,568 @@ +from __future__ import annotations +import argparse +import datetime +import json +import logging +import os +import sys +import signal +import time +from typing import List +from multiprocessing import Process, Manager +from multiprocessing import current_process +import lib_run_single +from desktop_env.desktop_env import DesktopEnv +from mm_agents.mano_agent import ManoAgent +import os + +# Global variables for signal handling +active_environments = [] +processes = [] +is_terminating = False + +# load the environment variables from .env file +if os.path.exists(".env"): + from dotenv import load_dotenv + load_dotenv() + +# Logger Configs {{{ # +def config() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run end-to-end evaluation on the benchmark" + ) + + # environment config + parser.add_argument("--path_to_vm", type=str, default=None) + parser.add_argument( + "--headless", action="store_true", help="Run in headless machine" + ) + parser.add_argument( + "--action_space", type=str, default="pyautogui", help="Action type" + ) + parser.add_argument( + "--observation_type", + choices=["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"], + default="screenshot", + help="Observation type", + ) + parser.add_argument("--sleep_after_execution", type=float, default=3.0) + parser.add_argument("--max_steps", type=int, default=15) + + # evaluation config + parser.add_argument( + "--test_config_base_dir", type=str, default="evaluation_examples" + ) + + # lm config + parser.add_argument("--model", type=str, default="mano") + parser.add_argument("--model_type", type=str, default="qwen25vl") + parser.add_argument("--infer_mode", type=str, default="qwen25vl_normal") + parser.add_argument("--prompt_style", type=str, default="qwen25vl_normal") + parser.add_argument("--input_swap", action="store_true", help="Use copy and paste to type content") + parser.add_argument("--language", type=str, default="Chinese") + parser.add_argument("--max_pixels", type=float, default=16384*28*28) + parser.add_argument("--min_pixels", type=float, default=100*28*28) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--top_p", type=float, default=0.9) + parser.add_argument("--top_k", type=int, default=-1) + parser.add_argument("--history_n", type=int, default=5) + parser.add_argument("--callusr_tolerance", type=int, default=3) + parser.add_argument("--max_tokens", type=int, default=1000) + parser.add_argument("--stop_token", type=str, default=None) + + parser.add_argument("--max_trajectory_length", type=int, default=None, help="The max number of trajectory steps.") + parser.add_argument("--max_image_history_length", type=int, default=5, help="The max number of images in the history.") + + # example config + parser.add_argument("--domain", type=str, default="all") + parser.add_argument( + "--test_all_meta_path", type=str, default="evaluation_examples/test_all.json" + ) + + # logging related + parser.add_argument("--result_dir", type=str, default="./results") + parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to run in parallel") + parser.add_argument("--log_level", type=str, choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], + default='INFO', help="Set the logging level") + # aws config + parser.add_argument( + "--region", type=str, default="us-east-1", help="AWS region for the VM" + ) + parser.add_argument( + "--provider_name", type=str, default="aws", choices=["aws", "virtualbox", "vmware", "docker", "azure"], help="Provider name" + ) + parser.add_argument( + "--client_password", type=str, default="", help="Client password" + ) + parser.add_argument( + "--screen_width", type=int, default=1920, help="Screen width" + ) + parser.add_argument( + "--screen_height", type=int, default=1080, help="Screen height" + ) + args = parser.parse_args() + return args + +args = config() # Get command line arguments first + +logger = logging.getLogger() +log_level = getattr(logging, args.log_level.upper()) +logger.setLevel(log_level) + +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) + +file_handler.setLevel(logging.INFO) +debug_handler.setLevel(logging.DEBUG) +stdout_handler.setLevel(log_level) + +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) + +stdout_handler.addFilter(logging.Filter("desktopenv")) + +logger.addHandler(file_handler) +logger.addHandler(debug_handler) +logger.addHandler(stdout_handler) +# }}} Logger Configs # + +logger = logging.getLogger("desktopenv.experiment") + + +def distribute_tasks(test_all_meta: dict) -> List[tuple]: + all_tasks = [] + for domain, examples in test_all_meta.items(): + for example_id in examples: + all_tasks.append((domain, example_id)) + return all_tasks + + +def process_signal_handler(signum, frame, env_idx): + """Signal handler for child processes to gracefully shut down their environments.""" + logger.info(f"Process {env_idx + 1} received signal {signum}. Shutting down...") + + # Get the active_environments from the caller's frame + local_vars = frame.f_locals + active_environments = local_vars.get('active_environments', []) + + # Close environment in the current process context + for env in active_environments: + if env is not None: + try: + logger.info(f"Process {env_idx + 1} closing environment...") + env.close() + logger.info(f"Process {env_idx + 1} environment closed successfully") + except Exception as e: + logger.error(f"Process {env_idx + 1} error closing environment: {e}") + + logger.info(f"Process {env_idx + 1} shutdown complete. Exiting.") + sys.exit(0) + +def run_env_tasks(task_queue: Queue, args: argparse.Namespace, shared_scores: list): + active_environments = [] + env = None + try: + #from desktop_env.providers.aws.manager import IMAGE_ID_MAP + #REGION = args.region + screen_size = (args.screen_width, args.screen_height) + #ami_id = IMAGE_ID_MAP[REGION].get(screen_size, IMAGE_ID_MAP[REGION][(1920, 1080)]) + env = DesktopEnv( + path_to_vm=args.path_to_vm, + action_space=args.action_space, + provider_name=args.provider_name, + #region=REGION, + #snapshot_name=ami_id, + screen_size=screen_size, + headless=args.headless, + os_type="Ubuntu", + require_a11y_tree=args.observation_type in ["a11y_tree", "screenshot_a11y_tree", "som"], + #enable_proxy=True, + #client_password=args.client_password + ) + + + active_environments.append(env) + args.max_trajectory_length = args.max_steps + if args.infer_mode == "qwen25vl_normal": + runtime_conf: dict = { + "infer_mode": "qwen25vl_normal", + "prompt_style": "qwen25vl_normal", + "input_swap": False, #True, + "language": "Chinese", + "history_n": 3, + "max_pixels": 16384*28*28, + "min_pixels": 100*28*28, + "callusr_tolerance": 3, + "temperature": 0.0, + "top_k": -1, + "top_p": 0.9, + "max_tokens": 1000 + + } + else: + raise ValueError(f"Unknown infer_mode: {args.infer_mode}") + + agent = ManoAgent( + model=args.model, + action_space=args.action_space, + observation_type=args.observation_type, + max_trajectory_length=args.max_trajectory_length, + model_type=args.model_type, + runtime_conf = runtime_conf + ) + + logger.info(f"Process {current_process().name} started.") + while True: + try: + item = task_queue.get(timeout=5) + except Exception: + break + domain, example_id = item + try: + config_file = os.path.join( + args.test_config_base_dir, f"examples/{domain}/{example_id}.json" + ) + with open(config_file, "r", encoding="utf-8") as f: + example = json.load(f) + logger.info(f"[{current_process().name}][Domain]: {domain}") + logger.info(f"[{current_process().name}][Example ID]: {example_id}") + logger.info(f"[{current_process().name}][Instruction]: {example['instruction']}") + example_result_dir = os.path.join( + args.result_dir, + args.action_space, + args.observation_type, + args.model, + domain, + example_id, + ) + os.makedirs(example_result_dir, exist_ok=True) + try: + lib_run_single.run_single_example_mano( + agent, + env, + example, + args.max_steps, + example["instruction"], + args, + example_result_dir, + shared_scores, + ) + except Exception as e: + import traceback + logger.error(f"Exception in {current_process().name} {domain}/{example_id}: {e}") + logger.error(traceback.format_exc()) + try: + env.controller.end_recording( + os.path.join(example_result_dir, "recording.mp4") + ) + except Exception as rec_e: + logger.error(f"Failed to end recording: {rec_e}") + with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f: + f.write( + json.dumps( + {"Error": f"{domain}/{example_id} - {e}"} + ) + ) + f.write("\n") + except Exception as e: + logger.error(f"Task-level error in {current_process().name}: {e}") + import traceback + logger.error(traceback.format_exc()) + except Exception as e: + logger.error(f"Process-level error in {current_process().name}: {e}") + import traceback + logger.error(traceback.format_exc()) + finally: + logger.info(f"{current_process().name} cleaning up environment...") + try: + if env: + env.close() + logger.info(f"{current_process().name} environment closed successfully") + except Exception as e: + logger.error(f"{current_process().name} error during environment cleanup: {e}") + + + +def signal_handler(signum, frame): + """Handle termination signals (SIGINT, SIGTERM) to gracefully shutdown environments.""" + global is_terminating, active_environments, processes + + # Avoid duplicate handling + if is_terminating: + return + + is_terminating = True + logger.info(f"Received signal {signum}. Gracefully shutting down...") + + # Close all registered environments in the main process + for env in active_environments: + try: + logger.info(f"Closing environment...") + env.close() + logger.info(f"Environment closed successfully") + except Exception as e: + logger.error(f"Error closing environment: {e}") + + # Send termination signal to all child processes first + for p in processes: + if p.is_alive(): + try: + logger.info(f"Sending termination signal to process {p.name}...") + p.terminate() + except Exception as e: + logger.error(f"Error sending termination signal to process: {e}") + + # Allow a short time for processes to handle their own cleanup + time.sleep(1) + + # Forcefully terminate any processes that didn't exit + for p in processes: + if p.is_alive(): + try: + logger.info(f"Forcefully terminating process {p.name}...") + import signal as sig + os.kill(p.pid, sig.SIGKILL) + except Exception as e: + logger.error(f"Error forcefully terminating process: {e}") + + logger.info("Shutdown complete. Exiting.") + sys.exit(0) + + +def test(args: argparse.Namespace, test_all_meta: dict) -> None: + global processes + logger.info("Args: %s", args) + all_tasks = distribute_tasks(test_all_meta) + logger.info(f"Total tasks: {len(all_tasks)}") + with Manager() as manager: + shared_scores = manager.list() + task_queue = manager.Queue() + for item in all_tasks: + task_queue.put(item) + num_envs = args.num_envs + processes = [] + for i in range(num_envs): + p = Process( + target=run_env_tasks, + args=(task_queue, args, shared_scores), + name=f"EnvProcess-{i+1}" + ) + p.daemon = True + p.start() + processes.append(p) + logger.info(f"Started process {p.name} with PID {p.pid}") + try: + while True: + alive_count = 0 + for idx, p in enumerate(processes): + if not p.is_alive(): + logger.warning(f"Process {p.name} died, restarting...") + new_p = Process( + target=run_env_tasks, + args=(task_queue, args, shared_scores), + name=f"EnvProcess-Restart-{idx+1}" + ) + new_p.daemon = True + new_p.start() + processes[idx] = new_p + logger.info(f"Restarted process {new_p.name} with PID {new_p.pid}") + else: + alive_count += 1 + if task_queue.empty(): + logger.info("All tasks finished.") + break + if alive_count == 0: + logger.error("All processes died, exiting.") + break + time.sleep(5) + for p in processes: + p.join() + except KeyboardInterrupt: + logger.info("Main process received KeyboardInterrupt. Initiating graceful shutdown...") + raise + except Exception as e: + logger.error(f"Unexpected error while waiting for processes: {e}", exc_info=True) + for p in processes: + if p.is_alive(): + try: + logger.info(f"Terminating process {p.name} due to error...") + p.terminate() + except Exception as term_e: + logger.error(f"Error terminating process {p.name}: {term_e}") + raise + scores = list(shared_scores) + logger.info(f"Average score: {sum(scores) / len(scores) if scores else 0}") + + +def get_unfinished( + action_space, use_model, observation_type, result_dir, total_file_json +): + target_dir = os.path.join(result_dir, action_space, observation_type, use_model) + + if not os.path.exists(target_dir): + return total_file_json + + finished = {} + for domain in os.listdir(target_dir): + finished[domain] = [] + domain_path = os.path.join(target_dir, domain) + if os.path.isdir(domain_path): + for example_id in os.listdir(domain_path): + if example_id == "onboard": + continue + example_path = os.path.join(domain_path, example_id) + if os.path.isdir(example_path): + if "result.txt" not in os.listdir(example_path): + # empty all files under example_id + for file in os.listdir(example_path): + os.remove(os.path.join(example_path, file)) + else: + finished[domain].append(example_id) + + if not finished: + return total_file_json + + for domain, examples in finished.items(): + if domain in total_file_json: + total_file_json[domain] = [ + x for x in total_file_json[domain] if x not in examples + ] + + return total_file_json + + +def get_result(action_space, use_model, observation_type, result_dir, total_file_json): + target_dir = os.path.join(result_dir, action_space, observation_type, use_model) + if not os.path.exists(target_dir): + print("New experiment, no result yet.") + return None + + all_result = [] + + for domain in os.listdir(target_dir): + domain_path = os.path.join(target_dir, domain) + if os.path.isdir(domain_path): + for example_id in os.listdir(domain_path): + example_path = os.path.join(domain_path, example_id) + if os.path.isdir(example_path): + if "result.txt" in os.listdir(example_path): + # empty all files under example_id + try: + all_result.append( + float( + open( + os.path.join(example_path, "result.txt"), "r" + ).read() + ) + ) + except: + all_result.append(0.0) + + if not all_result: + print("New experiment, no result yet.") + return None + else: + print("Current Success Rate:", sum(all_result) / len(all_result) * 100, "%") + return all_result + + +if __name__ == "__main__": + ####### The complete version of the list of examples ####### + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + # Register signal handlers for graceful termination + signal.signal(signal.SIGINT, signal_handler) # Handle Ctrl+C + signal.signal(signal.SIGTERM, signal_handler) # Handle termination signal + + try: + args = config() + + # save args to json in result_dir/action_space/observation_type/model/args.json + path_to_args = os.path.join( + args.result_dir, + args.action_space, + args.observation_type, + args.model, + "args.json", + ) + os.makedirs(os.path.dirname(path_to_args), exist_ok=True) + with open(path_to_args, "w", encoding="utf-8") as f: + json.dump(vars(args), f, indent=4) + + with open(args.test_all_meta_path, "r", encoding="utf-8") as f: + test_all_meta = json.load(f) + + if args.domain != "all": + test_all_meta = {args.domain: test_all_meta[args.domain]} + + test_file_list = get_unfinished( + args.action_space, + args.model, + args.observation_type, + args.result_dir, + test_all_meta, + ) + left_info = "" + for domain in test_file_list: + left_info += f"{domain}: {len(test_file_list[domain])}\n" + logger.info(f"Left tasks:\n{left_info}") + + get_result( + args.action_space, + args.model, + args.observation_type, + args.result_dir, + test_all_meta, + ) + test(args, test_file_list) + except KeyboardInterrupt: + logger.info("Main process received KeyboardInterrupt.") + # Signal handler will take care of cleanup + except Exception as e: + logger.error(f"Unexpected error in main process: {e}", exc_info=True) + # Also trigger cleanup for unhandled exceptions + signal_handler(signal.SIGTERM, None) + finally: + # Final cleanup in case any environments or processes remain + logger.info("Main process final cleanup...") + for env in active_environments: + if env is not None: + try: + logger.info(f"Closing environment in final cleanup...") + env.close() + logger.info(f"Environment closed successfully in final cleanup") + except Exception as e: + logger.error(f"Error during final environment cleanup: {e}") + + # First try gentle termination + for p in processes: + if p is not None and p.is_alive(): + try: + logger.info(f"Terminating process {p.name}...") + p.terminate() + except Exception as e: + logger.error(f"Error terminating process: {e}") + + # Wait a moment for processes to terminate + time.sleep(1) + + # Then force kill if needed + for p in processes: + if p is not None and p.is_alive(): + try: + logger.info(f"Force killing process {p.name}...") + os.kill(p.pid, signal.SIGKILL) + logger.info(f"Process {p.name} force killed") + except Exception as e: + logger.error(f"Error force killing process: {e}")