From c8d8b1d1673a022dc14954ace2fbb0be9f9c7f54 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Sun, 1 Mar 2026 17:33:05 -0500 Subject: [PATCH 01/29] Add Env EB-Alfred --- examples/evaluate/eb_alfred/config.yaml | 82 +++++ vagen/envs/eb_alfred/__init__.py | 0 vagen/envs/eb_alfred/eb_alfred_env.py | 383 ++++++++++++++++++++++++ vagen/envs/eb_alfred/handler.py | 95 ++++++ vagen/envs/eb_alfred/serve.py | 92 ++++++ vagen/envs/eb_alfred/utils/__init__.py | 0 vagen/envs/eb_alfred/utils/prompt.py | 200 +++++++++++++ vagen/envs/eb_alfred/utils/utils.py | 154 ++++++++++ 8 files changed, 1006 insertions(+) create mode 100644 examples/evaluate/eb_alfred/config.yaml create mode 100644 vagen/envs/eb_alfred/__init__.py create mode 100644 vagen/envs/eb_alfred/eb_alfred_env.py create mode 100644 vagen/envs/eb_alfred/handler.py create mode 100644 vagen/envs/eb_alfred/serve.py create mode 100644 vagen/envs/eb_alfred/utils/__init__.py create mode 100644 vagen/envs/eb_alfred/utils/prompt.py create mode 100644 vagen/envs/eb_alfred/utils/utils.py diff --git a/examples/evaluate/eb_alfred/config.yaml b/examples/evaluate/eb_alfred/config.yaml new file mode 100644 index 000000000..c1fa62d0a --- /dev/null +++ b/examples/evaluate/eb_alfred/config.yaml @@ -0,0 +1,82 @@ +# EB-ALFRED Evaluation Config (ERA-aligned) +# +# Uses the remote environment pattern: +# 1. Start the EB-ALFRED server (requires GPU + X display): +# DISPLAY=:0 python -m vagen.envs.eb_alfred.serve --port 8000 +# +# 2. Run evaluation: +# python -m vagen.evaluate.run_eval --config examples/evaluate/eb_alfred/config.yaml +# +# Key ERA-aligned settings: +# - No-concat mode: system prompt (with task + actions) + current obs only per turn +# - Multi-step planning: up to 20 actions per LLM call +# - Max 30 env steps per episode (matches ERA's _max_episode_steps) +# - ERA-style replan: break on action failure, model replans next turn +# - Task examples in system prompt (cleaning, slicing, heating patterns) + +envs: + - name: RemoteEnv + n_envs: 50 + tag_id: eb_alfred_eval + seed: [0, 50, 1] + split: test + concat_history: false # no-concat: system + current obs only per turn + max_turns: 30 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: base # options: base, common_sense, complex, long_horizon + x_display: "0" + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 # multi-step planning (ERA-aligned) + max_env_steps: 30 # total env actions cap (ERA-aligned) + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + +experiment: + dump_dir: ./rollouts/eval_eb_alfred + default_max_turns: 30 + +run: + backend: openai + base_seed: 0 + max_concurrent_jobs: 10 + resume: skip_completed + live_summary: true + +backends: + openai: + api_key: "" # uses OPENAI_API_KEY env var + base_url: null + model: "gpt-4.1" + max_concurrency: 100 + max_retries: 6 + min_backoff: 0.5 + max_backoff: 8.0 + + sglang: + base_url: "http://127.0.0.1:30000/v1" + api_key: "EMPTY" + model: "Qwen/Qwen2.5-VL-7B-Instruct" + max_concurrency: 50 + max_retries: 6 + min_backoff: 0.5 + max_backoff: 8.0 + + claude: + api_key: "" + base_url: null + model: "claude-sonnet-4-6" + max_concurrency: 10 + max_retries: 6 + min_backoff: 0.5 + max_backoff: 8.0 diff --git a/vagen/envs/eb_alfred/__init__.py b/vagen/envs/eb_alfred/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py new file mode 100644 index 000000000..20eed508b --- /dev/null +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -0,0 +1,383 @@ +""" +EB-ALFRED environment adapter for VAGEN. + +Wraps the EBAlfEnv from EmbodiedBench as a GymImageEnv, +enabling integration with VAGEN's RL training and evaluation pipeline. + +The underlying EBAlfEnv uses AI2-THOR for 3D household robot task simulation. +It requires a GPU-accelerated X server for rendering. +""" + +import asyncio +import os +import numpy as np +from PIL import Image +from dataclasses import dataclass, field +from typing import Any, Dict, Tuple, List, Optional + +from .utils.prompt import ( + system_prompt, + format_prompt, + init_observation_template, + action_template, +) +from .utils.utils import parse_response, match_action, numpy_to_pil + +from vagen.envs.gym_image_env import GymImageEnv + + +@dataclass +class EbAlfredEnvConfig: + """Configuration for EB-ALFRED environment.""" + + # Environment settings + eval_set: str = "base" + exp_name: str = "vagen_eval" + down_sample_ratio: float = 1.0 + resolution: int = 500 + x_display: str = "1" + selected_indexes: List[int] = field(default_factory=list) + detection_box: bool = False + + # Interaction settings + max_turns: int = 30 + max_actions_per_step: int = 20 + max_env_steps: int = 30 # Max total environment actions per episode (matches ERA) + action_sep: str = "," + image_placeholder: str = "" + prompt_format: str = "free_think" + use_example_in_sys_prompt: bool = True + + # Observation image settings + obs_image_size: Optional[int] = None # Resize obs image to this size (square). None = use original. + + # Reward settings + format_reward: float = 0.1 + success_reward: float = 1.0 + + +class EbAlfred(GymImageEnv): + """ + EB-ALFRED environment implementing the GymImageEnv async interface. + + Wraps EBAlfEnv from EmbodiedBench, which uses AI2-THOR for + 3D household robot task simulation (e.g., "Clean a rag, put it away"). + + Key features: + - 162+ discrete actions (find, pick up, put down, open, close, etc.) + - Dynamic action space per episode (multi-instance objects) + - Vision-only observations (RGB images from AI2-THOR) + - Dense reward via task progress + format reward + """ + + def __init__(self, env_config: Dict[str, Any]): + super().__init__(env_config) + + # Filter config keys to only those in the dataclass + valid_keys = EbAlfredEnvConfig.__dataclass_fields__ + filtered = {k: v for k, v in env_config.items() if k in valid_keys} + self.config = EbAlfredEnvConfig(**filtered) + + # Set X display before importing/creating EBAlfEnv + import embodiedbench.envs.eb_alfred.EBAlfEnv as ebalfenv_mod + ebalfenv_mod.X_DISPLAY = self.config.x_display + from embodiedbench.envs.eb_alfred.EBAlfEnv import EBAlfEnv + + self.env = EBAlfEnv( + eval_set=self.config.eval_set, + exp_name=self.config.exp_name, + down_sample_ratio=self.config.down_sample_ratio, + selected_indexes=self.config.selected_indexes, + detection_box=self.config.detection_box, + resolution=self.config.resolution, + ) + + # Adapter state (reset per episode) + self._total_turns: int = 0 + self._total_env_steps: int = 0 + self._last_action: str = "" + self._last_feedback: str = "" + self._action_list: List[str] = [] + self._action_map: Dict[str, str] = {} # lowercase -> original + + # ------------------------------------------------------------------ + # GymImageEnv abstract methods + # ------------------------------------------------------------------ + + async def close(self) -> None: + """Close AI2-THOR process.""" + await asyncio.to_thread(self.env.close) + + async def system_prompt(self) -> Dict[str, Any]: + """ + Return the system prompt with per-episode task and action list. + + Includes role description, action descriptions, guidelines, + the current task instruction, available actions, and format + instructions. This ensures no-concat mode always has access + to task and action information. + """ + sys_str = system_prompt( + task_instruction=self.env.episode_language_instruction, + action_list=self._action_list, + ) + fmt_str = format_prompt( + max_actions_per_step=self.config.max_actions_per_step, + action_sep=self.config.action_sep, + add_example=self.config.use_example_in_sys_prompt, + prompt_format=self.config.prompt_format, + ) + return {"obs_str": sys_str + "\n" + fmt_str} + + async def reset(self, seed: int) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Reset environment for a new episode. + + The seed selects which episode to load from the dataset + (seed % number_of_episodes). After reset, the observation + includes the task instruction, available actions, and + the initial RGB image from AI2-THOR. + """ + # Select episode based on seed + episode_idx = seed % self.env.number_of_episodes + self.env._current_episode_num = episode_idx + + await asyncio.to_thread(self.env.reset) + + # Reset adapter state + self._total_turns = 0 + self._total_env_steps = 0 + self._last_action = "" + self._last_feedback = "" + + # Build action lookup for this episode (action space is dynamic) + self._action_list = list(self.env.language_skill_set) + self._action_map = {a.lower(): a for a in self._action_list} + + # Build observation + obs = self._build_obs(init=True) + info = { + "task_instruction": self.env.episode_language_instruction, + "num_actions": len(self._action_list), + "eval_set": self.config.eval_set, + "episode_idx": episode_idx, + } + return obs, info + + async def step( + self, action_str: str + ) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]: + """ + Execute one step given the LLM's response. + + Parses ...... from action_str, + matches the action against the current action space, and + executes it in AI2-THOR. + """ + self._total_turns += 1 + + # Parse LLM response + parsed = parse_response( + response=action_str, + action_sep=self.config.action_sep, + max_actions=self.config.max_actions_per_step, + prompt_format=self.config.prompt_format, + ) + + reward = 0.0 + done = False + info: Dict[str, Any] = {} + info.update(parsed) + + actions = parsed.get("actions", []) + format_correct = parsed.get("format_correct", False) + + metrics = { + "turn_metrics": { + "action_is_valid": False, + "action_is_effective": False, + }, + "traj_metrics": { + "success": False, + }, + } + + if format_correct and actions: + reward += self.config.format_reward + + # Clip actions to remaining env step budget (ERA-style) + remaining = self.config.max_env_steps - self._total_env_steps + actions = actions[:remaining] if remaining > 0 else [] + + for action_name in actions: + matched = match_action(action_name, self._action_list, self._action_map) + + if matched is None: + # Action name not recognized + self._last_action = action_name + self._last_feedback = ( + f"Action '{action_name}' is not a recognized action." + ) + break + + metrics["turn_metrics"]["action_is_valid"] = True + + # Execute in AI2-THOR + self._total_env_steps += 1 + obs_raw, step_reward, step_done, step_info = ( + await asyncio.to_thread(self.env.step, matched) + ) + + self._last_action = matched + self._last_feedback = step_info.get("env_feedback", "") + + action_success = step_info.get("last_action_success", 0.0) + if action_success: + metrics["turn_metrics"]["action_is_effective"] = True + + task_success = step_info.get("task_success", 0.0) + if task_success: + done = True + reward += self.config.success_reward + metrics["traj_metrics"]["success"] = True + break + + if step_done: + done = True + break + + # ERA-style: break on action failure to replan + if not action_success: + break + + # Check env step limit + if self._total_env_steps >= self.config.max_env_steps: + done = True + break + else: + # Format error: no valid actions parsed + self._last_action = parsed.get("action_content", "") + self._last_feedback = ( + "Could not parse a valid action from your response. " + "Please use the format: ...action name" + ) + + # Check turn limit and env step limit + if self._total_turns >= self.config.max_turns: + done = True + if self._total_env_steps >= self.config.max_env_steps: + done = True + + info["metrics"] = metrics + info["success"] = metrics["traj_metrics"]["success"] + + obs = self._build_obs(init=False) + return obs, reward, done, info + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_obs(self, init: bool) -> Dict[str, Any]: + """Build observation dict with image and text.""" + frame = self.env.env.last_event.frame + img = numpy_to_pil(frame) + if self.config.obs_image_size is not None: + sz = self.config.obs_image_size + img = img.resize((sz, sz), Image.LANCZOS) + img_str = self.config.image_placeholder + + if init: + obs_str = init_observation_template( + img_str=img_str, + ) + else: + obs_str = action_template( + last_action=self._last_action, + env_feedback=self._last_feedback, + img_str=img_str, + ) + + return { + "obs_str": obs_str + "\n", + "multi_modal_input": { + self.config.image_placeholder: [img] + }, + } + + +# ------------------------------ +# Local async test (optional) +# ------------------------------ +if __name__ == "__main__": + import fire + import logging + + logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") + + async def main_async( + eval_set: str = "base", + resolution: int = 500, + x_display: str = "1", + save_path: str = "./test_eb_alfred", + prompt_format: str = "free_think", + ): + cfg = { + "eval_set": eval_set, + "resolution": resolution, + "x_display": x_display, + "prompt_format": prompt_format, + } + env = EbAlfred(cfg) + + print("System Prompt:") + sys_prompt = await env.system_prompt() + print(sys_prompt["obs_str"]) + print("\n" + "=" * 50 + "\n") + + obs, info = await env.reset(seed=0) + print(f"Task: {info['task_instruction']}") + print(f"Available actions: {info['num_actions']}") + print(f"Observation:\n{obs['obs_str'][:200]}...") + + step = 0 + os.makedirs(save_path, exist_ok=True) + if "multi_modal_input" in obs: + img = obs["multi_modal_input"][env.config.image_placeholder][0] + img.save(os.path.join(save_path, f"step_{step}.png")) + + while True: + step += 1 + print(f"\nStep {step}:") + try: + action_input = input("Enter action (or 'quit'): ") + except EOFError: + action_input = "quit" + + if action_input.lower() == "quit": + break + + if not action_input.startswith(""): + action_input = ( + f"Executing the action." + f"{action_input}" + ) + + obs, reward, done, info = await env.step(action_input) + if "multi_modal_input" in obs: + img = obs["multi_modal_input"][env.config.image_placeholder][0] + img.save(os.path.join(save_path, f"step_{step}.png")) + print(f"Reward: {reward}, Done: {done}") + print(f"Success: {info.get('success', False)}") + print(f"Observation:\n{obs['obs_str'][:200]}...") + + if done: + print("Episode finished!") + break + + await env.close() + + def main(**kwargs): + asyncio.run(main_async(**kwargs)) + + fire.Fire(main) diff --git a/vagen/envs/eb_alfred/handler.py b/vagen/envs/eb_alfred/handler.py new file mode 100644 index 000000000..7d99cf88b --- /dev/null +++ b/vagen/envs/eb_alfred/handler.py @@ -0,0 +1,95 @@ +""" +EB-ALFRED handler for the remote gym environment service. + +This is the only component that needs customization. +It implements create_env() to instantiate EB-ALFRED environments +with automatic multi-GPU load balancing. +""" + +import asyncio +import logging +import subprocess +from typing import Any, Dict, List, Optional + +from vagen.envs_remote.handler import BaseGymHandler +from .eb_alfred_env import EbAlfred + +LOGGER = logging.getLogger(__name__) + + +def detect_gpu_displays() -> List[str]: + """Auto-detect available GPUs via nvidia-smi, return display list. + + Assumes display :i maps to GPU i (standard convention for + multi-GPU X server setups with Xvfb or xinit per GPU). + """ + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0: + indices = [ + line.strip() + for line in result.stdout.strip().split("\n") + if line.strip() + ] + if indices: + return indices + except (FileNotFoundError, subprocess.TimeoutExpired): + pass + return ["0"] + + +class EbAlfredHandler(BaseGymHandler): + """Handler for EB-ALFRED with automatic multi-GPU load balancing. + + By default, auto-detects available GPUs and distributes new + sessions to the least-loaded GPU. Single-GPU is just the + special case where only one GPU is detected. + """ + + def __init__( + self, + x_displays: Optional[List[str]] = None, + **kwargs, + ): + """ + Args: + x_displays: List of X display IDs to use (e.g. ["0", "1"]). + None = auto-detect GPUs via nvidia-smi. + **kwargs: Passed to BaseGymHandler (session_timeout, max_sessions). + """ + super().__init__(**kwargs) + self._x_displays = x_displays if x_displays is not None else detect_gpu_displays() + LOGGER.info(f"[Handler] Using X displays: {self._x_displays}") + + def _least_loaded_display(self) -> str: + """Pick the display with the fewest active sessions.""" + counts = {d: 0 for d in self._x_displays} + for ctx in self._sessions.values(): + d = getattr(ctx.env, "_assigned_display", None) + if d in counts: + counts[d] += 1 + chosen = min(counts, key=counts.get) + LOGGER.debug(f"[Handler] GPU load: {counts}, assigning display :{chosen}") + return chosen + + async def create_env(self, env_config: Dict[str, Any]) -> Any: + """ + Create an EbAlfred environment on the least-loaded GPU. + + AI2-THOR startup is blocking, so we offload to a thread. + """ + display = self._least_loaded_display() + env_config = {**env_config, "x_display": display} + + env = await asyncio.to_thread(EbAlfred, env_config) + env._assigned_display = display + LOGGER.info( + f"[Handler] Created env on display :{display} " + f"(GPU load: { {d: sum(1 for c in self._sessions.values() if getattr(c.env, '_assigned_display', None) == d) for d in self._x_displays} })" + ) + return env diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py new file mode 100644 index 000000000..d422370f2 --- /dev/null +++ b/vagen/envs/eb_alfred/serve.py @@ -0,0 +1,92 @@ +""" +EB-ALFRED Remote Environment Server. + +Starts a FastAPI service that exposes EB-ALFRED as a remote gym environment. +The service can run on a machine with GPU + X server (for AI2-THOR rendering), +while VAGEN RL training runs on a separate machine using GymImageEnvClient. + +Multi-GPU is the default: GPUs are auto-detected and sessions are +distributed to the least-loaded GPU automatically. + +Usage: + # Auto-detect GPUs (default) + python -m vagen.envs.eb_alfred.serve --port 8000 + + # Override: use only specific GPUs + python -m vagen.envs.eb_alfred.serve --port 8000 --x-displays 0,1 + + # Then on the training machine, configure env_config: + # base_urls: ["http://:8000"] + # eval_set: "base" + # resolution: 500 +""" + +import argparse +import asyncio +import concurrent.futures +import uvicorn + +from vagen.envs_remote.service import build_gym_service +from .handler import EbAlfredHandler + + +def main(): + parser = argparse.ArgumentParser(description="EB-ALFRED Remote Environment Server") + parser.add_argument("--port", type=int, default=8000, help="Server port") + parser.add_argument("--host", type=str, default="0.0.0.0", help="Server host") + parser.add_argument( + "--x-displays", + type=str, + default=None, + help="X displays for GPU assignment (comma-separated, e.g. '0,1'). " + "Default: auto-detect all GPUs via nvidia-smi.", + ) + parser.add_argument( + "--session-timeout", + type=float, + default=3600.0, + help="Session timeout in seconds", + ) + parser.add_argument( + "--max-sessions", + type=int, + default=0, + help="Max concurrent sessions (0=unlimited)", + ) + parser.add_argument( + "--thread-workers", + type=int, + default=128, + help="Thread pool size for Unity instance creation (default: 128)", + ) + args = parser.parse_args() + + x_displays = args.x_displays.split(",") if args.x_displays else None + + handler = EbAlfredHandler( + x_displays=x_displays, + session_timeout=args.session_timeout, + max_sessions=args.max_sessions, + ) + app = build_gym_service(handler) + + # Expand the asyncio thread pool via FastAPI startup so concurrent Unity + # startups don't queue behind Python's default limit of min(32, cpu+4). + _thread_workers = args.thread_workers + + @app.on_event("startup") + async def _set_thread_pool(): + loop = asyncio.get_event_loop() + loop.set_default_executor( + concurrent.futures.ThreadPoolExecutor(max_workers=_thread_workers) + ) + + displays_str = ", ".join(f":{d}" for d in handler._x_displays) + print(f"Starting EB-ALFRED service on {args.host}:{args.port}") + print(f"GPU displays: [{displays_str}] (auto-balanced)") + print(f"Health check: http://localhost:{args.port}/health") + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/vagen/envs/eb_alfred/utils/__init__.py b/vagen/envs/eb_alfred/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vagen/envs/eb_alfred/utils/prompt.py b/vagen/envs/eb_alfred/utils/prompt.py new file mode 100644 index 000000000..8521c6716 --- /dev/null +++ b/vagen/envs/eb_alfred/utils/prompt.py @@ -0,0 +1,200 @@ +from typing import List, Optional + + +# ERA-aligned task examples teaching ALFRED-specific mechanics +# (cleaning, heating, slicing, storage, object placement) +TASK_EXAMPLES = [ + { + "task": "Pick up the alarm clock and turn on the lamp", + "think": "I need to find the alarm clock, pick it up, then find the desk lamp and turn it on.", + "actions": ["find a AlarmClock", "pick up the AlarmClock", "find a DeskLamp", "turn on the DeskLamp"], + }, + { + "task": "Set the box on the table", + "think": "I need to find the box, pick it up, then find the dining table and put it down.", + "actions": ["find a Box", "pick up the Box", "find a DiningTable", "put down the object in hand"], + }, + { + "task": "Move the towel on the hanger into the cabinet", + "think": "I need to find the hand towel, pick it up, find a cabinet, open it, put the towel inside, and close the cabinet.", + "actions": ["find a HandTowel", "pick up the HandTowel", "find a Cabinet", "open the Cabinet", "put down the object in hand", "close the Cabinet"], + }, + { + "task": "Put a clean pan in the refrigerator", + "think": "To clean the pan, I need to: pick it up, put it in the sink, turn on the faucet, turn off the faucet, then pick up the now-clean pan and put it in the fridge.", + "actions": ["find a Pan", "pick up the Pan", "find a Sink", "put down the object in hand", "find a Faucet", "turn on the Faucet", "turn off the Faucet", "find a Pan", "pick up the Pan", "find a Fridge", "open the Fridge", "put down the object in hand", "close the Fridge"], + }, + { + "task": "Slice a loaf of bread put a slice on the counter", + "think": "To slice bread I need a knife first. Pick up knife, find bread, slice it. Put knife away, then pick up a bread slice and place it on the counter.", + "actions": ["find a Knife", "pick up the Knife", "find a Bread", "slice the Bread", "find a CounterTop", "put down the object in hand", "find a Bread", "pick up the Bread", "find a CounterTop", "put down the object in hand"], + }, +] + + +def system_prompt(task_instruction: Optional[str] = None, action_list: Optional[List[str]] = None, add_task_examples: bool = True): + """ + System prompt for EB-ALFRED household robot tasks. + + When task_instruction and action_list are provided (after reset), + includes the per-episode task and available actions so that + no-concat mode always has access to them. + """ + base = """You are a robot operating in a home. Given a task, you must accomplish the task using a defined set of actions to achieve the desired outcome. + +## Action Descriptions and Validity Rules +- Find: Parameterized by the name of the receptacle to navigate to. Always valid if the object exists in the scene. +- Pick up: Parameterized by the name of the object to pick. Only valid if close to the object, not already holding something, and the object is not in a closed receptacle. +- Put down: Parameterized by the name of the object to put down to a nearby receptacle. Only valid if holding an object. +- Drop: Parameterized by the name of the object to put down. Different from 'put down' as this does not guarantee the held object will be put into a specified receptacle. +- Open: Parameterized by the name of the receptacle to open. Only valid if the receptacle is closed and close to the receptacle. +- Close: Parameterized by the name of the receptacle to close. Only valid if the receptacle is open and close to the receptacle. +- Turn on: Parameterized by the name of the object to turn on. Only valid if the object is turned off and close to the object. +- Turn off: Parameterized by the name of the object to turn off. Only valid if the object is turned on and close to the object. +- Slice: Parameterized by the name of the object to slice. Only valid if the object is sliceable and close to the object. + +## Guidelines +1. Output a plan of actions. Each plan should include no more than 20 actions. +2. Always locate an object using 'find' before interacting with it. +3. Make sure to match the action name and its corresponding action id in the output. Use 'put down' rather than 'drop' to place objects in specific receptacles. +4. Do not repeat the same failed action sequence. Try to modify the action sequence because previous actions did not lead to success. +5. Objects may have multiple instances (e.g., Cabinet_2, Cabinet_3). Explore different instances if needed. +6. Use environment feedback to refine your plan. If an action fails, reflect on the reason and adjust accordingly.""" + + if add_task_examples and TASK_EXAMPLES: + base += "\n\n## Task Examples" + for i, ex in enumerate(TASK_EXAMPLES): + actions_str = ", ".join(ex["actions"]) + base += f"\n\nExample {i+1}: {ex['task']}\n{ex['think']}\n{actions_str}" + + if task_instruction is not None: + base += f"\n\n## Current Task\n{task_instruction}" + + if action_list is not None: + actions_str = "\n".join(f"action id {i}: {a}" for i, a in enumerate(action_list)) + base += f"\n\n## Available Actions (0~{len(action_list) - 1})\n{actions_str}" + + return base + + +def init_observation_template(img_str): + """Template for initial observation after reset. + + Task instruction and available actions are now in the system prompt, + so the initial observation only contains the image. + """ + return f"""[Current Observation]: +{img_str} + +Decide your next action.""" + + +def action_template(last_action, env_feedback, img_str): + """Template for step observation with feedback. + + Encourages structured reasoning: describe what you see, + reflect on why the last action succeeded or failed, + then plan your next actions. + """ + return f"""[Last Action]: {last_action} +[Feedback]: {env_feedback} + +[Current Observation]: +{img_str} + +Describe what you see, reflect on the feedback, and plan your next actions.""" + + +def format_prompt(max_actions_per_step, action_sep, add_example=True, prompt_format="free_think"): + """Generate format prompt based on the specified format.""" + if prompt_format == "free_think": + return free_think_format_prompt(max_actions_per_step, action_sep, add_example) + elif prompt_format == "wm": + return wm_format_prompt(max_actions_per_step, action_sep, add_example) + else: + raise ValueError(f"Unknown prompt format: {prompt_format}") + + +def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True): + """Generate format prompt for free_think format.""" + if max_actions_per_step == 1: + base = """You should output 1 action at a time. +Output the action name exactly as listed in the available actions, or the action ID (integer). +Your response should be in the format of: +...action name or action ID""" + else: + base = f"""You should output a plan of up to {max_actions_per_step} actions at a time, separated by "{action_sep}". +Output the action name exactly as listed in the available actions, or the action ID (integer). +Your response should be in the format of: +...action1{action_sep} action2{action_sep} ...""" + + if add_example: + if max_actions_per_step == 1: + examples = """ +Example 1: +I need to find a mug first. Let me navigate to where mugs might be. +find a Mug + +Example 2: +The mug is nearby and I'm not holding anything. I should pick it up. +pick up the Mug + +Example 3: +I'm holding the mug and I'm near the table. Let me put it down. +put down the object in hand""" + else: + examples = f""" +Example 1 (multi-step plan): +I need to find the alarm clock, pick it up, then find the desk lamp and turn it on. +find a AlarmClock{action_sep} pick up the AlarmClock{action_sep} find a DeskLamp{action_sep} turn on the DeskLamp + +Example 2 (single action when unsure): +I am not sure where the mug is. Let me find it first. +find a Mug + +Example 3 (replanning after failure): +The last action failed because the cabinet was closed. I need to open it first, then pick up the object. +open the Cabinet{action_sep} pick up the Mug""" + return base + "\n" + examples + + return base + + +def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): + """Generate format prompt for wm format with observation and prediction tags.""" + base = f"""You should output {max_actions_per_step} action(s) at a time. +Output the action name exactly as listed in the available actions, or the action ID (integer). +Your response must be in the format of: +......action name or action ID.... + +Rules for : +- Describe the current scene: what objects you see, your position, what you are holding, and relevant receptacle states. + +Rules for : +- Predict what will change after your action: where you will be, what you will see, and the expected result. + +Rules for : +- Output exactly 1 action name or action ID.""" + + if add_example: + examples = """ +Example 1: +I see a kitchen with a counter, a microwave, and a mug on the counter. I am not holding anything. +I need to pick up the mug. First, I should find it to get close to it. +find a Mug +I will navigate to the mug and see it up close on the counter. + +Example 2: +I am close to a Mug on the counter. I am not holding anything. The mug is within reach. +The mug is nearby and I'm not holding anything. I should pick it up. +pick up the Mug +I will be holding the mug. The counter will no longer have the mug on it. + +Example 3: +I am holding a Mug. I see a table nearby with an empty spot. +I'm holding the mug and I'm near the table. Let me put it down. +put down the object in hand +The mug will be placed on the table. I will no longer be holding anything.""" + return base + "\n" + examples + + return base diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py new file mode 100644 index 000000000..de104a063 --- /dev/null +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -0,0 +1,154 @@ +import re +from typing import Dict, List, Optional +from PIL import Image +import numpy as np + + +def parse_free_think(response: str, action_sep: str = ",", max_actions: int = 1) -> Dict: + """ + Parse free_think format response: ...... + + For EB-ALFRED, the tag typically contains a single action name + (e.g., "find a Cabinet") or an action ID (e.g., "42"). + """ + pattern = r'(.*?)\s*(.*?)' + match = re.search(pattern, response, re.DOTALL) + + format_correct = match is not None + + if not match: + think_content = "" + action_content = "" + actions = [] + else: + think_content = match.group(1).strip() + action_content = match.group(2).strip() + + if max_actions == 1: + actions = [action_content.strip()] if action_content.strip() else [] + else: + actions = [a.strip() for a in action_content.split(action_sep) if a.strip()] + if len(actions) > max_actions: + actions = actions[:max_actions] + action_content = action_sep.join(actions) + + llm_response = f"{think_content}{action_content}" + + return { + "llm_raw_response": response, + "llm_response": llm_response, + "think_content": think_content, + "action_content": action_content, + "actions": actions, + "format_correct": format_correct, + } + + +def parse_wm(response: str, action_sep: str = ",", max_actions: int = 1) -> Dict: + """ + Parse wm format response: + ... + ... + ... + ... + """ + pattern = ( + r'(.*?)\s*' + r'(.*?)\s*' + r'(.*?)\s*' + r'(.*?)' + ) + + match = re.search(pattern, response, re.DOTALL) + format_correct = match is not None + + if not match: + observation_content = "" + think_content = "" + prediction_content = "" + action_content = "" + actions: List[str] = [] + else: + observation_content = match.group(1).strip() + think_content = match.group(2).strip() + action_content = match.group(3).strip() + prediction_content = match.group(4).strip() + + if max_actions == 1: + actions = [action_content.strip()] if action_content.strip() else [] + else: + actions = [a.strip() for a in action_content.split(action_sep) if a.strip()] + if len(actions) > max_actions: + actions = actions[:max_actions] + action_content = action_sep.join(actions) + + llm_response = ( + f"{observation_content}" + f"{think_content}" + f"{action_content}" + f"{prediction_content}" + ) + + reasoning_content = think_content + + return { + "llm_raw_response": response, + "llm_response": llm_response, + "observation_content": observation_content, + "think_content": think_content, + "reasoning_content": reasoning_content, + "prediction_content": prediction_content, + "action_content": action_content, + "actions": actions, + "format_correct": format_correct, + } + + +def parse_response( + response: str, + prompt_format: str = "free_think", + action_sep: str = ",", + max_actions: int = 1, +) -> Dict: + """Parse LLM response based on the specified prompt format.""" + if prompt_format == "free_think": + return parse_free_think(response, action_sep, max_actions) + elif prompt_format == "wm": + return parse_wm(response, action_sep, max_actions) + else: + raise ValueError(f"Unknown prompt format: {prompt_format}") + + +def match_action( + action_name: str, + action_list: List[str], + action_map: Dict[str, str], +) -> Optional[str]: + """ + Match a parsed action against the valid action set. + + Supports two formats: + - Action name (case-insensitive): "find a Cabinet" + - Action ID (integer): "42" + + Returns the original action string if matched, None otherwise. + """ + name = action_name.strip() + + # Try as integer action ID + try: + idx = int(name) + if 0 <= idx < len(action_list): + return action_list[idx] + except ValueError: + pass + + # Try exact match by name (case-insensitive) + return action_map.get(name.lower()) + + +def numpy_to_pil(numpy_array: np.ndarray) -> Image.Image: + """Convert numpy (H, W, 3) to PIL.Image in RGB.""" + if numpy_array.shape[-1] == 3: + return Image.fromarray(numpy_array.astype(np.uint8), mode="RGB") + raise ValueError(f"Unsupported channels: {numpy_array.shape[-1]}. Expected 3 (RGB).") From 1ef03b91f3eb8761ca4fea4b981a63c7057e3d05 Mon Sep 17 00:00:00 2001 From: YaningDylan Date: Mon, 9 Mar 2026 08:24:28 +0000 Subject: [PATCH 02/29] Update eb_alfred: capacity queuing, thread-safe GPU display, close timeout; fix envs_remote aclose parallel shutdown --- vagen/envs/eb_alfred/README.md | 163 ++++++++++++++ vagen/envs/eb_alfred/eb_alfred_env.py | 82 +++++-- vagen/envs/eb_alfred/handler.py | 303 +++++++++++++++++++++++++- vagen/envs/eb_alfred/serve.py | 22 ++ 4 files changed, 545 insertions(+), 25 deletions(-) create mode 100644 vagen/envs/eb_alfred/README.md diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md new file mode 100644 index 000000000..32dcc07a9 --- /dev/null +++ b/vagen/envs/eb_alfred/README.md @@ -0,0 +1,163 @@ +# EB-ALFRED Environment — Setup & Run Guide + +EB-ALFRED integrates [EmbodiedBench](https://github.com/EmbodiedBench/EmbodiedBench)'s +AI2-THOR household tasks into the VAGEN framework. + +--- + +## Installation (First-Time Only) + +### 1. Environment Installation + +```bash +cd ERA-rl/VAGEN/vagen/envs/eb_alfred/Embench_new +conda env create -f conda_envs/environment.yaml +conda activate embench +pip install -e . +``` + +### 2. Additional Installation + +Download the dataset from HuggingFace: + +```bash +conda activate embench +git clone https://huggingface.co/datasets/EmbodiedBench/EB-ALFRED +mv EB-ALFRED embodiedbench/envs/eb_alfred/data/json_2.1.0 +``` + +--- + +## One-Time Setup (per machine restart) + +These steps must be done **once per machine boot** before any evaluation or training run. + +### 1. Start GPU X Server + +AI2-THOR (Unity) requires a real display. We use a GPU-accelerated Xorg: + +```bash +# GPU 0 (PCI:1:0:0) → Display :2 +Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ + -config /tmp/xorg.conf :2 & +``` + +> **Note:** `/tmp/xorg.conf` is created during machine setup. Its contents: +> ``` +> Section "Device" +> Identifier "Device0" +> Driver "nvidia" +> BusID "PCI:1:0:0" +> EndSection +> Section "Screen" +> Identifier "Screen0" +> Device "Device0" +> DefaultDepth 24 +> Option "AllowEmptyInitialConfiguration" "True" +> SubSection "Display" +> Depth 24 +> Virtual 1024 768 +> EndSubSection +> EndSection +> Section "ServerLayout" +> Identifier "Layout0" +> Screen 0 "Screen0" 0 0 +> EndSection +> ``` +> If `/tmp/xorg.conf` is missing (e.g. after reboot), recreate it with the above. + +> **Alternative (slower):** Software rendering with Xvfb: +> ```bash +> Xvfb :1 -screen 0 1024x768x24 -ac & +> # Then use DISPLAY=:1 everywhere below +> ``` +> Xvfb is ~7× slower per reset but works without NVIDIA Xorg drivers. + +### 2. Start EB-ALFRED Environment Server + +The server manages AI2-THOR processes and load-balances sessions across GPUs. + +```bash +export PATH="/opt/miniforge3/bin:/usr/bin:/bin:$PATH" + +DISPLAY=:2 conda run -n embench \ + python -m vagen.envs.eb_alfred.serve \ + --port 8000 \ + --x-displays 2 +``` + +> **`--x-displays 2`** is required. Without it, the server auto-detects GPU indices +> (0, 1, ...) from `nvidia-smi` and tries displays `:0`, `:1`, which may not exist. +> Always specify the actual display number explicitly. + +Wait until you see: +``` +Starting EB-ALFRED service on 0.0.0.0:8000 +GPU displays: [:2] (auto-balanced) +Health check: http://localhost:8000/health +``` + +You can verify the server is up: +```bash +curl http://localhost:8000/health +``` + +--- + +## Running Evaluations + +With the server running, launch evaluations from a **separate terminal**: + +```bash +export PATH="/opt/miniforge3/bin:/usr/bin:/bin:$PATH" + +conda run -n vagen \ + python -m vagen.evaluate.run_eval \ + --config tests/.yaml +``` + +### Available Configs + +| Config | Episodes | Concurrency | Resolution | Notes | +|--------|:--------:|:-----------:|:----------:|-------| +| `tests/eval_eb_alfred_gpt41_20ep.yaml` | 20 | 3 | 300 | Quick reference run | +| `tests/eval_eb_alfred_gpt41_10ep_serial_500.yaml` | 10 | 1 (serial) | 500 | Serial baseline | +| `tests/eval_eb_alfred_gpt41_128ep_parallel_500.yaml` | 128 | 100 (parallel) | 500 | High-throughput parallel | + +--- + +## Environment Notes + +### Conda Environments + +| Env | Python | Purpose | +|-----|--------|---------| +| `embench` | 3.9 | AI2-THOR + EmbodiedBench (env server) | +| `vagen` | 3.10 | VAGEN framework (eval runner, training) | + +### Key Constraints + +- **Flask 1.1.4 + Werkzeug 1.0.1** must stay pinned in `embench` — newer versions break AI2-THOR's socket server. +- **GPU Xorg** requires the nvidia driver version to exactly match the kernel module (`580.95.05`). Do **not** `apt upgrade` the nvidia driver without redoing the `xorg.conf` setup. +- **Only one GPU can host an X server** in this container environment. All Unity instances share Xorg `:2` (GPU 0). GPU 1 is available for CUDA workloads but not for display. +- **Multiple concurrent Unity instances** are fine on the same Xorg. 20+ instances tested successfully; ~100 is feasible (CPU-bound at ~150% per instance on 160-core machine). + +### Speed Reference (this machine, 300×300) + +| Backend | Avg Reset | Avg Step | +|---------|:---------:|:--------:| +| GPU Xorg :2 (NVIDIA) | 0.6s | 0.039s | +| Xvfb :1 (software) | 4.3s | 0.211s | + +--- + +## Startup Checklist + +Before every eval run: + +- [ ] `ps aux | grep Xorg` — confirm Xorg :2 is running +- [ ] `curl http://localhost:8000/health` — confirm env server is up +- [ ] `ps aux | grep thor` — no leftover Unity processes from previous runs +- [ ] `nvidia-smi` — GPU memory is mostly free + +If Xorg or the server crashed, restart them per steps 1 & 2 above. diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py index 20eed508b..2d7c957d9 100644 --- a/vagen/envs/eb_alfred/eb_alfred_env.py +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -10,6 +10,8 @@ import asyncio import os +import signal +import threading import numpy as np from PIL import Image from dataclasses import dataclass, field @@ -25,6 +27,37 @@ from vagen.envs.gym_image_env import GymImageEnv +# Thread-local storage for passing x_display to ThorConnector without +# touching the process-global X_DISPLAY module variable. +# Each worker thread sets its own _tl.x_display before creating EBAlfEnv; +# the monkey-patched ThorConnector.__init__ reads it instead of the global. +_tl = threading.local() +_patched = False +_patch_lock = threading.Lock() + + +def _ensure_thor_patched(): + """One-time monkey-patch: make ThorConnector read display from thread-local.""" + global _patched + if _patched: + return + with _patch_lock: + if _patched: + return + from embodiedbench.envs.eb_alfred.thor_connector import ThorConnector + + _orig_init = ThorConnector.__init__ + + def _patched_init(self, x_display=None, **kwargs): + # Prefer thread-local display (set by EbAlfred.__init__) + tl_display = getattr(_tl, "x_display", None) + if tl_display is not None: + x_display = tl_display + _orig_init(self, x_display=x_display, **kwargs) + + ThorConnector.__init__ = _patched_init + _patched = True + @dataclass class EbAlfredEnvConfig: @@ -78,19 +111,24 @@ def __init__(self, env_config: Dict[str, Any]): filtered = {k: v for k, v in env_config.items() if k in valid_keys} self.config = EbAlfredEnvConfig(**filtered) - # Set X display before importing/creating EBAlfEnv - import embodiedbench.envs.eb_alfred.EBAlfEnv as ebalfenv_mod - ebalfenv_mod.X_DISPLAY = self.config.x_display + # Patch ThorConnector to read x_display from thread-local storage + # instead of the process-global X_DISPLAY. This allows fully + # parallel env creation across GPUs with no locks. + _ensure_thor_patched() from embodiedbench.envs.eb_alfred.EBAlfEnv import EBAlfEnv - self.env = EBAlfEnv( - eval_set=self.config.eval_set, - exp_name=self.config.exp_name, - down_sample_ratio=self.config.down_sample_ratio, - selected_indexes=self.config.selected_indexes, - detection_box=self.config.detection_box, - resolution=self.config.resolution, - ) + _tl.x_display = self.config.x_display + try: + self.env = EBAlfEnv( + eval_set=self.config.eval_set, + exp_name=self.config.exp_name, + down_sample_ratio=self.config.down_sample_ratio, + selected_indexes=self.config.selected_indexes, + detection_box=self.config.detection_box, + resolution=self.config.resolution, + ) + finally: + _tl.x_display = None # Adapter state (reset per episode) self._total_turns: int = 0 @@ -105,8 +143,23 @@ def __init__(self, env_config: Dict[str, Any]): # ------------------------------------------------------------------ async def close(self) -> None: - """Close AI2-THOR process.""" - await asyncio.to_thread(self.env.close) + """Close AI2-THOR process. + + Applies a 30-second timeout so that a hung Unity process or + WSGI server shutdown does not block the event loop forever. + """ + try: + await asyncio.wait_for( + asyncio.to_thread(self.env.close), timeout=30.0 + ) + except asyncio.TimeoutError: + # Force-kill the Unity process if graceful shutdown hangs + pid = getattr(self.env.env, "unity_pid", None) + if pid: + try: + os.kill(pid, signal.SIGKILL) + except OSError: + pass async def system_prompt(self) -> Dict[str, Any]: """ @@ -186,8 +239,7 @@ async def step( reward = 0.0 done = False - info: Dict[str, Any] = {} - info.update(parsed) + info = dict(parsed) actions = parsed.get("actions", []) format_correct = parsed.get("format_correct", False) diff --git a/vagen/envs/eb_alfred/handler.py b/vagen/envs/eb_alfred/handler.py index 7d99cf88b..7cd30b8b9 100644 --- a/vagen/envs/eb_alfred/handler.py +++ b/vagen/envs/eb_alfred/handler.py @@ -4,19 +4,40 @@ This is the only component that needs customization. It implements create_env() to instantiate EB-ALFRED environments with automatic multi-GPU load balancing. + +Capacity control: + When ``capacity`` is set (> 0), at most that many Unity environments + run concurrently. Extra ``/connect`` requests are accepted immediately + (returning session_id so the client does NOT retry) and queued. + Environments are created in the background as slots free up -- each + independently, not in batches. """ import asyncio import logging +import random import subprocess +import time +import uuid +from dataclasses import dataclass, field from typing import Any, Dict, List, Optional -from vagen.envs_remote.handler import BaseGymHandler +from vagen.envs_remote.handler import BaseGymHandler, HandlerResult, SessionContext from .eb_alfred_env import EbAlfred LOGGER = logging.getLogger(__name__) +@dataclass +class _DeferredSessionContext(SessionContext): + """SessionContext with extra fields for capacity-based deferred env creation.""" + + env_config: Dict[str, Any] = field(default_factory=dict) + _ready: Optional[asyncio.Event] = field(default=None, repr=False) + _holds_slot: bool = field(default=False, repr=False) + _error: Optional[str] = field(default=None, repr=False) + + def detect_gpu_displays() -> List[str]: """Auto-detect available GPUs via nvidia-smi, return display list. @@ -44,36 +65,79 @@ def detect_gpu_displays() -> List[str]: class EbAlfredHandler(BaseGymHandler): - """Handler for EB-ALFRED with automatic multi-GPU load balancing. + """Handler for EB-ALFRED with capacity-based queuing and multi-GPU load balancing. + + When capacity > 0: + - /connect returns session_id immediately (env creation is deferred) + - A background task waits for a capacity slot, then creates the env + - /call (reset/step) blocks until the env is ready + - /call (close) releases the slot so the next queued env can start + - Each env is independent: no batch waiting + + When capacity = 0 (default): + - Original behaviour: env is created synchronously on /connect - By default, auto-detects available GPUs and distributes new - sessions to the least-loaded GPU. Single-GPU is just the - special case where only one GPU is detected. + startup_concurrency controls how many Unity processes may be in the + startup phase simultaneously, independent of capacity. This prevents + a "startup storm" when many sessions are queued and capacity slots + become available at the same time. E.g. capacity=64, + startup_concurrency=8 means up to 64 envs run concurrently but at + most 8 are initialising at any given moment. """ def __init__( self, x_displays: Optional[List[str]] = None, + capacity: int = 16, + startup_concurrency: int = 8, **kwargs, ): """ Args: x_displays: List of X display IDs to use (e.g. ["0", "1"]). None = auto-detect GPUs via nvidia-smi. + capacity: Max concurrently running Unity environments (0 = unlimited). + startup_concurrency: Max Unity processes that may be starting up at + once (0 = unlimited). Prevents CPU spikes when many capacity + slots open simultaneously. Ignored when capacity = 0. **kwargs: Passed to BaseGymHandler (session_timeout, max_sessions). """ super().__init__(**kwargs) self._x_displays = x_displays if x_displays is not None else detect_gpu_displays() - LOGGER.info(f"[Handler] Using X displays: {self._x_displays}") + self._pending_counts: Dict[str, int] = {d: 0 for d in self._x_displays} + self._capacity = capacity + self._startup_concurrency = startup_concurrency + # Defer semaphore creation: it must be created on the running event loop, + # not during __init__ (which runs before uvicorn starts the loop). + self._capacity_sem: Optional[asyncio.Semaphore] = None + self._startup_sem: Optional[asyncio.Semaphore] = None + LOGGER.info( + f"[Handler] Using X displays: {self._x_displays}, " + f"capacity={capacity if capacity > 0 else 'unlimited'}, " + f"startup_concurrency={startup_concurrency if startup_concurrency > 0 else 'unlimited'}" + ) + + def _ensure_semaphore(self) -> None: + """Lazily create the capacity and startup semaphores on the running event loop.""" + if self._capacity_sem is None and self._capacity > 0: + self._capacity_sem = asyncio.Semaphore(self._capacity) + if self._startup_sem is None and self._startup_concurrency > 0 and self._capacity > 0: + self._startup_sem = asyncio.Semaphore(self._startup_concurrency) def _least_loaded_display(self) -> str: - """Pick the display with the fewest active sessions.""" - counts = {d: 0 for d in self._x_displays} + """Pick the display with the fewest active + pending sessions. + + On ties, randomly choose among the least-loaded displays to + avoid always funnelling to the first GPU. + """ + counts = {d: self._pending_counts.get(d, 0) for d in self._x_displays} for ctx in self._sessions.values(): d = getattr(ctx.env, "_assigned_display", None) if d in counts: counts[d] += 1 - chosen = min(counts, key=counts.get) + min_count = min(counts.values()) + candidates = [d for d, c in counts.items() if c == min_count] + chosen = random.choice(candidates) LOGGER.debug(f"[Handler] GPU load: {counts}, assigning display :{chosen}") return chosen @@ -84,12 +148,231 @@ async def create_env(self, env_config: Dict[str, Any]) -> Any: AI2-THOR startup is blocking, so we offload to a thread. """ display = self._least_loaded_display() + self._pending_counts[display] = self._pending_counts.get(display, 0) + 1 env_config = {**env_config, "x_display": display} - env = await asyncio.to_thread(EbAlfred, env_config) + try: + env = await asyncio.to_thread(EbAlfred, env_config) + finally: + self._pending_counts[display] = max(0, self._pending_counts.get(display, 1) - 1) + env._assigned_display = display LOGGER.info( f"[Handler] Created env on display :{display} " f"(GPU load: { {d: sum(1 for c in self._sessions.values() if getattr(c.env, '_assigned_display', None) == d) for d in self._x_displays} })" ) return env + + # ------------------------------------------------------------------ + # Capacity-aware connect / call / close + # ------------------------------------------------------------------ + + async def connect( + self, env_config: Dict[str, Any], seed: Optional[int] = None + ) -> HandlerResult: + """Accept session immediately; defer env creation if capacity-limited.""" + self._ensure_semaphore() + # No capacity limit → use original behaviour + if self._capacity_sem is None: + return await super().connect(env_config, seed=seed) + + # Check total session limit + if self.max_sessions > 0 and len(self._sessions) >= self.max_sessions: + raise RuntimeError( + f"Max sessions limit reached ({self.max_sessions}). " + f"Please try again later or close existing sessions." + ) + + session_id = uuid.uuid4().hex + ready_event = asyncio.Event() + ctx = _DeferredSessionContext( + session_id=session_id, + env=None, + created_at=time.time(), + last_access=time.time(), + env_config=env_config, + _ready=ready_event, + ) + self._sessions[session_id] = ctx + + # Start cleanup task if not running + if self._cleanup_task is None or self._cleanup_task.done(): + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + + # Fire-and-forget: wait for slot → create env + asyncio.create_task(self._deferred_create(ctx)) + + n_active = sum(1 for s in self._sessions.values() if s.env is not None) + n_queued = len(self._sessions) - n_active + # Estimate wait: (queued_ahead / capacity) * avg_episode_time + # Use a rough estimate of 15s per env creation cycle + estimated_wait = max(0, (n_queued - 1)) / max(1, self._capacity) * 15 + + LOGGER.info( + f"[Handler] Session {session_id} queued " + f"(active={n_active}, queued={n_queued}, capacity={self._capacity}, " + f"est_wait={estimated_wait:.0f}s)" + ) + + return HandlerResult(data={ + "session_id": session_id, + "status": "queued", + "estimated_wait_s": estimated_wait, + }) + + async def _deferred_create(self, ctx: _DeferredSessionContext) -> None: + """Background task: acquire capacity slot, then create env. + + Two-phase acquisition: + 1. capacity_sem – limits total running envs (held for env lifetime) + 2. startup_sem – limits concurrent Unity startups (held only during + EbAlfred.__init__, released as soon as the process + is running) + This prevents a "startup storm" when many capacity slots open at once. + """ + try: + LOGGER.info(f"[Handler] Session {ctx.session_id} waiting for capacity slot...") + await self._capacity_sem.acquire() + ctx._holds_slot = True + LOGGER.info(f"[Handler] Session {ctx.session_id} acquired capacity slot, waiting for startup slot...") + + if self._startup_sem is not None: + await self._startup_sem.acquire() + + LOGGER.info(f"[Handler] Session {ctx.session_id} starting Unity...") + try: + ctx.env = await self.create_env(ctx.env_config) + finally: + if self._startup_sem is not None: + self._startup_sem.release() + + LOGGER.info(f"[Handler] Session {ctx.session_id} env ready") + except Exception as e: + LOGGER.error(f"[Handler] Session {ctx.session_id} env creation failed: {e}") + ctx._error = str(e) + if ctx._holds_slot: + self._capacity_sem.release() + ctx._holds_slot = False + finally: + ctx._ready.set() + + async def _wait_env_ready(self, ctx: _DeferredSessionContext) -> None: + """Block until env is created (called by call() before dispatching).""" + if ctx._ready is not None and not ctx._ready.is_set(): + LOGGER.info(f"[Handler] Session {ctx.session_id} caller waiting for env...") + await ctx._ready.wait() + if ctx.env is None: + error = ctx._error or "Environment creation failed" + raise RuntimeError(f"Session {ctx.session_id}: {error}") + + async def call( + self, + session_id: str, + method: str, + params: Dict[str, Any], + images, + ) -> HandlerResult: + """Dispatch method call; wait for env if still queued.""" + if session_id not in self._sessions: + raise ValueError(f"Session {session_id} not found") + + ctx = self._sessions[session_id] + ctx.last_access = time.time() + + # Wait for env to be ready (no-op if capacity=0 / already ready) + if ctx.env is None and method != "close": + await self._wait_env_ready(ctx) + + return await super().call(session_id, method, params, images) + + async def _handle_close(self, ctx: SessionContext) -> HandlerResult: + """Close env and release capacity slot.""" + try: + if ctx.env is not None: + await ctx.env.close() + except Exception as e: + LOGGER.error(f"[Handler] Error closing env for session {ctx.session_id}: {e}") + finally: + if ctx._holds_slot and self._capacity_sem is not None: + self._capacity_sem.release() + ctx._holds_slot = False + self._sessions.pop(ctx.session_id, None) + + n_active = sum(1 for s in self._sessions.values() if s.env is not None) + n_queued = len(self._sessions) - n_active + LOGGER.info( + f"[Handler] Closed session {ctx.session_id} " + f"(active={n_active}, queued={n_queued}, capacity={self._capacity})" + ) + return HandlerResult(data={"closed": True}) + + def get_session_stats(self) -> Dict[str, Any]: + """Session stats with active/queued breakdown.""" + stats = super().get_session_stats() + n_active = sum(1 for s in self._sessions.values() if s.env is not None) + stats["active"] = n_active + stats["queued"] = len(self._sessions) - n_active + stats["capacity"] = self._capacity if self._capacity > 0 else "unlimited" + for s in stats.get("sessions", []): + sid = s["session_id"] + ctx = self._sessions.get(sid) + s["status"] = "active" if (ctx and ctx.env is not None) else "queued" + return stats + + async def _cleanup_loop(self): + """Cleanup timed-out sessions, releasing capacity slots.""" + while True: + try: + await asyncio.sleep(60) + now = time.time() + to_remove = [] + for session_id, ctx in self._sessions.items(): + if now - ctx.last_access > self.session_timeout: + to_remove.append(session_id) + LOGGER.warning(f"[Handler] Session {session_id} timed out") + + for session_id in to_remove: + ctx = self._sessions.get(session_id) + if ctx is None: + continue + try: + if ctx.env is not None: + await ctx.env.close() + except Exception as e: + LOGGER.error(f"[Handler] Cleanup error {session_id}: {e}") + finally: + if ctx._holds_slot and self._capacity_sem is not None: + self._capacity_sem.release() + ctx._holds_slot = False + self._sessions.pop(session_id, None) + except asyncio.CancelledError: + break + except Exception as e: + LOGGER.error(f"[Handler] Cleanup loop error: {e}") + + async def aclose(self): + """Shutdown: close all sessions, release all capacity slots.""" + if self._cleanup_task and not self._cleanup_task.done(): + self._cleanup_task.cancel() + try: + await self._cleanup_task + except asyncio.CancelledError: + pass + + async def _close_one(sid: str, ctx: SessionContext): + try: + if ctx.env is not None: + await ctx.env.close() + except Exception as e: + LOGGER.error(f"[Handler] Shutdown close error {sid}: {e}") + finally: + if ctx._holds_slot and self._capacity_sem is not None: + self._capacity_sem.release() + ctx._holds_slot = False + + if self._sessions: + await asyncio.gather( + *(_close_one(sid, ctx) for sid, ctx in self._sessions.items()) + ) + self._sessions.clear() + LOGGER.info("[Handler] All sessions closed") diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index d422370f2..db59762ea 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -53,6 +53,22 @@ def main(): default=0, help="Max concurrent sessions (0=unlimited)", ) + parser.add_argument( + "--capacity", + type=int, + default=16, + help="Max concurrently running Unity environments (0=unlimited). " + "Extra sessions are queued and created as slots free up.", + ) + parser.add_argument( + "--startup-concurrency", + type=int, + default=8, + help="Max Unity processes that may be starting up simultaneously (0=unlimited). " + "Prevents CPU spikes when many capacity slots open at once. " + "E.g. --capacity 64 --startup-concurrency 8 means 64 envs run " + "concurrently but startups are staggered 8 at a time.", + ) parser.add_argument( "--thread-workers", type=int, @@ -67,6 +83,8 @@ def main(): x_displays=x_displays, session_timeout=args.session_timeout, max_sessions=args.max_sessions, + capacity=args.capacity, + startup_concurrency=args.startup_concurrency, ) app = build_gym_service(handler) @@ -82,8 +100,12 @@ async def _set_thread_pool(): ) displays_str = ", ".join(f":{d}" for d in handler._x_displays) + cap_str = str(args.capacity) if args.capacity > 0 else "unlimited" + startup_str = str(args.startup_concurrency) if args.startup_concurrency > 0 else "unlimited" print(f"Starting EB-ALFRED service on {args.host}:{args.port}") print(f"GPU displays: [{displays_str}] (auto-balanced)") + print(f"Capacity: {cap_str} concurrent environments") + print(f"Startup concurrency: {startup_str} simultaneous Unity startups") print(f"Health check: http://localhost:{args.port}/health") uvicorn.run(app, host=args.host, port=args.port) From 14c67f12ac085ec1f2e2e6c472fc6a31a1801205 Mon Sep 17 00:00:00 2001 From: YaningDylan Date: Mon, 16 Mar 2026 02:00:00 +0000 Subject: [PATCH 03/29] update readme and env logic --- vagen/envs/eb_alfred/README.md | 230 ++++++++++++++++---------- vagen/envs/eb_alfred/eb_alfred_env.py | 8 +- 2 files changed, 148 insertions(+), 90 deletions(-) diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 32dcc07a9..58cf0e89c 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -3,27 +3,64 @@ EB-ALFRED integrates [EmbodiedBench](https://github.com/EmbodiedBench/EmbodiedBench)'s AI2-THOR household tasks into the VAGEN framework. +- 301 evaluation episodes across 6 eval sets +- 162 discrete actions (find / pick / put / open / close / slice / toggle) +- GPU-accelerated Xorg required — Xvfb does **not** work (no hardware OpenGL) + --- ## Installation (First-Time Only) -### 1. Environment Installation +### 1. Install EmbodiedBench + +```bash +# Clone EmbodiedBench +git clone https://github.com/EmbodiedBench/EmbodiedBench.git /root/EmbodiedBench + +# REQUIRED: the package is missing __init__.py — editable install breaks without it +touch /root/EmbodiedBench/embodiedbench/__init__.py + +pip install -e /root/EmbodiedBench +``` + +### 2. Install Required Packages (order matters) + +AI2-THOR 2.1.0 has strict version requirements. Install in this order: ```bash -cd ERA-rl/VAGEN/vagen/envs/eb_alfred/Embench_new -conda env create -f conda_envs/environment.yaml -conda activate embench -pip install -e . +# 1. PyTorch (match your CUDA driver; cu126 for driver >= 525) +pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126 + +# 2. ai2thor + core deps +pip install "ai2thor==2.1.0" "gym==0.23.0" "numpy<2.0" \ + scipy Pillow networkx revtok vocab h5py tqdm natsort pyquaternion + +# 3. Pin flask/werkzeug — ai2thor 2.1.0 uses an internal Flask server +# that is incompatible with Flask 2+ / Werkzeug 2+ +pip install "flask==1.1.4" "werkzeug==1.0.1" \ + "markupsafe<2.1" "jinja2<3.0" "itsdangerous<2.0" + +# 4. opencv — must be <4.9 (4.9+ requires numpy>=2, which conflicts with gym 0.23.0) +pip install "opencv-python-headless<4.9" + +# 5. Re-pin numpy (opencv/hydra may have upgraded it) +pip install "numpy<2.0" ``` -### 2. Additional Installation +> **Critical version constraints:** +> | Package | Required | Reason | +> |---------|----------|--------| +> | `flask` | `==1.1.4` | ai2thor 2.1.0 internal web server | +> | `werkzeug` | `==1.0.1` | same — Werkzeug 2+ breaks the socket bridge | +> | `numpy` | `<2.0` | gym 0.23.0 incompatible with numpy 2.x | +> | `opencv-python-headless` | `<4.9` | 4.9+ requires numpy>=2 | +> | `gym` | `==0.23.0` | required by EmbodiedBench | -Download the dataset from HuggingFace: +### 3. Download Dataset ```bash -conda activate embench git clone https://huggingface.co/datasets/EmbodiedBench/EB-ALFRED -mv EB-ALFRED embodiedbench/envs/eb_alfred/data/json_2.1.0 +mv EB-ALFRED /path/to/EmbodiedBench/embodiedbench/envs/eb_alfred/data/json_2.1.0 ``` --- @@ -34,70 +71,100 @@ These steps must be done **once per machine boot** before any evaluation or trai ### 1. Start GPU X Server -AI2-THOR (Unity) requires a real display. We use a GPU-accelerated Xorg: +AI2-THOR (Unity) requires a **GPU-accelerated Xorg** for OpenGL rendering. +**Xvfb will not work** — Unity falls back to CPU rendering and hangs. + +The Xorg config must include a `Monitor` section with a `Modeline` and `Modes` directive. +Without it, Unity sees "Desktop is 0 x 0 @ 0 Hz" and freezes. + +#### Single-GPU setup ```bash -# GPU 0 (PCI:1:0:0) → Display :2 +# Find your GPU BusID (convert hex to decimal: e.g. 0x41 → 65) +nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader + +cat > /tmp/xorg.conf << 'EOF' +Section "ServerFlags" + Option "AllowEmptyInput" "True" +EndSection + +Section "Device" + Identifier "Device0" + Driver "nvidia" + BusID "PCI:1:0:0" # replace with your GPU BusID + Option "AllowEmptyInitialConfiguration" "True" +EndSection + +Section "Monitor" + Identifier "Monitor0" + HorizSync 28.0-80.0 + VertRefresh 48.0-75.0 + Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 + Option "DPMS" +EndSection + +Section "Screen" + Identifier "Screen0" + Device "Device0" + Monitor "Monitor0" + DefaultDepth 24 + Option "AllowEmptyInitialConfiguration" "True" + Option "UseDisplayDevice" "none" + SubSection "Display" + Depth 24 + Modes "1920x1080" + Virtual 1920 1080 + EndSubSection +EndSection + +Section "ServerLayout" + Identifier "Layout0" + Screen 0 "Screen0" 0 0 +EndSection +EOF + Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ - -config /tmp/xorg.conf :2 & + -config /tmp/xorg.conf :1 &>/tmp/xorg1.log & + +# Verify resolution was detected +grep "Virtual screen size" /var/log/Xorg.1.log +# Expected: "Virtual screen size configured to be 1920 x 1080" ``` -> **Note:** `/tmp/xorg.conf` is created during machine setup. Its contents: -> ``` -> Section "Device" -> Identifier "Device0" -> Driver "nvidia" -> BusID "PCI:1:0:0" -> EndSection -> Section "Screen" -> Identifier "Screen0" -> Device "Device0" -> DefaultDepth 24 -> Option "AllowEmptyInitialConfiguration" "True" -> SubSection "Display" -> Depth 24 -> Virtual 1024 768 -> EndSubSection -> EndSection -> Section "ServerLayout" -> Identifier "Layout0" -> Screen 0 "Screen0" 0 0 -> EndSection -> ``` -> If `/tmp/xorg.conf` is missing (e.g. after reboot), recreate it with the above. - -> **Alternative (slower):** Software rendering with Xvfb: -> ```bash -> Xvfb :1 -screen 0 1024x768x24 -ac & -> # Then use DISPLAY=:1 everywhere below -> ``` -> Xvfb is ~7× slower per reset but works without NVIDIA Xorg drivers. +#### Dual-GPU setup (one Xorg per GPU) -### 2. Start EB-ALFRED Environment Server +```bash +# GPU 0 → display :0 GPU 1 → display :1 +# Write a single-GPU config for each (see above template, change BusID/Identifier) -The server manages AI2-THOR processes and load-balances sessions across GPUs. +Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ + -config /tmp/xorg_gpu0.conf :0 &>/tmp/xorg0.log & -```bash -export PATH="/opt/miniforge3/bin:/usr/bin:/bin:$PATH" +Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ + -config /tmp/xorg_gpu1.conf :1 &>/tmp/xorg1.log & +``` + +### 2. Start EB-ALFRED Environment Server -DISPLAY=:2 conda run -n embench \ - python -m vagen.envs.eb_alfred.serve \ +```bash +python -m vagen.envs.eb_alfred.serve \ --port 8000 \ - --x-displays 2 + --capacity 128 \ + --startup-concurrency 4 \ + --x-displays 1 # single GPU: --x-displays 1 + # dual GPU: --x-displays 0,1 ``` -> **`--x-displays 2`** is required. Without it, the server auto-detects GPU indices -> (0, 1, ...) from `nvidia-smi` and tries displays `:0`, `:1`, which may not exist. -> Always specify the actual display number explicitly. +> **`--startup-concurrency 4`** staggers Unity process startup to avoid CPU spikes. +> Even with `--capacity 128`, only 4 Unity instances start simultaneously; the rest queue. -Wait until you see: +Wait for: ``` Starting EB-ALFRED service on 0.0.0.0:8000 -GPU displays: [:2] (auto-balanced) -Health check: http://localhost:8000/health +GPU displays: [:0, :1] (auto-balanced) ``` -You can verify the server is up: +Verify: ```bash curl http://localhost:8000/health ``` @@ -106,24 +173,12 @@ curl http://localhost:8000/health ## Running Evaluations -With the server running, launch evaluations from a **separate terminal**: +With the server running, launch from a separate terminal: ```bash -export PATH="/opt/miniforge3/bin:/usr/bin:/bin:$PATH" - -conda run -n vagen \ - python -m vagen.evaluate.run_eval \ - --config tests/.yaml +python -m vagen.evaluate.run_eval --config tests/.yaml ``` -### Available Configs - -| Config | Episodes | Concurrency | Resolution | Notes | -|--------|:--------:|:-----------:|:----------:|-------| -| `tests/eval_eb_alfred_gpt41_20ep.yaml` | 20 | 3 | 300 | Quick reference run | -| `tests/eval_eb_alfred_gpt41_10ep_serial_500.yaml` | 10 | 1 (serial) | 500 | Serial baseline | -| `tests/eval_eb_alfred_gpt41_128ep_parallel_500.yaml` | 128 | 100 (parallel) | 500 | High-throughput parallel | - --- ## Environment Notes @@ -135,29 +190,28 @@ conda run -n vagen \ | `embench` | 3.9 | AI2-THOR + EmbodiedBench (env server) | | `vagen` | 3.10 | VAGEN framework (eval runner, training) | -### Key Constraints +### Known Issues -- **Flask 1.1.4 + Werkzeug 1.0.1** must stay pinned in `embench` — newer versions break AI2-THOR's socket server. -- **GPU Xorg** requires the nvidia driver version to exactly match the kernel module (`580.95.05`). Do **not** `apt upgrade` the nvidia driver without redoing the `xorg.conf` setup. -- **Only one GPU can host an X server** in this container environment. All Unity instances share Xorg `:2` (GPU 0). GPU 1 is available for CUDA workloads but not for display. -- **Multiple concurrent Unity instances** are fine on the same Xorg. 20+ instances tested successfully; ~100 is feasible (CPU-bound at ~150% per instance on 160-core machine). +**Unity hangs with "Desktop is 0 x 0 @ 0 Hz"** +The Xorg config is missing `Monitor` + `Modeline` + `Modes`. The `Virtual` directive alone +is not enough. Use the full config shown above. -### Speed Reference (this machine, 300×300) +**`embodiedbench` import fails after editable install** +The package directory has no `__init__.py`, which breaks pip's editable-install finder. +Fix: `touch embodiedbench/__init__.py` then `pip install -e .` -| Backend | Avg Reset | Avg Step | -|---------|:---------:|:--------:| -| GPU Xorg :2 (NVIDIA) | 0.6s | 0.039s | -| Xvfb :1 (software) | 4.3s | 0.211s | +**`opencv-python-headless` version conflict** +`opencv>=4.9` requires `numpy>=2`, but `gym==0.23.0` requires `numpy<2`. +Fix: `pip install "opencv-python-headless<4.9"` ---- - -## Startup Checklist +**AI2-THOR first-run download** +Unity binary (~390 MB) downloads to `~/.ai2thor/releases/` on first run. +Subsequent runs use the cache. -Before every eval run: +### Startup Checklist -- [ ] `ps aux | grep Xorg` — confirm Xorg :2 is running -- [ ] `curl http://localhost:8000/health` — confirm env server is up +Before every run: +- [ ] `ps aux | grep Xorg` — Xorg is running on the expected display(s) +- [ ] `curl http://localhost:8000/health` — env server is up - [ ] `ps aux | grep thor` — no leftover Unity processes from previous runs - [ ] `nvidia-smi` — GPU memory is mostly free - -If Xorg or the server crashed, restart them per steps 1 & 2 above. diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py index 2d7c957d9..083f5ae11 100644 --- a/vagen/envs/eb_alfred/eb_alfred_env.py +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -195,7 +195,9 @@ async def reset(self, seed: int) -> Tuple[Dict[str, Any], Dict[str, Any]]: episode_idx = seed % self.env.number_of_episodes self.env._current_episode_num = episode_idx - await asyncio.to_thread(self.env.reset) + await asyncio.wait_for( + asyncio.to_thread(self.env.reset), timeout=300.0 + ) # Reset adapter state self._total_turns = 0 @@ -277,7 +279,9 @@ async def step( # Execute in AI2-THOR self._total_env_steps += 1 obs_raw, step_reward, step_done, step_info = ( - await asyncio.to_thread(self.env.step, matched) + await asyncio.wait_for( + asyncio.to_thread(self.env.step, matched), timeout=60.0 + ) ) self._last_action = matched From 1ff784c673ba655ff227b2f547327bad263c3a4c Mon Sep 17 00:00:00 2001 From: YaningDylan Date: Tue, 17 Mar 2026 01:46:22 +0000 Subject: [PATCH 04/29] update service logic --- vagen/envs/eb_alfred/serve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index db59762ea..a0bb6c666 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -86,7 +86,7 @@ def main(): capacity=args.capacity, startup_concurrency=args.startup_concurrency, ) - app = build_gym_service(handler) + app = build_gym_service(handler, max_inflight=args.capacity) # Expand the asyncio thread pool via FastAPI startup so concurrent Unity # startups don't queue behind Python's default limit of min(32, cpu+4). From e27fbe4766cdeca0c48e11f3ff80a30cb9fa6973 Mon Sep 17 00:00:00 2001 From: YaningDylan Date: Tue, 17 Mar 2026 22:00:56 +0000 Subject: [PATCH 05/29] fix service semaphore event loop bug; set max_inflight=0; simplify README --- vagen/envs/eb_alfred/README.md | 180 +++++++-------------------------- vagen/envs/eb_alfred/serve.py | 2 +- 2 files changed, 35 insertions(+), 147 deletions(-) diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 58cf0e89c..23018ea90 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -1,217 +1,105 @@ -# EB-ALFRED Environment — Setup & Run Guide - -EB-ALFRED integrates [EmbodiedBench](https://github.com/EmbodiedBench/EmbodiedBench)'s -AI2-THOR household tasks into the VAGEN framework. - -- 301 evaluation episodes across 6 eval sets -- 162 discrete actions (find / pick / put / open / close / slice / toggle) -- GPU-accelerated Xorg required — Xvfb does **not** work (no hardware OpenGL) - ---- +# EB-ALFRED Environment ## Installation (First-Time Only) -### 1. Install EmbodiedBench - ```bash -# Clone EmbodiedBench git clone https://github.com/EmbodiedBench/EmbodiedBench.git /root/EmbodiedBench - -# REQUIRED: the package is missing __init__.py — editable install breaks without it touch /root/EmbodiedBench/embodiedbench/__init__.py - pip install -e /root/EmbodiedBench -``` - -### 2. Install Required Packages (order matters) - -AI2-THOR 2.1.0 has strict version requirements. Install in this order: -```bash -# 1. PyTorch (match your CUDA driver; cu126 for driver >= 525) pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126 - -# 2. ai2thor + core deps pip install "ai2thor==2.1.0" "gym==0.23.0" "numpy<2.0" \ scipy Pillow networkx revtok vocab h5py tqdm natsort pyquaternion - -# 3. Pin flask/werkzeug — ai2thor 2.1.0 uses an internal Flask server -# that is incompatible with Flask 2+ / Werkzeug 2+ pip install "flask==1.1.4" "werkzeug==1.0.1" \ "markupsafe<2.1" "jinja2<3.0" "itsdangerous<2.0" - -# 4. opencv — must be <4.9 (4.9+ requires numpy>=2, which conflicts with gym 0.23.0) pip install "opencv-python-headless<4.9" - -# 5. Re-pin numpy (opencv/hydra may have upgraded it) pip install "numpy<2.0" -``` - -> **Critical version constraints:** -> | Package | Required | Reason | -> |---------|----------|--------| -> | `flask` | `==1.1.4` | ai2thor 2.1.0 internal web server | -> | `werkzeug` | `==1.0.1` | same — Werkzeug 2+ breaks the socket bridge | -> | `numpy` | `<2.0` | gym 0.23.0 incompatible with numpy 2.x | -> | `opencv-python-headless` | `<4.9` | 4.9+ requires numpy>=2 | -> | `gym` | `==0.23.0` | required by EmbodiedBench | -### 3. Download Dataset - -```bash git clone https://huggingface.co/datasets/EmbodiedBench/EB-ALFRED mv EB-ALFRED /path/to/EmbodiedBench/embodiedbench/envs/eb_alfred/data/json_2.1.0 ``` --- -## One-Time Setup (per machine restart) - -These steps must be done **once per machine boot** before any evaluation or training run. +## Per-Boot Setup -### 1. Start GPU X Server - -AI2-THOR (Unity) requires a **GPU-accelerated Xorg** for OpenGL rendering. -**Xvfb will not work** — Unity falls back to CPU rendering and hangs. - -The Xorg config must include a `Monitor` section with a `Modeline` and `Modes` directive. -Without it, Unity sees "Desktop is 0 x 0 @ 0 Hz" and freezes. - -#### Single-GPU setup +### 1. Start Xorg (GPU-accelerated, Xvfb will not work) ```bash -# Find your GPU BusID (convert hex to decimal: e.g. 0x41 → 65) +# Get GPU BusID (hex → decimal, e.g. 0x41 → 65) nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader -cat > /tmp/xorg.conf << 'EOF' -Section "ServerFlags" - Option "AllowEmptyInput" "True" -EndSection - +cat > /tmp/xorg0.conf << 'EOF' Section "Device" - Identifier "Device0" + Identifier "GPU0" Driver "nvidia" - BusID "PCI:1:0:0" # replace with your GPU BusID + BusID "PCI:65:0:0" Option "AllowEmptyInitialConfiguration" "True" EndSection - Section "Monitor" Identifier "Monitor0" HorizSync 28.0-80.0 VertRefresh 48.0-75.0 Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 - Option "DPMS" EndSection - Section "Screen" Identifier "Screen0" - Device "Device0" + Device "GPU0" Monitor "Monitor0" DefaultDepth 24 - Option "AllowEmptyInitialConfiguration" "True" - Option "UseDisplayDevice" "none" SubSection "Display" Depth 24 Modes "1920x1080" Virtual 1920 1080 EndSubSection EndSection - Section "ServerLayout" Identifier "Layout0" - Screen 0 "Screen0" 0 0 + Screen 0 "Screen0" EndSection EOF -Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ - -config /tmp/xorg.conf :1 &>/tmp/xorg1.log & +# Single GPU +Xorg -noreset +extension GLX -config /tmp/xorg0.conf :0 & -# Verify resolution was detected -grep "Virtual screen size" /var/log/Xorg.1.log -# Expected: "Virtual screen size configured to be 1920 x 1080" -``` - -#### Dual-GPU setup (one Xorg per GPU) - -```bash -# GPU 0 → display :0 GPU 1 → display :1 -# Write a single-GPU config for each (see above template, change BusID/Identifier) +# Dual GPU (repeat with second config for :1) +Xorg -noreset +extension GLX -config /tmp/xorg1.conf :1 & -Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ - -config /tmp/xorg_gpu0.conf :0 &>/tmp/xorg0.log & - -Xorg -noreset +extension GLX +extension RANDR +extension RENDER \ - -config /tmp/xorg_gpu1.conf :1 &>/tmp/xorg1.log & +# Verify +grep "Virtual screen size" /var/log/Xorg.0.log ``` -### 2. Start EB-ALFRED Environment Server +### 2. Start Server ```bash python -m vagen.envs.eb_alfred.serve \ --port 8000 \ - --capacity 128 \ - --startup-concurrency 4 \ - --x-displays 1 # single GPU: --x-displays 1 - # dual GPU: --x-displays 0,1 -``` - -> **`--startup-concurrency 4`** staggers Unity process startup to avoid CPU spikes. -> Even with `--capacity 128`, only 4 Unity instances start simultaneously; the rest queue. - -Wait for: -``` -Starting EB-ALFRED service on 0.0.0.0:8000 -GPU displays: [:0, :1] (auto-balanced) -``` + --capacity 90 \ + --startup-concurrency 6 \ + --x-displays 0,1 -Verify: -```bash curl http://localhost:8000/health ``` ---- - -## Running Evaluations - -With the server running, launch from a separate terminal: +### 3. SSH Tunnel (if training machine is remote) ```bash -python -m vagen.evaluate.run_eval --config tests/.yaml +# On env server +ssh -p -R 8000:localhost:8000 \ + -o ServerAliveInterval=30 \ + -o ServerAliveCountMax=5 \ + -N -f user@training-machine-ip + +# On training machine — set once in /etc/ssh/sshd_config: +# MaxSessions 200 +# then: service ssh reload ``` --- -## Environment Notes - -### Conda Environments - -| Env | Python | Purpose | -|-----|--------|---------| -| `embench` | 3.9 | AI2-THOR + EmbodiedBench (env server) | -| `vagen` | 3.10 | VAGEN framework (eval runner, training) | - -### Known Issues - -**Unity hangs with "Desktop is 0 x 0 @ 0 Hz"** -The Xorg config is missing `Monitor` + `Modeline` + `Modes`. The `Virtual` directive alone -is not enough. Use the full config shown above. - -**`embodiedbench` import fails after editable install** -The package directory has no `__init__.py`, which breaks pip's editable-install finder. -Fix: `touch embodiedbench/__init__.py` then `pip install -e .` - -**`opencv-python-headless` version conflict** -`opencv>=4.9` requires `numpy>=2`, but `gym==0.23.0` requires `numpy<2`. -Fix: `pip install "opencv-python-headless<4.9"` - -**AI2-THOR first-run download** -Unity binary (~390 MB) downloads to `~/.ai2thor/releases/` on first run. -Subsequent runs use the cache. - -### Startup Checklist +## Checklist -Before every run: -- [ ] `ps aux | grep Xorg` — Xorg is running on the expected display(s) -- [ ] `curl http://localhost:8000/health` — env server is up -- [ ] `ps aux | grep thor` — no leftover Unity processes from previous runs -- [ ] `nvidia-smi` — GPU memory is mostly free +- [ ] `ps aux | grep Xorg` +- [ ] `curl http://localhost:8000/health` +- [ ] `ps aux | grep thor` — no leftover Unity processes +- [ ] `nvidia-smi` — GPU memory free diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index a0bb6c666..be1ec2f02 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -86,7 +86,7 @@ def main(): capacity=args.capacity, startup_concurrency=args.startup_concurrency, ) - app = build_gym_service(handler, max_inflight=args.capacity) + app = build_gym_service(handler, max_inflight=0) # Expand the asyncio thread pool via FastAPI startup so concurrent Unity # startups don't queue behind Python's default limit of min(32, cpu+4). From 7ad072d1bd9359c8f6dff732e15f1cc3beb69556 Mon Sep 17 00:00:00 2001 From: YaningDylan Date: Tue, 17 Mar 2026 22:10:50 +0000 Subject: [PATCH 06/29] add start_server.sh and xorg confs; simplify README --- vagen/envs/eb_alfred/README.md | 58 ++-------------------------- vagen/envs/eb_alfred/start_server.sh | 16 ++++++++ vagen/envs/eb_alfred/xorg0.conf | 27 +++++++++++++ vagen/envs/eb_alfred/xorg1.conf | 27 +++++++++++++ 4 files changed, 73 insertions(+), 55 deletions(-) create mode 100755 vagen/envs/eb_alfred/start_server.sh create mode 100644 vagen/envs/eb_alfred/xorg0.conf create mode 100644 vagen/envs/eb_alfred/xorg1.conf diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 23018ea90..4577d389b 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -23,65 +23,13 @@ mv EB-ALFRED /path/to/EmbodiedBench/embodiedbench/envs/eb_alfred/data/json_2.1.0 ## Per-Boot Setup -### 1. Start Xorg (GPU-accelerated, Xvfb will not work) +### 1. Start Xorg + Server ```bash -# Get GPU BusID (hex → decimal, e.g. 0x41 → 65) -nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader - -cat > /tmp/xorg0.conf << 'EOF' -Section "Device" - Identifier "GPU0" - Driver "nvidia" - BusID "PCI:65:0:0" - Option "AllowEmptyInitialConfiguration" "True" -EndSection -Section "Monitor" - Identifier "Monitor0" - HorizSync 28.0-80.0 - VertRefresh 48.0-75.0 - Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 -EndSection -Section "Screen" - Identifier "Screen0" - Device "GPU0" - Monitor "Monitor0" - DefaultDepth 24 - SubSection "Display" - Depth 24 - Modes "1920x1080" - Virtual 1920 1080 - EndSubSection -EndSection -Section "ServerLayout" - Identifier "Layout0" - Screen 0 "Screen0" -EndSection -EOF - -# Single GPU -Xorg -noreset +extension GLX -config /tmp/xorg0.conf :0 & - -# Dual GPU (repeat with second config for :1) -Xorg -noreset +extension GLX -config /tmp/xorg1.conf :1 & - -# Verify -grep "Virtual screen size" /var/log/Xorg.0.log -``` - -### 2. Start Server - -```bash -python -m vagen.envs.eb_alfred.serve \ - --port 8000 \ - --capacity 90 \ - --startup-concurrency 6 \ - --x-displays 0,1 - -curl http://localhost:8000/health +bash vagen/envs/eb_alfred/start_server.sh ``` -### 3. SSH Tunnel (if training machine is remote) +### 2. SSH Tunnel (if training machine is remote) ```bash # On env server diff --git a/vagen/envs/eb_alfred/start_server.sh b/vagen/envs/eb_alfred/start_server.sh new file mode 100755 index 000000000..969869730 --- /dev/null +++ b/vagen/envs/eb_alfred/start_server.sh @@ -0,0 +1,16 @@ +#!/bin/bash +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +cp "$SCRIPT_DIR/xorg0.conf" /tmp/xorg0.conf +cp "$SCRIPT_DIR/xorg1.conf" /tmp/xorg1.conf + +Xorg -noreset +extension GLX -config /tmp/xorg0.conf :0 & +Xorg -noreset +extension GLX -config /tmp/xorg1.conf :1 & + +sleep 2 + +python -m vagen.envs.eb_alfred.serve \ + --port 8000 \ + --capacity 90 \ + --startup-concurrency 6 \ + --x-displays 0,1 diff --git a/vagen/envs/eb_alfred/xorg0.conf b/vagen/envs/eb_alfred/xorg0.conf new file mode 100644 index 000000000..5b5d53691 --- /dev/null +++ b/vagen/envs/eb_alfred/xorg0.conf @@ -0,0 +1,27 @@ +Section "Device" + Identifier "GPU0" + Driver "nvidia" + BusID "PCI:65:0:0" + Option "AllowEmptyInitialConfiguration" "True" +EndSection +Section "Monitor" + Identifier "Monitor0" + HorizSync 28.0-80.0 + VertRefresh 48.0-75.0 + Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 +EndSection +Section "Screen" + Identifier "Screen0" + Device "GPU0" + Monitor "Monitor0" + DefaultDepth 24 + SubSection "Display" + Depth 24 + Modes "1920x1080" + Virtual 1920 1080 + EndSubSection +EndSection +Section "ServerLayout" + Identifier "Layout0" + Screen 0 "Screen0" +EndSection diff --git a/vagen/envs/eb_alfred/xorg1.conf b/vagen/envs/eb_alfred/xorg1.conf new file mode 100644 index 000000000..99e40b539 --- /dev/null +++ b/vagen/envs/eb_alfred/xorg1.conf @@ -0,0 +1,27 @@ +Section "Device" + Identifier "GPU1" + Driver "nvidia" + BusID "PCI:97:0:0" + Option "AllowEmptyInitialConfiguration" "True" +EndSection +Section "Monitor" + Identifier "Monitor1" + HorizSync 28.0-80.0 + VertRefresh 48.0-75.0 + Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 +EndSection +Section "Screen" + Identifier "Screen1" + Device "GPU1" + Monitor "Monitor1" + DefaultDepth 24 + SubSection "Display" + Depth 24 + Modes "1920x1080" + Virtual 1920 1080 + EndSubSection +EndSection +Section "ServerLayout" + Identifier "Layout1" + Screen 0 "Screen1" +EndSection From 7327cfc9132b1a3e941e32e6cf7d57064197f5a8 Mon Sep 17 00:00:00 2001 From: VAGEN Dev Date: Tue, 17 Mar 2026 22:54:22 +0000 Subject: [PATCH 07/29] Update eb_alfred serve.py to use GymService class API main refactored envs_remote/service.py from build_gym_service() function to GymService class; update serve.py accordingly. Co-Authored-By: Claude Sonnet 4.6 --- vagen/envs/eb_alfred/serve.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index be1ec2f02..aae71a9a1 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -26,7 +26,7 @@ import concurrent.futures import uvicorn -from vagen.envs_remote.service import build_gym_service +from vagen.envs_remote.service import GymService from .handler import EbAlfredHandler @@ -86,7 +86,7 @@ def main(): capacity=args.capacity, startup_concurrency=args.startup_concurrency, ) - app = build_gym_service(handler, max_inflight=0) + app = GymService(handler, max_inflight=0).build() # Expand the asyncio thread pool via FastAPI startup so concurrent Unity # startups don't queue behind Python's default limit of min(32, cpu+4). From 64ff6f957ec243b0a5bcd1bc08b378e8711046d7 Mon Sep 17 00:00:00 2001 From: VAGEN Dev Date: Tue, 17 Mar 2026 23:04:45 +0000 Subject: [PATCH 08/29] Auto-detect GPUs and start Xorg in eb_alfred serve.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Mirror navigation's _detect_gpus() (CUDA_VISIBLE_DEVICES or nvidia-smi) - Auto-generate /tmp/xorgN.conf from nvidia-smi PCI bus IDs (hex→decimal) - Auto-start Xorg :N for each GPU N, skip if already running - Switch from argparse to fire.Fire (consistent with navigation) - Rename --x-displays to --devices (List[int], consistent with navigation) - Remove hardcoded xorg0.conf / xorg1.conf (machine-specific BusIDs) - Simplify start_server.sh to a one-liner example Co-Authored-By: Claude Sonnet 4.6 --- vagen/envs/eb_alfred/serve.py | 242 ++++++++++++++++++--------- vagen/envs/eb_alfred/start_server.sh | 16 +- vagen/envs/eb_alfred/xorg0.conf | 27 --- vagen/envs/eb_alfred/xorg1.conf | 27 --- 4 files changed, 168 insertions(+), 144 deletions(-) delete mode 100644 vagen/envs/eb_alfred/xorg0.conf delete mode 100644 vagen/envs/eb_alfred/xorg1.conf diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index aae71a9a1..1327e15a9 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -5,15 +5,15 @@ The service can run on a machine with GPU + X server (for AI2-THOR rendering), while VAGEN RL training runs on a separate machine using GymImageEnvClient. -Multi-GPU is the default: GPUs are auto-detected and sessions are -distributed to the least-loaded GPU automatically. +GPUs and X servers are auto-detected and started automatically. You only +need to override them if the defaults don't work for your setup. Usage: - # Auto-detect GPUs (default) - python -m vagen.envs.eb_alfred.serve --port 8000 + # Auto-detect all GPUs, start Xorg automatically: + python -m vagen.envs.eb_alfred.serve - # Override: use only specific GPUs - python -m vagen.envs.eb_alfred.serve --port 8000 --x-displays 0,1 + # Override GPU list or other settings: + python -m vagen.envs.eb_alfred.serve --devices='[0,1]' --capacity=64 --port=8001 # Then on the training machine, configure env_config: # base_urls: ["http://:8000"] @@ -21,94 +21,180 @@ # resolution: 500 """ -import argparse +from __future__ import annotations + import asyncio import concurrent.futures +import logging +import os +import subprocess +import time +from typing import List, Optional + +import fire import uvicorn -from vagen.envs_remote.service import GymService +from vagen.envs_remote import GymService from .handler import EbAlfredHandler +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s") +LOGGER = logging.getLogger(__name__) + +_XORG_CONF_TEMPLATE = """\ +Section "Device" + Identifier "GPU{idx}" + Driver "nvidia" + BusID "{bus_id}" + Option "AllowEmptyInitialConfiguration" "True" +EndSection +Section "Monitor" + Identifier "Monitor{idx}" + HorizSync 28.0-80.0 + VertRefresh 48.0-75.0 + Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 +EndSection +Section "Screen" + Identifier "Screen{idx}" + Device "GPU{idx}" + Monitor "Monitor{idx}" + DefaultDepth 24 + SubSection "Display" + Depth 24 + Modes "1920x1080" + Virtual 1920 1080 + EndSubSection +EndSection +Section "ServerLayout" + Identifier "Layout{idx}" + Screen 0 "Screen{idx}" +EndSection +""" -def main(): - parser = argparse.ArgumentParser(description="EB-ALFRED Remote Environment Server") - parser.add_argument("--port", type=int, default=8000, help="Server port") - parser.add_argument("--host", type=str, default="0.0.0.0", help="Server host") - parser.add_argument( - "--x-displays", - type=str, - default=None, - help="X displays for GPU assignment (comma-separated, e.g. '0,1'). " - "Default: auto-detect all GPUs via nvidia-smi.", - ) - parser.add_argument( - "--session-timeout", - type=float, - default=3600.0, - help="Session timeout in seconds", - ) - parser.add_argument( - "--max-sessions", - type=int, - default=0, - help="Max concurrent sessions (0=unlimited)", - ) - parser.add_argument( - "--capacity", - type=int, - default=16, - help="Max concurrently running Unity environments (0=unlimited). " - "Extra sessions are queued and created as slots free up.", - ) - parser.add_argument( - "--startup-concurrency", - type=int, - default=8, - help="Max Unity processes that may be starting up simultaneously (0=unlimited). " - "Prevents CPU spikes when many capacity slots open at once. " - "E.g. --capacity 64 --startup-concurrency 8 means 64 envs run " - "concurrently but startups are staggered 8 at a time.", - ) - parser.add_argument( - "--thread-workers", - type=int, - default=128, - help="Thread pool size for Unity instance creation (default: 128)", + +def _detect_gpus() -> List[int]: + """Auto-detect NVIDIA GPU indices via CUDA_VISIBLE_DEVICES or nvidia-smi.""" + vis = os.environ.get("CUDA_VISIBLE_DEVICES") + if vis: + return [int(d) for d in vis.split(",") if d.strip()] + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], text=True + ) + return [int(line.strip()) for line in out.strip().split("\n") if line.strip()] + except Exception: + return [0] + + +def _get_pci_bus_id(gpu_index: int) -> str: + """Return xorg-format PCI BusID for a GPU (e.g. 'PCI:65:0:0'). + + nvidia-smi reports PCI IDs in hex (e.g. '0000:41:00.0'). + Xorg BusID uses decimal (e.g. 'PCI:65:0:0'). + """ + out = subprocess.check_output( + ["nvidia-smi", f"--id={gpu_index}", "--query-gpu=pci.bus_id", "--format=csv,noheader"], + text=True, + ).strip() + # Format: "0000:BUS:DEV.FUNC" (all hex) + _, bus_hex, dev_func = out.split(":") + dev_hex, func_hex = dev_func.split(".") + return f"PCI:{int(bus_hex, 16)}:{int(dev_hex, 16)}:{int(func_hex, 16)}" + + +def _xorg_running(display: int) -> bool: + """Check if an X server is already running on the given display.""" + return os.path.exists(f"/tmp/.X{display}-lock") + + +def _start_xorg(gpu_index: int, display: int) -> None: + """Generate xorg.conf and (re)start Xorg for a GPU/display pair. + + Skips silently if Xorg is already running on that display. + """ + if _xorg_running(display): + LOGGER.info(f"Xorg already running on :{display}, skipping") + return + + pci_bus_id = _get_pci_bus_id(gpu_index) + conf_path = f"/tmp/xorg{display}.conf" + with open(conf_path, "w") as f: + f.write(_XORG_CONF_TEMPLATE.format(idx=display, bus_id=pci_bus_id)) + + LOGGER.info(f"Starting Xorg :{display} for GPU {gpu_index} (BusID={pci_bus_id})") + subprocess.Popen( + ["Xorg", "-noreset", "+extension", "GLX", "-config", conf_path, f":{display}"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, ) - args = parser.parse_args() - x_displays = args.x_displays.split(",") if args.x_displays else None + for _ in range(15): + time.sleep(1) + if _xorg_running(display): + LOGGER.info(f"Xorg :{display} ready") + return + + raise RuntimeError(f"Xorg :{display} did not become ready within 15 seconds") + + +def main( + host: str = "0.0.0.0", + port: int = 8000, + # GPU device indices. None = auto-detect via CUDA_VISIBLE_DEVICES or nvidia-smi. + # Convention: GPU i is assigned to X display :i. + devices: Optional[List[int]] = None, + # Max concurrently running Unity environments (0 = unlimited). + # Extra /connect requests are queued and served as slots free up. + capacity: int = 16, + # Max Unity processes starting up simultaneously (0 = unlimited). + # Prevents CPU spikes when many capacity slots open at once. + startup_concurrency: int = 8, + # Thread pool for asyncio.to_thread(). Should be >= capacity. + thread_pool_size: int = 128, + # Session idle timeout before auto-cleanup (seconds). + session_timeout: float = 3600.0, + # Max total sessions (0 = unlimited). + max_sessions: int = 0, + # API key for authentication. Empty = no auth. + api_key: str = "", + # Uvicorn workers. Keep at 1 (handler state is in-process). + workers: int = 1, +): + """Start the EB-ALFRED environment server.""" + if devices is None: + devices = _detect_gpus() + + # Start one Xorg server per GPU (display :i = GPU i) + for gpu_idx in devices: + _start_xorg(gpu_idx, display=gpu_idx) + x_displays = [str(gpu_idx) for gpu_idx in devices] + + executor = concurrent.futures.ThreadPoolExecutor(max_workers=thread_pool_size) + + LOGGER.info( + f"GPUs: {devices} | displays: {x_displays} | " + f"capacity: {capacity} | startup_concurrency: {startup_concurrency} | " + f"threads: {thread_pool_size}" + ) handler = EbAlfredHandler( x_displays=x_displays, - session_timeout=args.session_timeout, - max_sessions=args.max_sessions, - capacity=args.capacity, - startup_concurrency=args.startup_concurrency, + capacity=capacity, + startup_concurrency=startup_concurrency, + session_timeout=session_timeout, + max_sessions=max_sessions, ) - app = GymService(handler, max_inflight=0).build() - - # Expand the asyncio thread pool via FastAPI startup so concurrent Unity - # startups don't queue behind Python's default limit of min(32, cpu+4). - _thread_workers = args.thread_workers + app = GymService(handler, api_key=api_key).build() @app.on_event("startup") - async def _set_thread_pool(): - loop = asyncio.get_event_loop() - loop.set_default_executor( - concurrent.futures.ThreadPoolExecutor(max_workers=_thread_workers) - ) + async def _configure_executor(): + asyncio.get_running_loop().set_default_executor(executor) + + @app.on_event("shutdown") + def _shutdown_executor(): + executor.shutdown(wait=True) - displays_str = ", ".join(f":{d}" for d in handler._x_displays) - cap_str = str(args.capacity) if args.capacity > 0 else "unlimited" - startup_str = str(args.startup_concurrency) if args.startup_concurrency > 0 else "unlimited" - print(f"Starting EB-ALFRED service on {args.host}:{args.port}") - print(f"GPU displays: [{displays_str}] (auto-balanced)") - print(f"Capacity: {cap_str} concurrent environments") - print(f"Startup concurrency: {startup_str} simultaneous Unity startups") - print(f"Health check: http://localhost:{args.port}/health") - uvicorn.run(app, host=args.host, port=args.port) + uvicorn.run(app, host=host, port=port, workers=workers) if __name__ == "__main__": - main() + fire.Fire(main) diff --git a/vagen/envs/eb_alfred/start_server.sh b/vagen/envs/eb_alfred/start_server.sh index 969869730..0a8be6431 100755 --- a/vagen/envs/eb_alfred/start_server.sh +++ b/vagen/envs/eb_alfred/start_server.sh @@ -1,16 +1,8 @@ #!/bin/bash -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -cp "$SCRIPT_DIR/xorg0.conf" /tmp/xorg0.conf -cp "$SCRIPT_DIR/xorg1.conf" /tmp/xorg1.conf - -Xorg -noreset +extension GLX -config /tmp/xorg0.conf :0 & -Xorg -noreset +extension GLX -config /tmp/xorg1.conf :1 & - -sleep 2 - +# Start the EB-ALFRED server. +# GPUs and Xorg servers are auto-detected and started by serve.py. +# Override with --devices='[0,1]' if needed. python -m vagen.envs.eb_alfred.serve \ --port 8000 \ --capacity 90 \ - --startup-concurrency 6 \ - --x-displays 0,1 + --startup_concurrency 6 diff --git a/vagen/envs/eb_alfred/xorg0.conf b/vagen/envs/eb_alfred/xorg0.conf deleted file mode 100644 index 5b5d53691..000000000 --- a/vagen/envs/eb_alfred/xorg0.conf +++ /dev/null @@ -1,27 +0,0 @@ -Section "Device" - Identifier "GPU0" - Driver "nvidia" - BusID "PCI:65:0:0" - Option "AllowEmptyInitialConfiguration" "True" -EndSection -Section "Monitor" - Identifier "Monitor0" - HorizSync 28.0-80.0 - VertRefresh 48.0-75.0 - Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 -EndSection -Section "Screen" - Identifier "Screen0" - Device "GPU0" - Monitor "Monitor0" - DefaultDepth 24 - SubSection "Display" - Depth 24 - Modes "1920x1080" - Virtual 1920 1080 - EndSubSection -EndSection -Section "ServerLayout" - Identifier "Layout0" - Screen 0 "Screen0" -EndSection diff --git a/vagen/envs/eb_alfred/xorg1.conf b/vagen/envs/eb_alfred/xorg1.conf deleted file mode 100644 index 99e40b539..000000000 --- a/vagen/envs/eb_alfred/xorg1.conf +++ /dev/null @@ -1,27 +0,0 @@ -Section "Device" - Identifier "GPU1" - Driver "nvidia" - BusID "PCI:97:0:0" - Option "AllowEmptyInitialConfiguration" "True" -EndSection -Section "Monitor" - Identifier "Monitor1" - HorizSync 28.0-80.0 - VertRefresh 48.0-75.0 - Modeline "1920x1080" 172.80 1920 2040 2248 2576 1080 1081 1084 1118 -EndSection -Section "Screen" - Identifier "Screen1" - Device "GPU1" - Monitor "Monitor1" - DefaultDepth 24 - SubSection "Display" - Depth 24 - Modes "1920x1080" - Virtual 1920 1080 - EndSubSection -EndSection -Section "ServerLayout" - Identifier "Layout1" - Screen 0 "Screen1" -EndSection From 5b458c83f1c09b657bba8358bf857c6cf7155979 Mon Sep 17 00:00:00 2001 From: VAGEN Dev Date: Wed, 18 Mar 2026 00:42:01 +0000 Subject: [PATCH 09/29] Rewrite eb_alfred README to match navigation style Concise sections: Running Service / Evaluation / Training / Prompt Formats / Datasets / Interactive Test / Checklist. Remove verbose prose, add auto-detect GPU/Xorg note, SSH tunnel tip. Co-Authored-By: Claude Sonnet 4.6 --- vagen/envs/eb_alfred/README.md | 97 ++++++++++++++++++++++++++++------ 1 file changed, 80 insertions(+), 17 deletions(-) diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 4577d389b..27ce9f3c0 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -1,8 +1,17 @@ # EB-ALFRED Environment -## Installation (First-Time Only) +AI2-THOR based household robot task environment from [EmbodiedBench](https://github.com/EmbodiedBench/EmbodiedBench). The agent receives egocentric RGB images and executes multi-step tasks (cleaning, heating, slicing, storing objects). + +## Running the Service + +The environment runs on a **separate GPU machine** with a physical or virtual display. AI2-THOR requires X11 rendering (CloudRendering is not supported on ai2thor 2.1.0). + +**One-time setup — create the conda environment:** ```bash +conda create -n embodiedbench python=3.9 -y +conda activate embodiedbench + git clone https://github.com/EmbodiedBench/EmbodiedBench.git /root/EmbodiedBench touch /root/EmbodiedBench/embodiedbench/__init__.py pip install -e /root/EmbodiedBench @@ -13,41 +22,95 @@ pip install "ai2thor==2.1.0" "gym==0.23.0" "numpy<2.0" \ pip install "flask==1.1.4" "werkzeug==1.0.1" \ "markupsafe<2.1" "jinja2<3.0" "itsdangerous<2.0" pip install "opencv-python-headless<4.9" -pip install "numpy<2.0" +# Download dataset git clone https://huggingface.co/datasets/EmbodiedBench/EB-ALFRED -mv EB-ALFRED /path/to/EmbodiedBench/embodiedbench/envs/eb_alfred/data/json_2.1.0 +mv EB-ALFRED /root/EmbodiedBench/embodiedbench/envs/eb_alfred/data/json_2.1.0 ``` ---- +**Start the server** (GPUs and Xorg are auto-detected and started): -## Per-Boot Setup +```bash +conda activate embodiedbench +python -m vagen.envs.eb_alfred.serve +``` -### 1. Start Xorg + Server +Key parameters: +- `devices`: GPU indices (default: auto-detect via `CUDA_VISIBLE_DEVICES` or `nvidia-smi`) +- `capacity`: max concurrent Unity environments (default: 16) +- `startup_concurrency`: max Unity processes starting simultaneously, prevents CPU spikes (default: 8) +- `session_timeout`: idle session cleanup in seconds (default: 3600) ```bash -bash vagen/envs/eb_alfred/start_server.sh +# Example: 2 GPUs, higher capacity +python -m vagen.envs.eb_alfred.serve --devices='[0,1]' --capacity=90 --startup_concurrency=6 --port=8000 ``` -### 2. SSH Tunnel (if training machine is remote) +**SSH tunnel** (if training machine is remote): ```bash -# On env server +# Run on the env server — forwards port 8000 to training machine ssh -p -R 8000:localhost:8000 \ - -o ServerAliveInterval=30 \ - -o ServerAliveCountMax=5 \ + -o ServerAliveInterval=30 -o ServerAliveCountMax=5 \ -N -f user@training-machine-ip -# On training machine — set once in /etc/ssh/sshd_config: +# On training machine, allow many tunnels — add to /etc/ssh/sshd_config: # MaxSessions 200 # then: service ssh reload ``` ---- +## Evaluation + +```bash +conda activate vagen + +# Terminal 1 (env server): start service +python -m vagen.envs.eb_alfred.serve --devices='[0,1]' --capacity=90 + +# Terminal 2 (training machine): run eval +python -m vagen.evaluate.run_eval --config examples/evaluate/eb_alfred/config.yaml +``` + +Config: `examples/evaluate/eb_alfred/config.yaml` + +## Training + +```bash +conda activate vagen + +# Terminal 1 (env server): start service +python -m vagen.envs.eb_alfred.serve --devices='[0,1]' --capacity=90 + +# Terminal 2 (training machine): run training +cd VAGEN +bash examples/train/eb_alfred/train_grpo_qwen25vl3b.sh +``` + +Configs: `examples/train/eb_alfred/` + +## Prompt Formats + +- `free_think`: `......` +- `wm`: `............` + +## Datasets + +The `eval_set` parameter selects which episode split to use: +- `base` — standard household tasks +- `long` — longer horizon tasks + +## Interactive Test + +```bash +conda activate embodiedbench +python -m vagen.envs.eb_alfred.eb_alfred_env --x_display 0 --eval_set base --seed 0 +``` ## Checklist -- [ ] `ps aux | grep Xorg` -- [ ] `curl http://localhost:8000/health` -- [ ] `ps aux | grep thor` — no leftover Unity processes -- [ ] `nvidia-smi` — GPU memory free +```bash +ps aux | grep Xorg # Xorg running per GPU +curl http://localhost:8000/health +ps aux | grep thor # no leftover Unity processes from crashed sessions +nvidia-smi # GPU memory available +``` From 068dede143764bcdafe487dd91659357b50621b1 Mon Sep 17 00:00:00 2001 From: VAGEN Dev Date: Wed, 18 Mar 2026 00:59:54 +0000 Subject: [PATCH 10/29] Trim eb_alfred README: remove Prompt Formats, move dataset into setup, drop Interactive Test Co-Authored-By: Claude Sonnet 4.6 --- vagen/envs/eb_alfred/README.md | 23 ++++------------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 27ce9f3c0..98200fc79 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -22,8 +22,11 @@ pip install "ai2thor==2.1.0" "gym==0.23.0" "numpy<2.0" \ pip install "flask==1.1.4" "werkzeug==1.0.1" \ "markupsafe<2.1" "jinja2<3.0" "itsdangerous<2.0" pip install "opencv-python-headless<4.9" +``` + +**Download dataset** (`eval_set` selects the split: `base` — standard tasks, `long` — longer horizon): -# Download dataset +```bash git clone https://huggingface.co/datasets/EmbodiedBench/EB-ALFRED mv EB-ALFRED /root/EmbodiedBench/embodiedbench/envs/eb_alfred/data/json_2.1.0 ``` @@ -88,24 +91,6 @@ bash examples/train/eb_alfred/train_grpo_qwen25vl3b.sh Configs: `examples/train/eb_alfred/` -## Prompt Formats - -- `free_think`: `......` -- `wm`: `............` - -## Datasets - -The `eval_set` parameter selects which episode split to use: -- `base` — standard household tasks -- `long` — longer horizon tasks - -## Interactive Test - -```bash -conda activate embodiedbench -python -m vagen.envs.eb_alfred.eb_alfred_env --x_display 0 --eval_set base --seed 0 -``` - ## Checklist ```bash From 24d5de4b0950aaeb0c1e6e4102bee283a6355f99 Mon Sep 17 00:00:00 2001 From: VAGEN Dev Date: Wed, 18 Mar 2026 01:00:24 +0000 Subject: [PATCH 11/29] Remove Checklist section from eb_alfred README Co-Authored-By: Claude Sonnet 4.6 --- vagen/envs/eb_alfred/README.md | 8 -------- 1 file changed, 8 deletions(-) diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 98200fc79..99b6f7a97 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -91,11 +91,3 @@ bash examples/train/eb_alfred/train_grpo_qwen25vl3b.sh Configs: `examples/train/eb_alfred/` -## Checklist - -```bash -ps aux | grep Xorg # Xorg running per GPU -curl http://localhost:8000/health -ps aux | grep thor # no leftover Unity processes from crashed sessions -nvidia-smi # GPU memory available -``` From ef9dd73a4d50367b7f9fa6916db91082c4aede63 Mon Sep 17 00:00:00 2001 From: VAGEN Dev Date: Wed, 18 Mar 2026 02:05:21 +0000 Subject: [PATCH 12/29] update service and scripts --- examples/evaluate/eb_alfred/config.yaml | 77 +++++++++-------- examples/evaluate/eb_alfred/run_eval.sh | 14 +++ .../eb_alfred/train_eb_alfred_vision.yaml | 66 ++++++++++++++ .../train_ppo_no_concat_qwen25vl3b.sh | 86 +++++++++++++++++++ .../train/eb_alfred/val_eb_alfred_vision.yaml | 44 ++++++++++ 5 files changed, 253 insertions(+), 34 deletions(-) create mode 100755 examples/evaluate/eb_alfred/run_eval.sh create mode 100644 examples/train/eb_alfred/train_eb_alfred_vision.yaml create mode 100755 examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh create mode 100644 examples/train/eb_alfred/val_eb_alfred_vision.yaml diff --git a/examples/evaluate/eb_alfred/config.yaml b/examples/evaluate/eb_alfred/config.yaml index c1fa62d0a..1d43c6eb4 100644 --- a/examples/evaluate/eb_alfred/config.yaml +++ b/examples/evaluate/eb_alfred/config.yaml @@ -1,37 +1,46 @@ -# EB-ALFRED Evaluation Config (ERA-aligned) -# -# Uses the remote environment pattern: -# 1. Start the EB-ALFRED server (requires GPU + X display): -# DISPLAY=:0 python -m vagen.envs.eb_alfred.serve --port 8000 -# -# 2. Run evaluation: -# python -m vagen.evaluate.run_eval --config examples/evaluate/eb_alfred/config.yaml -# -# Key ERA-aligned settings: -# - No-concat mode: system prompt (with task + actions) + current obs only per turn -# - Multi-step planning: up to 20 actions per LLM call -# - Max 30 env steps per episode (matches ERA's _max_episode_steps) -# - ERA-style replan: break on action failure, model replans next turn -# - Task examples in system prompt (cleaning, slicing, heating patterns) +fileroot: ${oc.env:HOME}/projects/vagen envs: - name: RemoteEnv n_envs: 50 - tag_id: eb_alfred_eval + data_source: eb_alfred + tag_id: eb_alfred_val_common seed: [0, 50, 1] - split: test - concat_history: false # no-concat: system + current obs only per turn - max_turns: 30 + max_turns: 6 config: base_urls: - "http://localhost:8000" timeout: 600 - eval_set: base # options: base, common_sense, complex, long_horizon - x_display: "0" + eval_set: common_sense obs_image_size: 500 - max_turns: 30 - max_actions_per_step: 20 # multi-step planning (ERA-aligned) - max_env_steps: 30 # total env actions cap (ERA-aligned) + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_val_spatial + seed: [0, 50, 1] + max_turns: 6 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: spatial + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 action_sep: "," prompt_format: free_think use_example_in_sys_prompt: true @@ -43,22 +52,22 @@ envs: top_p: 1.0 experiment: - dump_dir: ./rollouts/eval_eb_alfred - default_max_turns: 30 + dump_dir: ${fileroot}/rollouts/eval_eb_alfred + default_max_turns: 6 run: - backend: openai + backend: "openai" base_seed: 0 - max_concurrent_jobs: 10 + max_concurrent_jobs: 4 resume: skip_completed live_summary: true backends: openai: - api_key: "" # uses OPENAI_API_KEY env var + api_key: "" # or env OPENAI_API_KEY base_url: null - model: "gpt-4.1" - max_concurrency: 100 + model: "gpt-4o-mini" + max_concurrency: 2 max_retries: 6 min_backoff: 0.5 max_backoff: 8.0 @@ -67,7 +76,7 @@ backends: base_url: "http://127.0.0.1:30000/v1" api_key: "EMPTY" model: "Qwen/Qwen2.5-VL-7B-Instruct" - max_concurrency: 50 + max_concurrency: 2 max_retries: 6 min_backoff: 0.5 max_backoff: 8.0 @@ -75,8 +84,8 @@ backends: claude: api_key: "" base_url: null - model: "claude-sonnet-4-6" - max_concurrency: 10 + model: "claude-3-5-sonnet-latest" + max_concurrency: 2 max_retries: 6 min_backoff: 0.5 max_backoff: 8.0 diff --git a/examples/evaluate/eb_alfred/run_eval.sh b/examples/evaluate/eb_alfred/run_eval.sh new file mode 100755 index 000000000..bdafead66 --- /dev/null +++ b/examples/evaluate/eb_alfred/run_eval.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Before running, start the eb_alfred server in another terminal: +# python -m vagen.envs.eb_alfred.serve + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CONFIG="${1:-$SCRIPT_DIR/config.yaml}" +shift 2>/dev/null || true + +LOG_FILE="run.log" + +python -m vagen.evaluate.run_eval --config "$CONFIG" "$@" \ + 2>&1 | tee "${LOG_FILE}" diff --git a/examples/train/eb_alfred/train_eb_alfred_vision.yaml b/examples/train/eb_alfred/train_eb_alfred_vision.yaml new file mode 100644 index 000000000..d471ab1f6 --- /dev/null +++ b/examples/train/eb_alfred/train_eb_alfred_vision.yaml @@ -0,0 +1,66 @@ +envs: + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_train_base + seed: [0, 200, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: base + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_train_complex + seed: [0, 200, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: complex_instruction + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_train_visual + seed: [0, 200, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: visual_appearance + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 diff --git a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh new file mode 100755 index 000000000..9bf8fa42c --- /dev/null +++ b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh @@ -0,0 +1,86 @@ +#!/bin/bash + +set -x + +PROJECT_NAME="vagen_experiments" +EXPERIMENT_NAME="ppo_eb_alfred_no_concat" + +BASEDIR=$(pwd) +SCRIPTDIR=$(dirname "$0") +EXPERIMENT_DIR=${BASEDIR}/exps/${PROJECT_NAME}/${EXPERIMENT_NAME} +SAVE_CHECKPOINT_DIR=${EXPERIMENT_DIR}/verl_checkpoints +DATASET_TRAIN=${SCRIPTDIR}/train_eb_alfred_vision.yaml +DATASET_VAL=${SCRIPTDIR}/val_eb_alfred_vision.yaml +agent_loop_config_path=${BASEDIR}/vagen/configs/agent_no_concat.yaml +REF_MODEL_PATH=Qwen/Qwen2.5-VL-3B-Instruct +mkdir -p ${EXPERIMENT_DIR} + +export HF_HOME=/workspace/.hf_home +export PATH=/venv/vagen/bin:$PATH + +PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ + --config-path=${BASEDIR}/vagen/configs \ + --config-name='vagen_multiturn' \ + data.train_files=${DATASET_TRAIN} \ + data.val_files=${DATASET_VAL} \ + data.train_batch_size=128 \ + data.max_prompt_length=2048 \ + data.max_response_length=512 \ + +data.max_trajectory_length=7000 \ + algorithm.adv_estimator=no_concat_gae_first \ + algorithm.kl_ctrl.kl_coef=0.0 \ + actor_rollout_ref.model.path=${REF_MODEL_PATH} \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.use_fused_kernels=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.kl_loss_coef=0.0 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0.0 \ + actor_rollout_ref.actor.checkpoint.save_contents=['model','hf_model','optimizer','extra'] \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.freeze_vision_tower=True \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=sglang \ + actor_rollout_ref.rollout.mode=async \ + actor_rollout_ref.rollout.n=1 \ + actor_rollout_ref.rollout.max_num_batched_tokens=10000 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.multi_turn.enable=True \ + actor_rollout_ref.rollout.agent.agent_loop_config_path=$agent_loop_config_path \ + actor_rollout_ref.rollout.disable_log_stats=False \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + critic.enable=True \ + critic.optim.lr=1e-5 \ + critic.model.use_remove_padding=True \ + critic.model.path=${REF_MODEL_PATH} \ + critic.model.enable_gradient_checkpointing=True \ + critic.ppo_micro_batch_size_per_gpu=1 \ + critic.model.fsdp_config.param_offload=True \ + critic.model.fsdp_config.optimizer_offload=True \ + trainer.critic_warmup=0 \ + trainer.logger=['console','wandb'] \ + trainer.val_before_train=True \ + trainer.n_gpus_per_node=4 \ + trainer.nnodes=1 \ + trainer.save_freq=100 \ + trainer.test_freq=20 \ + trainer.project_name=${PROJECT_NAME} \ + trainer.experiment_name=${EXPERIMENT_NAME} \ + trainer.default_local_dir=${SAVE_CHECKPOINT_DIR} \ + trainer.validation_data_dir=${EXPERIMENT_DIR}/validation \ + trainer.rollout_data_dir=${EXPERIMENT_DIR}/rollout_data \ + trainer.log_val_generations=32 \ + +trainer.concat_multi_turn=False \ + trainer.total_training_steps=400 2>&1 | \ + tee ${EXPERIMENT_DIR}/${PROJECT_NAME}_${EXPERIMENT_NAME}.log >(tee ${BASEDIR}/${PROJECT_NAME}_${EXPERIMENT_NAME}.log >/dev/null) diff --git a/examples/train/eb_alfred/val_eb_alfred_vision.yaml b/examples/train/eb_alfred/val_eb_alfred_vision.yaml new file mode 100644 index 000000000..233fc3e21 --- /dev/null +++ b/examples/train/eb_alfred/val_eb_alfred_vision.yaml @@ -0,0 +1,44 @@ +envs: + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_val_common + seed: [0, 50, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: common_sense + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_val_spatial + seed: [0, 50, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: spatial + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 From d5e0c648ccbafa03b11c671887cbdc422b358498 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Thu, 19 Mar 2026 17:32:53 -0400 Subject: [PATCH 13/29] update prompt logic --- vagen/envs/eb_alfred/eb_alfred_env.py | 1 + vagen/envs/eb_alfred/serve.py | 3 +- vagen/envs/eb_alfred/utils/prompt.py | 59 ++++++++++++++++----------- vagen/envs/eb_alfred/utils/utils.py | 24 ++++++++++- 4 files changed, 61 insertions(+), 26 deletions(-) diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py index 083f5ae11..5434fdffc 100644 --- a/vagen/envs/eb_alfred/eb_alfred_env.py +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -352,6 +352,7 @@ def _build_obs(self, init: bool) -> Dict[str, Any]: last_action=self._last_action, env_feedback=self._last_feedback, img_str=img_str, + task_instruction=self.env.episode_language_instruction, ) return { diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index 1327e15a9..e8b6fe741 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -103,7 +103,8 @@ def _get_pci_bus_id(gpu_index: int) -> str: def _xorg_running(display: int) -> bool: """Check if an X server is already running on the given display.""" - return os.path.exists(f"/tmp/.X{display}-lock") + return (os.path.exists(f"/tmp/.X{display}-lock") + or os.path.exists(f"/tmp/.X11-unix/X{display}")) def _start_xorg(gpu_index: int, display: int) -> None: diff --git a/vagen/envs/eb_alfred/utils/prompt.py b/vagen/envs/eb_alfred/utils/prompt.py index 8521c6716..2b0f1a75b 100644 --- a/vagen/envs/eb_alfred/utils/prompt.py +++ b/vagen/envs/eb_alfred/utils/prompt.py @@ -64,14 +64,26 @@ def system_prompt(task_instruction: Optional[str] = None, action_list: Optional[ if add_task_examples and TASK_EXAMPLES: base += "\n\n## Task Examples" for i, ex in enumerate(TASK_EXAMPLES): - actions_str = ", ".join(ex["actions"]) + if action_list is not None: + # Build action-to-id lookup from the current episode's action list + name_to_id = {a.lower(): idx for idx, a in enumerate(action_list)} + parts = [] + for a in ex["actions"]: + aid = name_to_id.get(a.lower()) + if aid is not None: + parts.append(f"[{aid}, {a}]") + else: + parts.append(a) + actions_str = "| ".join(parts) + else: + actions_str = "| ".join(ex["actions"]) base += f"\n\nExample {i+1}: {ex['task']}\n{ex['think']}\n{actions_str}" if task_instruction is not None: base += f"\n\n## Current Task\n{task_instruction}" if action_list is not None: - actions_str = "\n".join(f"action id {i}: {a}" for i, a in enumerate(action_list)) + actions_str = "\n".join(f"[{i}, {a}]" for i, a in enumerate(action_list)) base += f"\n\n## Available Actions (0~{len(action_list) - 1})\n{actions_str}" return base @@ -89,20 +101,21 @@ def init_observation_template(img_str): Decide your next action.""" -def action_template(last_action, env_feedback, img_str): +def action_template(last_action, env_feedback, img_str, task_instruction=None): """Template for step observation with feedback. - Encourages structured reasoning: describe what you see, + Encourages structured reasoning: analyze the feedback, reflect on why the last action succeeded or failed, - then plan your next actions. + then plan the next logical step. """ + task_line = f"\n[Task]: {task_instruction}\n" if task_instruction else "" return f"""[Last Action]: {last_action} [Feedback]: {env_feedback} - +{task_line} [Current Observation]: {img_str} -Describe what you see, reflect on the feedback, and plan your next actions.""" +You MUST first analyze the feedback above. If the action succeeded, plan the next logical step to complete the task. If it failed, explain why and try a different approach. Do NOT repeat the same failed action.""" def format_prompt(max_actions_per_step, action_sep, add_example=True, prompt_format="free_think"): @@ -119,42 +132,42 @@ def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True) """Generate format prompt for free_think format.""" if max_actions_per_step == 1: base = """You should output 1 action at a time. -Output the action name exactly as listed in the available actions, or the action ID (integer). +Output the action as [action_id, action_name] using the ID from the available actions list. Your response should be in the format of: -...action name or action ID""" +...[N, action name]""" else: base = f"""You should output a plan of up to {max_actions_per_step} actions at a time, separated by "{action_sep}". -Output the action name exactly as listed in the available actions, or the action ID (integer). +Output each action as [action_id, action_name] using the ID from the available actions list. Your response should be in the format of: -...action1{action_sep} action2{action_sep} ...""" +...[N1, action1]{action_sep} [N2, action2]{action_sep} ...""" if add_example: if max_actions_per_step == 1: examples = """ Example 1: I need to find a mug first. Let me navigate to where mugs might be. -find a Mug +[5, find a Mug] Example 2: The mug is nearby and I'm not holding anything. I should pick it up. -pick up the Mug +[12, pick up the Mug] Example 3: I'm holding the mug and I'm near the table. Let me put it down. -put down the object in hand""" +[38, put down the object in hand]""" else: examples = f""" Example 1 (multi-step plan): I need to find the alarm clock, pick it up, then find the desk lamp and turn it on. -find a AlarmClock{action_sep} pick up the AlarmClock{action_sep} find a DeskLamp{action_sep} turn on the DeskLamp +[3, find a AlarmClock]{action_sep} [15, pick up the AlarmClock]{action_sep} [7, find a DeskLamp]{action_sep} [42, turn on the DeskLamp] Example 2 (single action when unsure): I am not sure where the mug is. Let me find it first. -find a Mug +[5, find a Mug] Example 3 (replanning after failure): The last action failed because the cabinet was closed. I need to open it first, then pick up the object. -open the Cabinet{action_sep} pick up the Mug""" +[20, open the Cabinet]{action_sep} [12, pick up the Mug]""" return base + "\n" + examples return base @@ -163,9 +176,9 @@ def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True) def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): """Generate format prompt for wm format with observation and prediction tags.""" base = f"""You should output {max_actions_per_step} action(s) at a time. -Output the action name exactly as listed in the available actions, or the action ID (integer). +Output the action as [action_id, action_name] using the ID from the available actions list. Your response must be in the format of: -......action name or action ID.... +......[N, action name].... Rules for : - Describe the current scene: what objects you see, your position, what you are holding, and relevant receptacle states. @@ -174,26 +187,26 @@ def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): - Predict what will change after your action: where you will be, what you will see, and the expected result. Rules for : -- Output exactly 1 action name or action ID.""" +- Output exactly 1 action as [action_id, action_name].""" if add_example: examples = """ Example 1: I see a kitchen with a counter, a microwave, and a mug on the counter. I am not holding anything. I need to pick up the mug. First, I should find it to get close to it. -find a Mug +[5, find a Mug] I will navigate to the mug and see it up close on the counter. Example 2: I am close to a Mug on the counter. I am not holding anything. The mug is within reach. The mug is nearby and I'm not holding anything. I should pick it up. -pick up the Mug +[12, pick up the Mug] I will be holding the mug. The counter will no longer have the mug on it. Example 3: I am holding a Mug. I see a table nearby with an empty spot. I'm holding the mug and I'm near the table. Let me put it down. -put down the object in hand +[38, put down the object in hand] The mug will be placed on the table. I will no longer be holding anything.""" return base + "\n" + examples diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py index de104a063..c6bf9b363 100644 --- a/vagen/envs/eb_alfred/utils/utils.py +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -127,14 +127,34 @@ def match_action( """ Match a parsed action against the valid action set. - Supports two formats: + Supports multiple formats (in priority order): + - ERA-style [id, action_name]: "[42, find a Cabinet]" + - Legacy (id: N) suffix: "find a Cabinet (id: 42)" + - Plain action ID: "42" - Action name (case-insensitive): "find a Cabinet" - - Action ID (integer): "42" Returns the original action string if matched, None otherwise. """ name = action_name.strip() + # Try ERA-style [id, action_name] format + bracket_match = re.match(r'^\[(\d+),\s*(.+?)\]$', name) + if bracket_match: + idx = int(bracket_match.group(1)) + if 0 <= idx < len(action_list): + return action_list[idx] + # ID out of range; try name part + fallback_name = bracket_match.group(2).strip() + return action_map.get(fallback_name.lower()) + + # Try legacy "(id: N)" suffix + id_match = re.search(r'\(id:\s*(\d+)\)\s*$', name) + if id_match: + idx = int(id_match.group(1)) + if 0 <= idx < len(action_list): + return action_list[idx] + name = name[:id_match.start()].strip() + # Try as integer action ID try: idx = int(name) From f3723e47c270e7c19083759f8a8e4edae959befd Mon Sep 17 00:00:00 2001 From: tmp Date: Thu, 19 Mar 2026 21:49:07 +0000 Subject: [PATCH 14/29] tmp --- vagen/envs/eb_alfred/utils/utils.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py index c6bf9b363..7a442d191 100644 --- a/vagen/envs/eb_alfred/utils/utils.py +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -104,6 +104,21 @@ def parse_wm(response: str, action_sep: str = ",", max_actions: int = 1) -> Dict } +def normalize_era_tokens(response: str) -> str: + """ + Convert ERA special tokens to VAGEN plain tags so the parser can handle + models trained with ERA's SFT format. + + ERA format: <|think_start|>...<|think_end|><|action_start|>...<|action_end|> + VAGEN format: ...... + """ + response = response.replace("<|think_start|>", "") + response = response.replace("<|think_end|>", "") + response = response.replace("<|action_start|>", "") + response = response.replace("<|action_end|>", "") + return response + + def parse_response( response: str, prompt_format: str = "free_think", @@ -111,6 +126,7 @@ def parse_response( max_actions: int = 1, ) -> Dict: """Parse LLM response based on the specified prompt format.""" + response = normalize_era_tokens(response) if prompt_format == "free_think": return parse_free_think(response, action_sep, max_actions) elif prompt_format == "wm": From ea2ae53a18798bc909a74051dbf5588f33f02d61 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Thu, 19 Mar 2026 17:41:33 -0400 Subject: [PATCH 15/29] Align action format to ERA [id, 'action_name'] style Change all prompt examples, available actions list, and format instructions from `action_name (id: N)` to `[N, 'action_name']` to match ERA SFT data format exactly. Update match_action parser to handle quoted names. Co-Authored-By: Claude Opus 4.6 (1M context) --- vagen/envs/eb_alfred/utils/prompt.py | 30 ++++++++++++++-------------- vagen/envs/eb_alfred/utils/utils.py | 4 ++-- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/vagen/envs/eb_alfred/utils/prompt.py b/vagen/envs/eb_alfred/utils/prompt.py index 2b0f1a75b..b162cf56c 100644 --- a/vagen/envs/eb_alfred/utils/prompt.py +++ b/vagen/envs/eb_alfred/utils/prompt.py @@ -71,7 +71,7 @@ def system_prompt(task_instruction: Optional[str] = None, action_list: Optional[ for a in ex["actions"]: aid = name_to_id.get(a.lower()) if aid is not None: - parts.append(f"[{aid}, {a}]") + parts.append(f"[{aid}, '{a}']") else: parts.append(a) actions_str = "| ".join(parts) @@ -83,7 +83,7 @@ def system_prompt(task_instruction: Optional[str] = None, action_list: Optional[ base += f"\n\n## Current Task\n{task_instruction}" if action_list is not None: - actions_str = "\n".join(f"[{i}, {a}]" for i, a in enumerate(action_list)) + actions_str = "\n".join(f"[{i}, '{a}']" for i, a in enumerate(action_list)) base += f"\n\n## Available Actions (0~{len(action_list) - 1})\n{actions_str}" return base @@ -134,40 +134,40 @@ def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True) base = """You should output 1 action at a time. Output the action as [action_id, action_name] using the ID from the available actions list. Your response should be in the format of: -...[N, action name]""" +...[N, 'action name']""" else: base = f"""You should output a plan of up to {max_actions_per_step} actions at a time, separated by "{action_sep}". Output each action as [action_id, action_name] using the ID from the available actions list. Your response should be in the format of: -...[N1, action1]{action_sep} [N2, action2]{action_sep} ...""" +...[N1, 'action1']{action_sep} [N2, 'action2']{action_sep} ...""" if add_example: if max_actions_per_step == 1: examples = """ Example 1: I need to find a mug first. Let me navigate to where mugs might be. -[5, find a Mug] +[5, 'find a Mug'] Example 2: The mug is nearby and I'm not holding anything. I should pick it up. -[12, pick up the Mug] +[12, 'pick up the Mug'] Example 3: I'm holding the mug and I'm near the table. Let me put it down. -[38, put down the object in hand]""" +[38, 'put down the object in hand']""" else: examples = f""" Example 1 (multi-step plan): I need to find the alarm clock, pick it up, then find the desk lamp and turn it on. -[3, find a AlarmClock]{action_sep} [15, pick up the AlarmClock]{action_sep} [7, find a DeskLamp]{action_sep} [42, turn on the DeskLamp] +[3, 'find a AlarmClock']{action_sep} [15, 'pick up the AlarmClock']{action_sep} [7, 'find a DeskLamp']{action_sep} [42, 'turn on the DeskLamp'] Example 2 (single action when unsure): I am not sure where the mug is. Let me find it first. -[5, find a Mug] +[5, 'find a Mug'] Example 3 (replanning after failure): The last action failed because the cabinet was closed. I need to open it first, then pick up the object. -[20, open the Cabinet]{action_sep} [12, pick up the Mug]""" +[20, 'open the Cabinet']{action_sep} [12, 'pick up the Mug']""" return base + "\n" + examples return base @@ -178,7 +178,7 @@ def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): base = f"""You should output {max_actions_per_step} action(s) at a time. Output the action as [action_id, action_name] using the ID from the available actions list. Your response must be in the format of: -......[N, action name].... +......[N, 'action name'].... Rules for : - Describe the current scene: what objects you see, your position, what you are holding, and relevant receptacle states. @@ -187,26 +187,26 @@ def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): - Predict what will change after your action: where you will be, what you will see, and the expected result. Rules for : -- Output exactly 1 action as [action_id, action_name].""" +- Output exactly 1 action as [action_id, 'action_name'].""" if add_example: examples = """ Example 1: I see a kitchen with a counter, a microwave, and a mug on the counter. I am not holding anything. I need to pick up the mug. First, I should find it to get close to it. -[5, find a Mug] +[5, 'find a Mug'] I will navigate to the mug and see it up close on the counter. Example 2: I am close to a Mug on the counter. I am not holding anything. The mug is within reach. The mug is nearby and I'm not holding anything. I should pick it up. -[12, pick up the Mug] +[12, 'pick up the Mug'] I will be holding the mug. The counter will no longer have the mug on it. Example 3: I am holding a Mug. I see a table nearby with an empty spot. I'm holding the mug and I'm near the table. Let me put it down. -[38, put down the object in hand] +[38, 'put down the object in hand'] The mug will be placed on the table. I will no longer be holding anything.""" return base + "\n" + examples diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py index 7a442d191..e01b5e339 100644 --- a/vagen/envs/eb_alfred/utils/utils.py +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -153,8 +153,8 @@ def match_action( """ name = action_name.strip() - # Try ERA-style [id, action_name] format - bracket_match = re.match(r'^\[(\d+),\s*(.+?)\]$', name) + # Try ERA-style [id, 'action_name'] format (with or without quotes) + bracket_match = re.match(r"^\[(\d+),\s*['\"]?(.+?)['\"]?\s*\]$", name) if bracket_match: idx = int(bracket_match.group(1)) if 0 <= idx < len(action_list): From c5c08f1118a4c2c3adc7f32a6cdf49473976d084 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Fri, 20 Mar 2026 10:31:21 -0400 Subject: [PATCH 16/29] update prompt logic --- vagen/envs/eb_alfred/eb_alfred_env.py | 10 ++ vagen/envs/eb_alfred/utils/prompt.py | 244 +++++++++++--------------- vagen/envs/eb_alfred/utils/utils.py | 22 ++- 3 files changed, 126 insertions(+), 150 deletions(-) diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py index 5434fdffc..46c8e4922 100644 --- a/vagen/envs/eb_alfred/eb_alfred_env.py +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -135,6 +135,8 @@ def __init__(self, env_config: Dict[str, Any]): self._total_env_steps: int = 0 self._last_action: str = "" self._last_feedback: str = "" + self._last_thinking: str = "" + self._last_action_id: Optional[int] = None self._action_list: List[str] = [] self._action_map: Dict[str, str] = {} # lowercase -> original @@ -204,6 +206,8 @@ async def reset(self, seed: int) -> Tuple[Dict[str, Any], Dict[str, Any]]: self._total_env_steps = 0 self._last_action = "" self._last_feedback = "" + self._last_thinking = "" + self._last_action_id = None # Build action lookup for this episode (action space is dynamic) self._action_list = list(self.env.language_skill_set) @@ -245,6 +249,7 @@ async def step( actions = parsed.get("actions", []) format_correct = parsed.get("format_correct", False) + self._last_thinking = parsed.get("think_content", "") metrics = { "turn_metrics": { @@ -285,6 +290,7 @@ async def step( ) self._last_action = matched + self._last_action_id = self._action_list.index(matched) if matched in self._action_list else None self._last_feedback = step_info.get("env_feedback", "") action_success = step_info.get("last_action_success", 0.0) @@ -346,6 +352,7 @@ def _build_obs(self, init: bool) -> Dict[str, Any]: if init: obs_str = init_observation_template( img_str=img_str, + task_instruction=self.env.episode_language_instruction, ) else: obs_str = action_template( @@ -353,6 +360,9 @@ def _build_obs(self, init: bool) -> Dict[str, Any]: env_feedback=self._last_feedback, img_str=img_str, task_instruction=self.env.episode_language_instruction, + step_id=self._total_turns - 1, + thinking=self._last_thinking, + action_id=self._last_action_id, ) return { diff --git a/vagen/envs/eb_alfred/utils/prompt.py b/vagen/envs/eb_alfred/utils/prompt.py index b162cf56c..6fed279b7 100644 --- a/vagen/envs/eb_alfred/utils/prompt.py +++ b/vagen/envs/eb_alfred/utils/prompt.py @@ -1,125 +1,115 @@ from typing import List, Optional -# ERA-aligned task examples teaching ALFRED-specific mechanics -# (cleaning, heating, slicing, storage, object placement) -TASK_EXAMPLES = [ - { - "task": "Pick up the alarm clock and turn on the lamp", - "think": "I need to find the alarm clock, pick it up, then find the desk lamp and turn it on.", - "actions": ["find a AlarmClock", "pick up the AlarmClock", "find a DeskLamp", "turn on the DeskLamp"], - }, - { - "task": "Set the box on the table", - "think": "I need to find the box, pick it up, then find the dining table and put it down.", - "actions": ["find a Box", "pick up the Box", "find a DiningTable", "put down the object in hand"], - }, - { - "task": "Move the towel on the hanger into the cabinet", - "think": "I need to find the hand towel, pick it up, find a cabinet, open it, put the towel inside, and close the cabinet.", - "actions": ["find a HandTowel", "pick up the HandTowel", "find a Cabinet", "open the Cabinet", "put down the object in hand", "close the Cabinet"], - }, - { - "task": "Put a clean pan in the refrigerator", - "think": "To clean the pan, I need to: pick it up, put it in the sink, turn on the faucet, turn off the faucet, then pick up the now-clean pan and put it in the fridge.", - "actions": ["find a Pan", "pick up the Pan", "find a Sink", "put down the object in hand", "find a Faucet", "turn on the Faucet", "turn off the Faucet", "find a Pan", "pick up the Pan", "find a Fridge", "open the Fridge", "put down the object in hand", "close the Fridge"], - }, - { - "task": "Slice a loaf of bread put a slice on the counter", - "think": "To slice bread I need a knife first. Pick up knife, find bread, slice it. Put knife away, then pick up a bread slice and place it on the counter.", - "actions": ["find a Knife", "pick up the Knife", "find a Bread", "slice the Bread", "find a CounterTop", "put down the object in hand", "find a Bread", "pick up the Bread", "find a CounterTop", "put down the object in hand"], - }, -] - - -def system_prompt(task_instruction: Optional[str] = None, action_list: Optional[List[str]] = None, add_task_examples: bool = True): - """ - System prompt for EB-ALFRED household robot tasks. +# ────────────────────────────────────────────────────────────────────── +# ERA-aligned system prompt +# ────────────────────────────────────────────────────────────────────── +# This matches the prompt the ERA EPL-Only model was SFT'd on, so the +# model stays in-distribution. The output (ERA special tokens) is +# normalised back to VAGEN tags by normalize_era_tokens() in utils.py. +# ────────────────────────────────────────────────────────────────────── - When task_instruction and action_list are provided (after reset), - includes the per-episode task and available actions so that - no-concat mode always has access to them. - """ - base = """You are a robot operating in a home. Given a task, you must accomplish the task using a defined set of actions to achieve the desired outcome. +ERA_SYSTEM_PROMPT_TEMPLATE = """\ +## You are a robot operating in a home. Given a task, you must accomplish the task using a defined set of actions to achieve the desired outcome. ## Action Descriptions and Validity Rules -- Find: Parameterized by the name of the receptacle to navigate to. Always valid if the object exists in the scene. -- Pick up: Parameterized by the name of the object to pick. Only valid if close to the object, not already holding something, and the object is not in a closed receptacle. -- Put down: Parameterized by the name of the object to put down to a nearby receptacle. Only valid if holding an object. -- Drop: Parameterized by the name of the object to put down. Different from 'put down' as this does not guarantee the held object will be put into a specified receptacle. -- Open: Parameterized by the name of the receptacle to open. Only valid if the receptacle is closed and close to the receptacle. -- Close: Parameterized by the name of the receptacle to close. Only valid if the receptacle is open and close to the receptacle. -- Turn on: Parameterized by the name of the object to turn on. Only valid if the object is turned off and close to the object. -- Turn off: Parameterized by the name of the object to turn off. Only valid if the object is turned on and close to the object. -- Slice: Parameterized by the name of the object to slice. Only valid if the object is sliceable and close to the object. +- Find: Parameterized by the name of the receptacle to navigate to. So long as the object is present in the scene, this skill is always valid +- Pick up: Parameterized by the name of the object to pick. Only valid if the robot is close to the object, not holding another object, and the object is not inside a closed receptacle. +- Put down: Parameterized by the name of the object to put down to a nearby receptacle. Only valid if the robot is holding an object. +- Drop: Parameterized by the name of the object to put down. It is different from Put down action, as this does not guarantee the held object will be put into a specified receptacle. +- Open: Parameterized by the name of the receptacle to open. Only valid if the receptacle is closed and the robot is close to the receptacle. +- Close: Parameterized by the name of the receptacle to close. Only valid if the receptacle is open and the robot is close to the receptacle. +- Turn on: Parameterized by the name of the object to turn on. Only valid if the object is turned off and the robot is close to the object. +- Turn off: Parameterized by the name of the object to turn off. Only valid if the object is turned on and the robot is close to the object. +- Slice: Parameterized by the name of the object to slice. Only valid if the object is sliceable and the robot is close to the object. + +## The available action id (0 ~ {max_action_id}) and action names are: {available_actions}. ## Guidelines -1. Output a plan of actions. Each plan should include no more than 20 actions. -2. Always locate an object using 'find' before interacting with it. -3. Make sure to match the action name and its corresponding action id in the output. Use 'put down' rather than 'drop' to place objects in specific receptacles. -4. Do not repeat the same failed action sequence. Try to modify the action sequence because previous actions did not lead to success. -5. Objects may have multiple instances (e.g., Cabinet_2, Cabinet_3). Explore different instances if needed. -6. Use environment feedback to refine your plan. If an action fails, reflect on the reason and adjust accordingly.""" - - if add_task_examples and TASK_EXAMPLES: - base += "\n\n## Task Examples" - for i, ex in enumerate(TASK_EXAMPLES): - if action_list is not None: - # Build action-to-id lookup from the current episode's action list - name_to_id = {a.lower(): idx for idx, a in enumerate(action_list)} - parts = [] - for a in ex["actions"]: - aid = name_to_id.get(a.lower()) - if aid is not None: - parts.append(f"[{aid}, '{a}']") - else: - parts.append(a) - actions_str = "| ".join(parts) - else: - actions_str = "| ".join(ex["actions"]) - base += f"\n\nExample {i+1}: {ex['task']}\n{ex['think']}\n{actions_str}" - - if task_instruction is not None: - base += f"\n\n## Current Task\n{task_instruction}" - - if action_list is not None: - actions_str = "\n".join(f"[{i}, '{a}']" for i, a in enumerate(action_list)) - base += f"\n\n## Available Actions (0~{len(action_list) - 1})\n{actions_str}" - - return base +1. **Output Plan**: Avoid generating empty plan. Each plan should include no more than 20 actions. +2. **Visibility**: Always locate a visible object by the 'find' action before interacting with it. +3. **Action Guidelines**: Make sure match the action name and its corresponding action id in the output. Avoid performing actions that do not meet the defined validity criteria. For instance, if you want to put object in a receptacle, use 'put down' rather than 'drop' actions. +4. **Prevent Repeating Action Sequences**: Do not repeatedly execute the same action or sequence of actions. Try to modify the action sequence because previous actions do not lead to success. +5. **Multiple Instances**: There may be multiple instances of the same object, distinguished by an index following their names, e.g., Cabinet_2, Cabinet_3. You can explore these instances if you do not find the desired object in the current receptacle. +6. **Reflection on History and Feedback**: Use interaction history and feedback from the environment to refine and improve your current plan. If the last action is invalid, reflect on the reason, such as not adhering to action rules or missing preliminary actions, and adjust your plan accordingly. + + ** Generation Guide ** + - Include the thinking process between <|think_start|> and <|think_end|> + - Include only the target action in <|action_start|> and <|action_end|>, i.e. the content inside <|action_start|> and <|action_end|> should be nothing more than [action_id, 'action_name'], where the action id is an integer and the action name is the corresponding name. Do not include any other thing, such as '"'. + """ -def init_observation_template(img_str): - """Template for initial observation after reset. +def system_prompt( + task_instruction: Optional[str] = None, + action_list: Optional[List[str]] = None, + add_task_examples: bool = True, +): + """ + Build ERA-aligned system prompt for EB-ALFRED. - Task instruction and available actions are now in the system prompt, - so the initial observation only contains the image. + The prompt exactly matches what the ERA EPL-Only model was trained on, + so the model stays in-distribution. Output normalisation (ERA tokens → + VAGEN tags) happens downstream in normalize_era_tokens(). """ - return f"""[Current Observation]: -{img_str} + if action_list is not None: + max_id = len(action_list) - 1 + available = ", ".join( + f"[{i}, '{a}']" for i, a in enumerate(action_list) + ) + else: + max_id = "?" + available = "(not yet available)" -Decide your next action.""" + return ERA_SYSTEM_PROMPT_TEMPLATE.format( + max_action_id=max_id, + available_actions=available, + ) -def action_template(last_action, env_feedback, img_str, task_instruction=None): - """Template for step observation with feedback. +def init_observation_template(img_str, task_instruction=None): + """ERA-style initial user message. - Encourages structured reasoning: analyze the feedback, - reflect on why the last action succeeded or failed, - then plan the next logical step. + ERA format: \\n instruction: {task} \\n interaction_history: [] ... """ - task_line = f"\n[Task]: {task_instruction}\n" if task_instruction else "" - return f"""[Last Action]: {last_action} -[Feedback]: {env_feedback} -{task_line} -[Current Observation]: -{img_str} - -You MUST first analyze the feedback above. If the action succeeded, plan the next logical step to complete the task. If it failed, explain why and try a different approach. Do NOT repeat the same failed action.""" + inst = task_instruction or "" + return ( + f"{img_str}\n" + f" instruction: {inst} \n" + f" interaction_history: [] \n" + "Based on the above information, please provide the action " + "for the next step to complete the task. Think, then act." + ) + + +def action_template(last_action, env_feedback, img_str, task_instruction=None, + step_id=0, thinking="", action_id=None): + """ERA-style step user message with structured interaction history. + + Matches ERA's exact format: interaction_history is a list of dicts with + step_id, thinking, action [id, name], and env_feedback. + """ + if action_id is not None: + action_field = [action_id, last_action] + else: + action_field = last_action + history = [{"step_id": step_id, "thinking": thinking, + "action": action_field, "env_feedback": env_feedback}] + inst = task_instruction or "" + return ( + f"{img_str}\n" + f" instruction: {inst} \n" + f" interaction_history: {history} \n" + "Based on the above information, please provide the action " + "for the next step to complete the task. Think, then act." + ) def format_prompt(max_actions_per_step, action_sep, add_example=True, prompt_format="free_think"): - """Generate format prompt based on the specified format.""" + """Generate format prompt based on the specified format. + + For ERA-aligned mode, the generation guide is already in the system + prompt, so we return a minimal reminder. + """ if prompt_format == "free_think": return free_think_format_prompt(max_actions_per_step, action_sep, add_example) elif prompt_format == "wm": @@ -129,54 +119,16 @@ def format_prompt(max_actions_per_step, action_sep, add_example=True, prompt_for def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True): - """Generate format prompt for free_think format.""" - if max_actions_per_step == 1: - base = """You should output 1 action at a time. -Output the action as [action_id, action_name] using the ID from the available actions list. -Your response should be in the format of: -...[N, 'action name']""" - else: - base = f"""You should output a plan of up to {max_actions_per_step} actions at a time, separated by "{action_sep}". -Output each action as [action_id, action_name] using the ID from the available actions list. -Your response should be in the format of: -...[N1, 'action1']{action_sep} [N2, 'action2']{action_sep} ...""" - - if add_example: - if max_actions_per_step == 1: - examples = """ -Example 1: -I need to find a mug first. Let me navigate to where mugs might be. -[5, 'find a Mug'] - -Example 2: -The mug is nearby and I'm not holding anything. I should pick it up. -[12, 'pick up the Mug'] - -Example 3: -I'm holding the mug and I'm near the table. Let me put it down. -[38, 'put down the object in hand']""" - else: - examples = f""" -Example 1 (multi-step plan): -I need to find the alarm clock, pick it up, then find the desk lamp and turn it on. -[3, 'find a AlarmClock']{action_sep} [15, 'pick up the AlarmClock']{action_sep} [7, 'find a DeskLamp']{action_sep} [42, 'turn on the DeskLamp'] - -Example 2 (single action when unsure): -I am not sure where the mug is. Let me find it first. -[5, 'find a Mug'] - -Example 3 (replanning after failure): -The last action failed because the cabinet was closed. I need to open it first, then pick up the object. -[20, 'open the Cabinet']{action_sep} [12, 'pick up the Mug']""" - return base + "\n" + examples - - return base + """Minimal format prompt — the ERA system prompt already has the Generation Guide.""" + # The ERA system prompt already instructs the model on output format. + # Adding extra format instructions can confuse an SFT model, so keep it short. + return "" def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): - """Generate format prompt for wm format with observation and prediction tags.""" + """World-model format prompt (not used by ERA, kept for compatibility).""" base = f"""You should output {max_actions_per_step} action(s) at a time. -Output the action as [action_id, action_name] using the ID from the available actions list. +Output the action as [action_id, 'action_name'] using the ID from the available actions list. Your response must be in the format of: ......[N, 'action name'].... diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py index e01b5e339..843237738 100644 --- a/vagen/envs/eb_alfred/utils/utils.py +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -109,13 +109,27 @@ def normalize_era_tokens(response: str) -> str: Convert ERA special tokens to VAGEN plain tags so the parser can handle models trained with ERA's SFT format. - ERA format: <|think_start|>...<|think_end|><|action_start|>...<|action_end|> - VAGEN format: ...... + ERA model outputs multi-turn style: + <|think_start|>...<|think_end|><|im_end|>\n<|im_start|>assistant\n<|think_start|>[id, 'action']<|action_end|> + We need to convert this to: + ...[id, 'action'] """ - response = response.replace("<|think_start|>", "") - response = response.replace("<|think_end|>", "") + # Strip multi-turn separators the model sometimes emits + response = re.sub(r'<\|im_end\|>\s*<\|im_start\|>assistant\s*', '', response) + # The model may use <|think_start|> for both think and action blocks; + # after stripping im tokens, the second <|think_start|> is the action block response = response.replace("<|action_start|>", "") response = response.replace("<|action_end|>", "") + # Convert think tokens — but the second one (action) should become + parts = response.split("<|think_start|>") + if len(parts) >= 3: + # parts[0]=before, parts[1]=think content, parts[2]=action content + think_part = parts[1].replace("<|think_end|>", "") + action_part = parts[2] + response = f"{think_part.strip()}{action_part}" + else: + response = response.replace("<|think_start|>", "") + response = response.replace("<|think_end|>", "") return response From 1ce537e28922a4259fc5d7b59287402254e8f5d4 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Fri, 20 Mar 2026 10:56:06 -0400 Subject: [PATCH 17/29] update script Co-Authored-By: Claude Opus 4.6 (1M context) --- examples/evaluate/eb_alfred/config.yaml | 59 +++++++++++++++++++------ 1 file changed, 45 insertions(+), 14 deletions(-) diff --git a/examples/evaluate/eb_alfred/config.yaml b/examples/evaluate/eb_alfred/config.yaml index 1d43c6eb4..271ff374e 100644 --- a/examples/evaluate/eb_alfred/config.yaml +++ b/examples/evaluate/eb_alfred/config.yaml @@ -1,22 +1,50 @@ fileroot: ${oc.env:HOME}/projects/vagen envs: + - name: RemoteEnv + n_envs: 51 + data_source: eb_alfred + tag_id: eb_alfred_train_base + seed: [0, 51, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: base + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + - name: RemoteEnv n_envs: 50 data_source: eb_alfred - tag_id: eb_alfred_val_common + tag_id: eb_alfred_train_complex seed: [0, 50, 1] - max_turns: 6 + max_turns: 30 + concat_multi_turn: false config: base_urls: - "http://localhost:8000" timeout: 600 - eval_set: common_sense + eval_set: complex_instruction obs_image_size: 500 - max_turns: 6 + max_turns: 30 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -25,23 +53,25 @@ envs: temperature: 0 max_tokens: 2048 top_p: 1.0 + stop: ["<|diff_marker|>"] - name: RemoteEnv n_envs: 50 data_source: eb_alfred - tag_id: eb_alfred_val_spatial + tag_id: eb_alfred_train_visual seed: [0, 50, 1] - max_turns: 6 + max_turns: 30 + concat_multi_turn: false config: base_urls: - "http://localhost:8000" timeout: 600 - eval_set: spatial + eval_set: visual_appearance obs_image_size: 500 - max_turns: 6 + max_turns: 30 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -50,13 +80,14 @@ envs: temperature: 0 max_tokens: 2048 top_p: 1.0 + stop: ["<|diff_marker|>"] experiment: dump_dir: ${fileroot}/rollouts/eval_eb_alfred - default_max_turns: 6 + default_max_turns: 30 run: - backend: "openai" + backend: "sglang" base_seed: 0 max_concurrent_jobs: 4 resume: skip_completed @@ -66,7 +97,7 @@ backends: openai: api_key: "" # or env OPENAI_API_KEY base_url: null - model: "gpt-4o-mini" + model: "gpt-4o" max_concurrency: 2 max_retries: 6 min_backoff: 0.5 @@ -75,7 +106,7 @@ backends: sglang: base_url: "http://127.0.0.1:30000/v1" api_key: "EMPTY" - model: "Qwen/Qwen2.5-VL-7B-Instruct" + model: "/home/march/workspace/Yaning/models/EPL-Only-Model_EB-Alfred" max_concurrency: 2 max_retries: 6 min_backoff: 0.5 From b76736141739392905842355d405af24a27df5c5 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Fri, 20 Mar 2026 10:58:58 -0400 Subject: [PATCH 18/29] update script Co-Authored-By: Claude Opus 4.6 (1M context) --- examples/evaluate/eb_alfred/config.yaml | 90 +++++++--- .../evaluate/eb_alfred/config_era_all.yaml | 158 ++++++++++++++++++ 2 files changed, 221 insertions(+), 27 deletions(-) create mode 100644 examples/evaluate/eb_alfred/config_era_all.yaml diff --git a/examples/evaluate/eb_alfred/config.yaml b/examples/evaluate/eb_alfred/config.yaml index 271ff374e..c67db66a0 100644 --- a/examples/evaluate/eb_alfred/config.yaml +++ b/examples/evaluate/eb_alfred/config.yaml @@ -1,10 +1,10 @@ -fileroot: ${oc.env:HOME}/projects/vagen +fileroot: ${oc.env:HOME}/workspace/Yaning/VAGEN envs: - name: RemoteEnv n_envs: 51 data_source: eb_alfred - tag_id: eb_alfred_train_base + tag_id: eb_alfred_eval_era_base seed: [0, 51, 1] max_turns: 30 concat_multi_turn: false @@ -31,7 +31,34 @@ envs: - name: RemoteEnv n_envs: 50 data_source: eb_alfred - tag_id: eb_alfred_train_complex + tag_id: eb_alfred_eval_era_common_sense + seed: [0, 50, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: common_sense + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_complex_instruction seed: [0, 50, 1] max_turns: 30 concat_multi_turn: false @@ -58,7 +85,34 @@ envs: - name: RemoteEnv n_envs: 50 data_source: eb_alfred - tag_id: eb_alfred_train_visual + tag_id: eb_alfred_eval_era_spatial + seed: [0, 50, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: spatial + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_visual_appearance seed: [0, 50, 1] max_turns: 30 concat_multi_turn: false @@ -83,40 +137,22 @@ envs: stop: ["<|diff_marker|>"] experiment: - dump_dir: ${fileroot}/rollouts/eval_eb_alfred + dump_dir: ${fileroot}/rollouts/eval_eb_alfred_era default_max_turns: 30 run: - backend: "sglang" + backend: "openai" base_seed: 0 - max_concurrent_jobs: 4 + max_concurrent_jobs: 1 resume: skip_completed live_summary: true backends: openai: - api_key: "" # or env OPENAI_API_KEY - base_url: null - model: "gpt-4o" - max_concurrency: 2 - max_retries: 6 - min_backoff: 0.5 - max_backoff: 8.0 - - sglang: base_url: "http://127.0.0.1:30000/v1" api_key: "EMPTY" - model: "/home/march/workspace/Yaning/models/EPL-Only-Model_EB-Alfred" - max_concurrency: 2 - max_retries: 6 - min_backoff: 0.5 - max_backoff: 8.0 - - claude: - api_key: "" - base_url: null - model: "claude-3-5-sonnet-latest" - max_concurrency: 2 + model: "EPL-Only-Model_EB-Alfred" + max_concurrency: 1 max_retries: 6 min_backoff: 0.5 max_backoff: 8.0 diff --git a/examples/evaluate/eb_alfred/config_era_all.yaml b/examples/evaluate/eb_alfred/config_era_all.yaml new file mode 100644 index 000000000..c67db66a0 --- /dev/null +++ b/examples/evaluate/eb_alfred/config_era_all.yaml @@ -0,0 +1,158 @@ +fileroot: ${oc.env:HOME}/workspace/Yaning/VAGEN + +envs: + - name: RemoteEnv + n_envs: 51 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_base + seed: [0, 51, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: base + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_common_sense + seed: [0, 50, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: common_sense + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_complex_instruction + seed: [0, 50, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: complex_instruction + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_spatial + seed: [0, 50, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: spatial + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + + - name: RemoteEnv + n_envs: 50 + data_source: eb_alfred + tag_id: eb_alfred_eval_era_visual_appearance + seed: [0, 50, 1] + max_turns: 30 + concat_multi_turn: false + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: visual_appearance + obs_image_size: 500 + max_turns: 30 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "|" + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + chat_config: + temperature: 0 + max_tokens: 2048 + top_p: 1.0 + stop: ["<|diff_marker|>"] + +experiment: + dump_dir: ${fileroot}/rollouts/eval_eb_alfred_era + default_max_turns: 30 + +run: + backend: "openai" + base_seed: 0 + max_concurrent_jobs: 1 + resume: skip_completed + live_summary: true + +backends: + openai: + base_url: "http://127.0.0.1:30000/v1" + api_key: "EMPTY" + model: "EPL-Only-Model_EB-Alfred" + max_concurrency: 1 + max_retries: 6 + min_backoff: 0.5 + max_backoff: 8.0 From 698a3f4fc9e659d7d3bb1c8b662ddcc0271dc67e Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Fri, 20 Mar 2026 13:40:29 -0400 Subject: [PATCH 19/29] udpate parse logic --- examples/evaluate/eb_alfred/config.yaml | 4 +- vagen/envs/eb_alfred/utils/utils.py | 60 ++++++++++++++++--------- 2 files changed, 42 insertions(+), 22 deletions(-) diff --git a/examples/evaluate/eb_alfred/config.yaml b/examples/evaluate/eb_alfred/config.yaml index c67db66a0..5be4b0ad8 100644 --- a/examples/evaluate/eb_alfred/config.yaml +++ b/examples/evaluate/eb_alfred/config.yaml @@ -141,14 +141,14 @@ experiment: default_max_turns: 30 run: - backend: "openai" + backend: "sglang" base_seed: 0 max_concurrent_jobs: 1 resume: skip_completed live_summary: true backends: - openai: + sglang: base_url: "http://127.0.0.1:30000/v1" api_key: "EMPTY" model: "EPL-Only-Model_EB-Alfred" diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py index 843237738..8a030b849 100644 --- a/vagen/envs/eb_alfred/utils/utils.py +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -109,27 +109,47 @@ def normalize_era_tokens(response: str) -> str: Convert ERA special tokens to VAGEN plain tags so the parser can handle models trained with ERA's SFT format. - ERA model outputs multi-turn style: - <|think_start|>...<|think_end|><|im_end|>\n<|im_start|>assistant\n<|think_start|>[id, 'action']<|action_end|> - We need to convert this to: - ...[id, 'action'] + Handles two cases: + 1. Special tokens present (e.g. from Flask server / transformers): + <|think_start|>...<|think_end|><|im_end|>\n<|im_start|>assistant\n<|think_start|>[id, 'action']<|action_end|> + 2. Special tokens stripped by sglang/vllm (plain text separators): + visual_description: ... reasoning: ... language_plan: ...\nassistant\n[id, 'action'] + Both are converted to: ...[id, 'action'] """ - # Strip multi-turn separators the model sometimes emits - response = re.sub(r'<\|im_end\|>\s*<\|im_start\|>assistant\s*', '', response) - # The model may use <|think_start|> for both think and action blocks; - # after stripping im tokens, the second <|think_start|> is the action block - response = response.replace("<|action_start|>", "") - response = response.replace("<|action_end|>", "") - # Convert think tokens — but the second one (action) should become - parts = response.split("<|think_start|>") - if len(parts) >= 3: - # parts[0]=before, parts[1]=think content, parts[2]=action content - think_part = parts[1].replace("<|think_end|>", "") - action_part = parts[2] - response = f"{think_part.strip()}{action_part}" - else: - response = response.replace("<|think_start|>", "") - response = response.replace("<|think_end|>", "") + # Case 1: special tokens present + if "<|think_start|>" in response or "<|action_start|>" in response: + response = re.sub(r'<\|im_end\|>\s*<\|im_start\|>assistant\s*', '', response) + response = response.replace("<|action_start|>", "") + response = response.replace("<|action_end|>", "") + parts = response.split("<|think_start|>") + if len(parts) >= 3: + think_part = parts[1].replace("<|think_end|>", "") + action_part = parts[2] + response = f"{think_part.strip()}{action_part}" + else: + response = response.replace("<|think_start|>", "") + response = response.replace("<|think_end|>", "") + return response + + # Case 2: sglang/vllm stripped special tokens, plain text separators + # Pattern: "thinking text\nassistant\n[id, 'action']" or just "[id, 'action']" + # Try splitting on "\nassistant\n" + assistant_split = re.split(r'\nassistant\s*\n', response) + if len(assistant_split) >= 2: + think_part = assistant_split[0].strip() + action_part = assistant_split[-1].strip() + return f"{think_part}{action_part}" + + # Fallback: look for [id, 'action'] at the end + action_match = re.search(r'(\[(\d+),\s*[\'"]?.+?[\'"]?\s*\])\s*$', response) + if action_match: + think_part = response[:action_match.start()].strip() + action_part = action_match.group(1) + if think_part: + return f"{think_part}{action_part}" + return f"{action_part}" + + # Nothing matched, return as-is (will fail format_correct) return response From fcdb7e9954bc6e4f128f0d90faf0c19c7979747b Mon Sep 17 00:00:00 2001 From: YaningDylan <20082585d@gmail.com> Date: Fri, 20 Mar 2026 22:33:35 +0000 Subject: [PATCH 20/29] fix --- .../train_ppo_no_concat_qwen25vl3b.sh | 12 ++-- .../train/eb_alfred/val_eb_alfred_vision.yaml | 70 ++++++++++++++++++- vagen/envs/eb_alfred/README.md | 15 ++++ vagen/gym_agent_dataset.py | 3 +- 4 files changed, 91 insertions(+), 9 deletions(-) diff --git a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh index 9bf8fa42c..341071f7f 100755 --- a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh +++ b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh @@ -24,10 +24,10 @@ PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ data.train_files=${DATASET_TRAIN} \ data.val_files=${DATASET_VAL} \ data.train_batch_size=128 \ - data.max_prompt_length=2048 \ + data.max_prompt_length=5000 \ data.max_response_length=512 \ - +data.max_trajectory_length=7000 \ - algorithm.adv_estimator=no_concat_gae_first \ + +data.max_trajectory_length=100000 \ + algorithm.adv_estimator=no_concat_gae \ algorithm.kl_ctrl.kl_coef=0.0 \ actor_rollout_ref.model.path=${REF_MODEL_PATH} \ actor_rollout_ref.model.use_remove_padding=True \ @@ -51,7 +51,7 @@ PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ actor_rollout_ref.rollout.mode=async \ actor_rollout_ref.rollout.n=1 \ actor_rollout_ref.rollout.max_num_batched_tokens=10000 \ - actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \ actor_rollout_ref.rollout.enforce_eager=True \ actor_rollout_ref.rollout.free_cache_engine=True \ actor_rollout_ref.rollout.enable_chunked_prefill=True \ @@ -71,7 +71,7 @@ PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ trainer.critic_warmup=0 \ trainer.logger=['console','wandb'] \ trainer.val_before_train=True \ - trainer.n_gpus_per_node=4 \ + trainer.n_gpus_per_node=8 \ trainer.nnodes=1 \ trainer.save_freq=100 \ trainer.test_freq=20 \ @@ -81,6 +81,6 @@ PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ trainer.validation_data_dir=${EXPERIMENT_DIR}/validation \ trainer.rollout_data_dir=${EXPERIMENT_DIR}/rollout_data \ trainer.log_val_generations=32 \ - +trainer.concat_multi_turn=False \ + trainer.concat_multi_turn=False \ trainer.total_training_steps=400 2>&1 | \ tee ${EXPERIMENT_DIR}/${PROJECT_NAME}_${EXPERIMENT_NAME}.log >(tee ${BASEDIR}/${PROJECT_NAME}_${EXPERIMENT_NAME}.log >/dev/null) diff --git a/examples/train/eb_alfred/val_eb_alfred_vision.yaml b/examples/train/eb_alfred/val_eb_alfred_vision.yaml index 233fc3e21..17deb8e80 100644 --- a/examples/train/eb_alfred/val_eb_alfred_vision.yaml +++ b/examples/train/eb_alfred/val_eb_alfred_vision.yaml @@ -1,6 +1,6 @@ envs: - name: RemoteEnv - n_envs: 50 + n_envs: 10 data_source: eb_alfred tag_id: eb_alfred_val_common seed: [0, 50, 1] @@ -22,7 +22,7 @@ envs: success_reward: 1.0 - name: RemoteEnv - n_envs: 50 + n_envs: 10 data_source: eb_alfred tag_id: eb_alfred_val_spatial seed: [0, 50, 1] @@ -42,3 +42,69 @@ envs: use_example_in_sys_prompt: true format_reward: 0.1 success_reward: 1.0 + + - name: RemoteEnv + n_envs: 10 + data_source: eb_alfred + tag_id: eb_alfred_train_base + seed: [0, 200, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: base + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + + - name: RemoteEnv + n_envs: 10 + data_source: eb_alfred + tag_id: eb_alfred_train_complex + seed: [0, 200, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: complex_instruction + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 + + - name: RemoteEnv + n_envs: 10 + data_source: eb_alfred + tag_id: eb_alfred_train_visual + seed: [0, 200, 1] + max_turns: 6 + response_length_per_turn: 512 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: visual_appearance + obs_image_size: 500 + max_turns: 6 + max_actions_per_step: 20 + max_env_steps: 30 + action_sep: "," + prompt_format: free_think + use_example_in_sys_prompt: true + format_reward: 0.1 + success_reward: 1.0 diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md index 99b6f7a97..e5c8af0be 100644 --- a/vagen/envs/eb_alfred/README.md +++ b/vagen/envs/eb_alfred/README.md @@ -22,6 +22,9 @@ pip install "ai2thor==2.1.0" "gym==0.23.0" "numpy<2.0" \ pip install "flask==1.1.4" "werkzeug==1.0.1" \ "markupsafe<2.1" "jinja2<3.0" "itsdangerous<2.0" pip install "opencv-python-headless<4.9" +pip install fire uvicorn httpx fastapi python-multipart +apt-get install -y xorg + ``` **Download dataset** (`eval_set` selects the split: `base` — standard tasks, `long` — longer horizon): @@ -38,6 +41,18 @@ conda activate embodiedbench python -m vagen.envs.eb_alfred.serve ``` +> **Note (this machine):** The `embodiedbench` conda env is at `/venv/embodiedbench`. Use `PYTHONPATH` to inject the VAGEN repo without reinstalling and risking version conflicts: +> +> ```bash +> # Extra dep needed once +> /venv/embodiedbench/bin/pip install hydra-core +> +> # Start server +> PYTHONPATH=/workspace/VAGEN /venv/embodiedbench/bin/python -m vagen.envs.eb_alfred.serve +> ``` +> +> AI2-THOR will auto-download Unity (~390MB) on first connect. Verified working: server starts, `/connect` returns session_id, Unity initializes on display `:3`, session reaches `env ready` state. + Key parameters: - `devices`: GPU indices (default: auto-detect via `CUDA_VISIBLE_DEVICES` or `nvidia-smi`) - `capacity`: max concurrent Unity environments (default: 16) diff --git a/vagen/gym_agent_dataset.py b/vagen/gym_agent_dataset.py index 24de20b39..3d7545633 100644 --- a/vagen/gym_agent_dataset.py +++ b/vagen/gym_agent_dataset.py @@ -40,7 +40,8 @@ class EnvSpecs: def load_envspecs(yaml_path: str) -> EnvSpecs: print(yaml_path) cfg = OmegaConf.load(yaml_path) - specs = [EnvSpec(**OmegaConf.to_container(s, resolve=True)) for s in cfg.get("envs", [])] + valid_fields = {f.name for f in EnvSpec.__dataclass_fields__.values()} + specs = [EnvSpec(**{k: v for k, v in OmegaConf.to_container(s, resolve=True).items() if k in valid_fields}) for s in cfg.get("envs", [])] return EnvSpecs(specs=specs) # Upper bound used for RNG sampling when only a base seed is provided From 9639f0ba5ffb3f054f7af90838e06a8fa2edae67 Mon Sep 17 00:00:00 2001 From: Jingnan Ma Date: Sun, 22 Mar 2026 13:20:52 -0500 Subject: [PATCH 21/29] tmp --- .../submit_ppo_no_concat_qwen25vl3b.sh | 40 +++++++++++++++++++ .../train_ppo_no_concat_qwen25vl3b.sh | 2 +- 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100755 examples/train/eb_alfred/submit_ppo_no_concat_qwen25vl3b.sh diff --git a/examples/train/eb_alfred/submit_ppo_no_concat_qwen25vl3b.sh b/examples/train/eb_alfred/submit_ppo_no_concat_qwen25vl3b.sh new file mode 100755 index 000000000..8fbd6a25c --- /dev/null +++ b/examples/train/eb_alfred/submit_ppo_no_concat_qwen25vl3b.sh @@ -0,0 +1,40 @@ +#!/bin/bash +#SBATCH --job-name=vagen_ppo_eb_alfred +#SBATCH --partition=gpuA100x4 +#SBATCH --account=bgig-delta +#SBATCH --nodes=1 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=64 +#SBATCH --mem=200G +#SBATCH --time=48:00:00 +#SBATCH --output=%x_%j.out +#SBATCH --error=%x_%j.err + +# --------------------------------------------------------------- +# Before submitting, start the env server on your local machine: +# python -m vagen.envs.eb_alfred.serve --port 8000 +# +# Then create a reverse SSH tunnel from your local machine to the +# Delta login node so the compute node can reach your env server: +# ssh -R 8000:localhost:8000 jma6@dt-login01.delta.ncsa.illinois.edu +# +# Keep that tunnel open for the duration of the job. +# --------------------------------------------------------------- + +set -x + +# Load modules +source /sw/rh9.4/python/miniforge3/etc/profile.d/conda.sh +module load miniforge3-python +conda activate /scratch/bgig/jma6/envs/vagen + +# Forward the login node's port 8000 to this compute node's localhost:8000 +# so the training script can connect to the env server via localhost. +ssh -f -N -L 8000:localhost:8000 dt-login01.delta.ncsa.illinois.edu +echo "SSH tunnel established: localhost:8000 -> dt-login01:8000 -> your local env server" + +# Wait a moment for tunnel to be ready +sleep 3 + +cd /u/jma6/workspace/VAGEN +bash examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh diff --git a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh index 341071f7f..6c1b9ded0 100755 --- a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh +++ b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh @@ -12,7 +12,7 @@ SAVE_CHECKPOINT_DIR=${EXPERIMENT_DIR}/verl_checkpoints DATASET_TRAIN=${SCRIPTDIR}/train_eb_alfred_vision.yaml DATASET_VAL=${SCRIPTDIR}/val_eb_alfred_vision.yaml agent_loop_config_path=${BASEDIR}/vagen/configs/agent_no_concat.yaml -REF_MODEL_PATH=Qwen/Qwen2.5-VL-3B-Instruct +REF_MODEL_PATH=/scratch/bgig/jma6/Qwen2.5-VL-3B-Instruct-EB-ALFRED-vagen-sft mkdir -p ${EXPERIMENT_DIR} export HF_HOME=/workspace/.hf_home From f869c156cf629b59f0282565b49582740ede1ef6 Mon Sep 17 00:00:00 2001 From: Jingnan Ma Date: Sun, 22 Mar 2026 13:42:24 -0500 Subject: [PATCH 22/29] tmp --- examples/train/eb_alfred/train_eb_alfred_vision.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/train/eb_alfred/train_eb_alfred_vision.yaml b/examples/train/eb_alfred/train_eb_alfred_vision.yaml index d471ab1f6..2af2d3822 100644 --- a/examples/train/eb_alfred/train_eb_alfred_vision.yaml +++ b/examples/train/eb_alfred/train_eb_alfred_vision.yaml @@ -15,7 +15,7 @@ envs: max_turns: 6 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -37,7 +37,7 @@ envs: max_turns: 6 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -59,7 +59,7 @@ envs: max_turns: 6 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 From 137a9d6629e558655335e18ae01bfa6dc3987ca5 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Tue, 24 Mar 2026 17:03:40 -0400 Subject: [PATCH 23/29] update reset logic Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/test_env_pool.py | 271 ++++++++++++++++++++++++++++++++ vagen/envs/eb_alfred/handler.py | 121 ++++++++++---- vagen/envs/eb_alfred/serve.py | 6 +- 3 files changed, 366 insertions(+), 32 deletions(-) create mode 100644 tests/test_env_pool.py diff --git a/tests/test_env_pool.py b/tests/test_env_pool.py new file mode 100644 index 000000000..c5b8a74b2 --- /dev/null +++ b/tests/test_env_pool.py @@ -0,0 +1,271 @@ +""" +Test env pool logic in EbAlfredHandler. + +Uses a mock env to verify pool lifecycle without needing AI2-THOR. +""" +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import sys, os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from vagen.envs.eb_alfred.handler import EbAlfredHandler + + +class FakeEnv: + """Mock env that tracks create/close calls.""" + _count = 0 + + def __init__(self): + FakeEnv._count += 1 + self.id = FakeEnv._count + self.closed = False + self._assigned_display = "0" + + async def close(self): + self.closed = True + + async def reset(self, seed): + return {"obs_str": f"obs-{seed}"}, {"seed": seed} + + async def system_prompt(self): + return {"obs_str": "system prompt"} + + async def step(self, action): + return {"obs_str": "step"}, 0.0, False, {} + + +async def test_pool_basic(): + """Verify envs are pooled on close and reused on connect.""" + FakeEnv._count = 0 + handler = EbAlfredHandler( + x_displays=["0"], + capacity=4, + startup_concurrency=4, + pool_size=4, + ) + + # Monkey-patch create_env to return FakeEnv + async def fake_create(config): + env = FakeEnv() + env._assigned_display = "0" + return env + + handler.create_env = fake_create + + # Create 4 sessions + results = [] + for i in range(4): + r = await handler.connect({"eval_set": "base"}, seed=i) + results.append(r) + + # Wait for all envs to be ready + for sid, ctx in handler._sessions.items(): + if hasattr(ctx, '_ready') and ctx._ready: + await ctx._ready.wait() + + assert len(handler._sessions) == 4, f"Expected 4 sessions, got {len(handler._sessions)}" + assert FakeEnv._count == 4, f"Expected 4 envs created, got {FakeEnv._count}" + assert len(handler._env_pool) == 0 + + # Close all 4 → should go to pool + sids = list(handler._sessions.keys()) + for sid in sids: + ctx = handler._sessions[sid] + await handler._handle_close(ctx) + + assert len(handler._sessions) == 0 + assert len(handler._env_pool) == 4, f"Expected 4 pooled, got {len(handler._env_pool)}" + assert FakeEnv._count == 4, "No new envs should be created" + + # Create 4 more sessions → should reuse from pool + for i in range(4): + await handler.connect({"eval_set": "base"}, seed=i + 100) + + # Wait for all envs to be ready + for sid, ctx in handler._sessions.items(): + if hasattr(ctx, '_ready') and ctx._ready: + await ctx._ready.wait() + + assert len(handler._sessions) == 4 + assert len(handler._env_pool) == 0, f"Pool should be empty, got {len(handler._env_pool)}" + assert FakeEnv._count == 4, f"Should reuse, not create new! Got {FakeEnv._count}" + + print("PASS: test_pool_basic") + + +async def test_pool_overflow(): + """When pool is full, env should be actually closed.""" + FakeEnv._count = 0 + handler = EbAlfredHandler( + x_displays=["0"], + capacity=4, + startup_concurrency=4, + pool_size=2, # Only keep 2 in pool + ) + + async def fake_create(config): + env = FakeEnv() + env._assigned_display = "0" + return env + + handler.create_env = fake_create + + # Create 4 sessions + for i in range(4): + await handler.connect({"eval_set": "base"}, seed=i) + + for sid, ctx in handler._sessions.items(): + if hasattr(ctx, '_ready') and ctx._ready: + await ctx._ready.wait() + + # Close all 4 → first 2 pooled, last 2 actually closed + sids = list(handler._sessions.keys()) + closed_envs = [] + for sid in sids: + ctx = handler._sessions[sid] + if ctx.env: + closed_envs.append(ctx.env) + await handler._handle_close(ctx) + + assert len(handler._env_pool) == 2, f"Expected 2 pooled, got {len(handler._env_pool)}" + actually_closed = sum(1 for e in closed_envs if e.closed) + assert actually_closed == 2, f"Expected 2 closed, got {actually_closed}" + + print("PASS: test_pool_overflow") + + +async def test_no_deadlock_with_queuing(): + """With batch_size > capacity, verify no deadlock: queued sessions + should be served as envs are pooled and permits released.""" + FakeEnv._count = 0 + handler = EbAlfredHandler( + x_displays=["0"], + capacity=2, + startup_concurrency=2, + pool_size=2, + ) + + async def fake_create(config): + await asyncio.sleep(0.01) # Simulate short startup + env = FakeEnv() + env._assigned_display = "0" + return env + + handler.create_env = fake_create + + # Create 4 sessions (capacity=2, so 2 will queue) + connect_tasks = [] + for i in range(4): + connect_tasks.append(handler.connect({"eval_set": "base"}, seed=i)) + await asyncio.gather(*connect_tasks) + + assert len(handler._sessions) == 4 + + # Wait for first 2 to be ready + await asyncio.sleep(0.1) + ready_count = sum(1 for ctx in handler._sessions.values() if ctx.env is not None) + assert ready_count == 2, f"Expected 2 ready, got {ready_count}" + + # Close 1 session → frees permit → queued session should get env from pool + first_sid = None + for sid, ctx in handler._sessions.items(): + if ctx.env is not None: + first_sid = sid + break + await handler._handle_close(handler._sessions[first_sid]) + + await asyncio.sleep(0.05) # Let queued task run + + # Now should have 2 ready (1 original + 1 newly unblocked that reused pool) + ready_count = sum(1 for ctx in handler._sessions.values() if ctx.env is not None) + assert ready_count == 2, f"Expected 2 ready after close+reuse, got {ready_count}" + # Pool should have been used (one env went in, one came out) + assert FakeEnv._count <= 3, f"Should reuse pool, only created {FakeEnv._count}" + + # Close another → unblock last queued session too + second_sid = None + for sid, ctx in handler._sessions.items(): + if ctx.env is not None: + second_sid = sid + break + await handler._handle_close(handler._sessions[second_sid]) + await asyncio.sleep(0.05) + + # Both remaining sessions should now be ready (last queued got unblocked) + ready_count = sum(1 for ctx in handler._sessions.values() if ctx.env is not None) + assert ready_count == 2, f"Expected 2 ready, got {ready_count}" + + # Cleanup + for sid in list(handler._sessions.keys()): + ctx = handler._sessions[sid] + await handler._handle_close(ctx) + + total_created = FakeEnv._count + print(f"PASS: test_no_deadlock_with_queuing (created {total_created} envs for 4 sessions)") + + +async def test_batch_cycle(): + """Simulate 2 training batches: batch_size=4, capacity=2. + Second batch should reuse all pooled envs.""" + FakeEnv._count = 0 + handler = EbAlfredHandler( + x_displays=["0"], + capacity=2, + startup_concurrency=2, + pool_size=2, + ) + + async def fake_create(config): + await asyncio.sleep(0.05) # Simulate startup + env = FakeEnv() + env._assigned_display = "0" + return env + + handler.create_env = fake_create + + async def run_episode(handler, seed): + """Simulate one episode: connect → wait ready → close.""" + result = await handler.connect({"eval_set": "base"}, seed=seed) + sid = result.data["session_id"] + ctx = handler._sessions[sid] + if hasattr(ctx, '_ready') and ctx._ready: + await ctx._ready.wait() + # Simulate some work + await asyncio.sleep(0.02) + await handler._handle_close(ctx) + + # Batch 1: 4 episodes + t0 = time.time() + await asyncio.gather(*[run_episode(handler, i) for i in range(4)]) + batch1_time = time.time() - t0 + batch1_created = FakeEnv._count + + assert len(handler._env_pool) == 2, f"Expected 2 pooled after batch 1, got {len(handler._env_pool)}" + + # Batch 2: 4 more episodes → should reuse pool + t0 = time.time() + await asyncio.gather(*[run_episode(handler, i + 100) for i in range(4)]) + batch2_time = time.time() - t0 + batch2_created = FakeEnv._count - batch1_created + + print(f" Batch 1: created {batch1_created} envs, took {batch1_time:.3f}s") + print(f" Batch 2: created {batch2_created} envs, took {batch2_time:.3f}s") + assert batch2_created == 0, f"Batch 2 should create 0 new envs, created {batch2_created}" + + # Cleanup + await handler.aclose() + print("PASS: test_batch_cycle") + + +async def main(): + await test_pool_basic() + await test_pool_overflow() + await test_no_deadlock_with_queuing() + await test_batch_cycle() + print("\nAll tests passed!") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/vagen/envs/eb_alfred/handler.py b/vagen/envs/eb_alfred/handler.py index 7cd30b8b9..0505060fc 100644 --- a/vagen/envs/eb_alfred/handler.py +++ b/vagen/envs/eb_alfred/handler.py @@ -65,7 +65,19 @@ def detect_gpu_displays() -> List[str]: class EbAlfredHandler(BaseGymHandler): - """Handler for EB-ALFRED with capacity-based queuing and multi-GPU load balancing. + """Handler for EB-ALFRED with capacity-based queuing, multi-GPU load + balancing, and **persistent env pooling**. + + Env pooling (enabled by default): + - When a session closes, its Unity process is returned to an idle + pool instead of being destroyed. + - When a new session connects, an idle env is taken from the pool + (near-instant) instead of spawning a new Unity process (~90 s). + - The pool implicitly holds capacity-semaphore permits: a pooled + env still counts against ``capacity`` because the Unity process + is alive and consuming GPU memory. + - Set ``pool_size=0`` together with ``capacity=0`` to disable + pooling entirely (original behaviour). When capacity > 0: - /connect returns session_id immediately (env creation is deferred) @@ -90,6 +102,7 @@ def __init__( x_displays: Optional[List[str]] = None, capacity: int = 16, startup_concurrency: int = 8, + pool_size: int = -1, **kwargs, ): """ @@ -100,6 +113,8 @@ def __init__( startup_concurrency: Max Unity processes that may be starting up at once (0 = unlimited). Prevents CPU spikes when many capacity slots open simultaneously. Ignored when capacity = 0. + pool_size: Max idle envs kept alive in the pool. -1 (default) = + same as capacity (keep every env alive). 0 = disable pooling. **kwargs: Passed to BaseGymHandler (session_timeout, max_sessions). """ super().__init__(**kwargs) @@ -107,6 +122,10 @@ def __init__( self._pending_counts: Dict[str, int] = {d: 0 for d in self._x_displays} self._capacity = capacity self._startup_concurrency = startup_concurrency + # Env pool: idle Unity processes available for immediate reuse. + # Each pooled env implicitly holds one capacity-semaphore permit. + self._pool_size = pool_size if pool_size >= 0 else max(capacity, 1) + self._env_pool: List[Any] = [] # Defer semaphore creation: it must be created on the running event loop, # not during __init__ (which runs before uvicorn starts the loop). self._capacity_sem: Optional[asyncio.Semaphore] = None @@ -114,7 +133,8 @@ def __init__( LOGGER.info( f"[Handler] Using X displays: {self._x_displays}, " f"capacity={capacity if capacity > 0 else 'unlimited'}, " - f"startup_concurrency={startup_concurrency if startup_concurrency > 0 else 'unlimited'}" + f"startup_concurrency={startup_concurrency if startup_concurrency > 0 else 'unlimited'}, " + f"pool_size={self._pool_size}" ) def _ensure_semaphore(self) -> None: @@ -221,19 +241,29 @@ async def connect( }) async def _deferred_create(self, ctx: _DeferredSessionContext) -> None: - """Background task: acquire capacity slot, then create env. - - Two-phase acquisition: - 1. capacity_sem – limits total running envs (held for env lifetime) - 2. startup_sem – limits concurrent Unity startups (held only during - EbAlfred.__init__, released as soon as the process - is running) - This prevents a "startup storm" when many capacity slots open at once. + """Background task: acquire capacity slot, then reuse pooled env or create new. + + Always acquires a capacity permit first (so queued sessions unblock + as soon as any session releases its permit via close/pool). After + acquiring the permit, checks the pool for an idle env: + - Pool hit → instant reuse, skip Unity startup + - Pool miss → two-phase creation with startup_sem throttle """ try: LOGGER.info(f"[Handler] Session {ctx.session_id} waiting for capacity slot...") await self._capacity_sem.acquire() ctx._holds_slot = True + + # ---- fast path: reuse from pool ---- + if self._env_pool: + ctx.env = self._env_pool.pop() + LOGGER.info( + f"[Handler] Session {ctx.session_id} reused pooled env " + f"(pool: {len(self._env_pool)} remaining)" + ) + return + + # ---- slow path: create new Unity process ---- LOGGER.info(f"[Handler] Session {ctx.session_id} acquired capacity slot, waiting for startup slot...") if self._startup_sem is not None: @@ -285,34 +315,61 @@ async def call( return await super().call(session_id, method, params, images) - async def _handle_close(self, ctx: SessionContext) -> HandlerResult: - """Close env and release capacity slot.""" + async def _release_env(self, ctx: SessionContext) -> None: + """Pool or close the env, then always release the capacity permit. + + The pool is a pure cache — it does NOT hold capacity permits. + This ensures queued sessions always unblock when a session closes, + regardless of whether the env was pooled or destroyed. + + When a new session later acquires a permit, it checks the pool + first (fast path) before creating a new Unity process (slow path). + """ try: - if ctx.env is not None: - await ctx.env.close() + if ctx.env is not None and len(self._env_pool) < self._pool_size: + # Return to pool — keep Unity alive for reuse + self._env_pool.append(ctx.env) + LOGGER.info( + f"[Handler] Session {ctx.session_id} returned env to pool " + f"(pool: {len(self._env_pool)}/{self._pool_size})" + ) + ctx.env = None + else: + # Pool full (or no env) — actually close Unity + if ctx.env is not None: + await ctx.env.close() + ctx.env = None except Exception as e: - LOGGER.error(f"[Handler] Error closing env for session {ctx.session_id}: {e}") + LOGGER.error(f"[Handler] Error releasing env for session {ctx.session_id}: {e}") finally: + # Always release capacity permit so queued sessions can proceed if ctx._holds_slot and self._capacity_sem is not None: self._capacity_sem.release() ctx._holds_slot = False - self._sessions.pop(ctx.session_id, None) + + async def _handle_close(self, ctx: SessionContext) -> HandlerResult: + """Close session: return env to pool or destroy it.""" + await self._release_env(ctx) + self._sessions.pop(ctx.session_id, None) n_active = sum(1 for s in self._sessions.values() if s.env is not None) n_queued = len(self._sessions) - n_active LOGGER.info( f"[Handler] Closed session {ctx.session_id} " - f"(active={n_active}, queued={n_queued}, capacity={self._capacity})" + f"(active={n_active}, queued={n_queued}, " + f"pool={len(self._env_pool)}, capacity={self._capacity})" ) return HandlerResult(data={"closed": True}) def get_session_stats(self) -> Dict[str, Any]: - """Session stats with active/queued breakdown.""" + """Session stats with active/queued/pool breakdown.""" stats = super().get_session_stats() n_active = sum(1 for s in self._sessions.values() if s.env is not None) stats["active"] = n_active stats["queued"] = len(self._sessions) - n_active stats["capacity"] = self._capacity if self._capacity > 0 else "unlimited" + stats["pool_size"] = len(self._env_pool) + stats["pool_max"] = self._pool_size for s in stats.get("sessions", []): sid = s["session_id"] ctx = self._sessions.get(sid) @@ -320,7 +377,7 @@ def get_session_stats(self) -> Dict[str, Any]: return stats async def _cleanup_loop(self): - """Cleanup timed-out sessions, releasing capacity slots.""" + """Cleanup timed-out sessions, releasing capacity slots or pooling envs.""" while True: try: await asyncio.sleep(60) @@ -335,23 +392,15 @@ async def _cleanup_loop(self): ctx = self._sessions.get(session_id) if ctx is None: continue - try: - if ctx.env is not None: - await ctx.env.close() - except Exception as e: - LOGGER.error(f"[Handler] Cleanup error {session_id}: {e}") - finally: - if ctx._holds_slot and self._capacity_sem is not None: - self._capacity_sem.release() - ctx._holds_slot = False - self._sessions.pop(session_id, None) + await self._release_env(ctx) + self._sessions.pop(session_id, None) except asyncio.CancelledError: break except Exception as e: LOGGER.error(f"[Handler] Cleanup loop error: {e}") async def aclose(self): - """Shutdown: close all sessions, release all capacity slots.""" + """Shutdown: close all sessions and pooled envs, release all capacity slots.""" if self._cleanup_task and not self._cleanup_task.done(): self._cleanup_task.cancel() try: @@ -375,4 +424,14 @@ async def _close_one(sid: str, ctx: SessionContext): *(_close_one(sid, ctx) for sid, ctx in self._sessions.items()) ) self._sessions.clear() - LOGGER.info("[Handler] All sessions closed") + + # Close all pooled envs (they don't hold capacity permits) + n_pooled = len(self._env_pool) + for env in self._env_pool: + try: + await env.close() + except Exception as e: + LOGGER.error(f"[Handler] Shutdown pool close error: {e}") + self._env_pool.clear() + + LOGGER.info(f"[Handler] All sessions closed, {n_pooled} pooled envs released") diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index e8b6fe741..73030abbb 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -149,6 +149,9 @@ def main( # Max Unity processes starting up simultaneously (0 = unlimited). # Prevents CPU spikes when many capacity slots open at once. startup_concurrency: int = 8, + # Max idle envs kept alive in pool for instant reuse. + # -1 = same as capacity (default), 0 = disable pooling. + pool_size: int = -1, # Thread pool for asyncio.to_thread(). Should be >= capacity. thread_pool_size: int = 128, # Session idle timeout before auto-cleanup (seconds). @@ -174,13 +177,14 @@ def main( LOGGER.info( f"GPUs: {devices} | displays: {x_displays} | " f"capacity: {capacity} | startup_concurrency: {startup_concurrency} | " - f"threads: {thread_pool_size}" + f"pool_size: {pool_size} | threads: {thread_pool_size}" ) handler = EbAlfredHandler( x_displays=x_displays, capacity=capacity, startup_concurrency=startup_concurrency, + pool_size=pool_size, session_timeout=session_timeout, max_sessions=max_sessions, ) From 192fa727abf03f9b11f191c0b90eeb4ebf0b53e0 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Tue, 24 Mar 2026 20:39:10 -0400 Subject: [PATCH 24/29] add preload logic --- tests/test_env_pool.py | 65 +++++++++++++++++++++++++++++++++ vagen/envs/eb_alfred/handler.py | 30 +++++++++++++++ vagen/envs/eb_alfred/serve.py | 7 ++++ 3 files changed, 102 insertions(+) diff --git a/tests/test_env_pool.py b/tests/test_env_pool.py index c5b8a74b2..f13337acb 100644 --- a/tests/test_env_pool.py +++ b/tests/test_env_pool.py @@ -259,11 +259,76 @@ async def run_episode(handler, seed): print("PASS: test_batch_cycle") +async def test_preload(): + """Verify preload fills the pool before any client connects.""" + FakeEnv._count = 0 + handler = EbAlfredHandler( + x_displays=["0"], + capacity=4, + startup_concurrency=4, + pool_size=4, + ) + + async def fake_create(config): + await asyncio.sleep(0.01) + env = FakeEnv() + env._assigned_display = "0" + return env + + handler.create_env = fake_create + + # Preload 4 envs + await handler.preload(4, {"eval_set": "base"}) + assert len(handler._env_pool) == 4, f"Expected 4 preloaded, got {len(handler._env_pool)}" + assert FakeEnv._count == 4 + + # Connect 4 sessions → all should reuse from pool instantly + handler._ensure_semaphore() + for i in range(4): + await handler.connect({"eval_set": "base"}, seed=i) + + for ctx in handler._sessions.values(): + if hasattr(ctx, '_ready') and ctx._ready: + await ctx._ready.wait() + + assert FakeEnv._count == 4, f"Should reuse preloaded, got {FakeEnv._count}" + assert len(handler._env_pool) == 0 + + await handler.aclose() + print("PASS: test_preload") + + +async def test_preload_capped_by_pool_size(): + """Preload(n) should be capped at pool_size.""" + FakeEnv._count = 0 + handler = EbAlfredHandler( + x_displays=["0"], + capacity=8, + startup_concurrency=8, + pool_size=3, + ) + + async def fake_create(config): + env = FakeEnv() + env._assigned_display = "0" + return env + + handler.create_env = fake_create + + await handler.preload(10, {"eval_set": "base"}) # Request 10, capped to 3 + assert len(handler._env_pool) == 3, f"Expected 3 (capped), got {len(handler._env_pool)}" + + await handler.aclose() + print("PASS: test_preload_capped_by_pool_size") + + async def main(): await test_pool_basic() await test_pool_overflow() await test_no_deadlock_with_queuing() await test_batch_cycle() + await test_preload() + await test_preload_capped_by_pool_size() print("\nAll tests passed!") diff --git a/vagen/envs/eb_alfred/handler.py b/vagen/envs/eb_alfred/handler.py index 0505060fc..cb24cc4aa 100644 --- a/vagen/envs/eb_alfred/handler.py +++ b/vagen/envs/eb_alfred/handler.py @@ -144,6 +144,36 @@ def _ensure_semaphore(self) -> None: if self._startup_sem is None and self._startup_concurrency > 0 and self._capacity > 0: self._startup_sem = asyncio.Semaphore(self._startup_concurrency) + async def preload(self, n: int, env_config: Dict[str, Any]) -> None: + """Pre-create *n* environments and place them in the idle pool. + + Called once during server startup so that the first training batch + gets instant env assignment instead of waiting ~90 s per Unity + process. Envs are created with ``startup_concurrency`` throttling. + + Args: + n: Number of envs to pre-create (capped at pool_size). + env_config: Config dict forwarded to ``create_env()``. + """ + n = min(n, self._pool_size) + if n <= 0: + return + + sem = asyncio.Semaphore(self._startup_concurrency or n) + + async def _create_one(idx: int): + async with sem: + LOGGER.info(f"[Preload] Creating env {idx + 1}/{n} ...") + env = await self.create_env(env_config) + return env + + t0 = time.time() + LOGGER.info(f"[Preload] Pre-creating {n} envs (concurrency={self._startup_concurrency or n}) ...") + envs = await asyncio.gather(*[_create_one(i) for i in range(n)]) + self._env_pool.extend(envs) + elapsed = time.time() - t0 + LOGGER.info(f"[Preload] {n} envs ready in {elapsed:.1f}s (pool: {len(self._env_pool)})") + def _least_loaded_display(self) -> str: """Pick the display with the fewest active + pending sessions. diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index 73030abbb..96a74d11e 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -152,6 +152,11 @@ def main( # Max idle envs kept alive in pool for instant reuse. # -1 = same as capacity (default), 0 = disable pooling. pool_size: int = -1, + # Pre-create this many envs on server startup (0 = lazy). + # Fills the pool so the first training batch doesn't wait ~90s per env. + preload: int = 0, + # eval_set used for preloaded envs (must match training config). + preload_eval_set: str = "base", # Thread pool for asyncio.to_thread(). Should be >= capacity. thread_pool_size: int = 128, # Session idle timeout before auto-cleanup (seconds). @@ -193,6 +198,8 @@ def main( @app.on_event("startup") async def _configure_executor(): asyncio.get_running_loop().set_default_executor(executor) + if preload > 0: + await handler.preload(preload, {"eval_set": preload_eval_set}) @app.on_event("shutdown") def _shutdown_executor(): From 39c31b11af6b09d5380012e80a786485a6bdee5e Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Tue, 24 Mar 2026 21:14:23 -0400 Subject: [PATCH 25/29] support dynamic eval_set per reset for env pooling --- vagen/envs/eb_alfred/eb_alfred_env.py | 49 ++++++++++++++++++++++----- vagen/envs/eb_alfred/handler.py | 8 ++++- vagen/envs/eb_alfred/serve.py | 5 ++- 3 files changed, 49 insertions(+), 13 deletions(-) diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py index 46c8e4922..06b09b015 100644 --- a/vagen/envs/eb_alfred/eb_alfred_env.py +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -9,8 +9,10 @@ """ import asyncio +import json import os import signal +import time import threading import numpy as np from PIL import Image @@ -130,6 +132,19 @@ def __init__(self, env_config: Dict[str, Any]): finally: _tl.x_display = None + # Replace single-split dataset with ALL splits so that a pooled + # env can serve any eval_set on reset without recreating Unity. + # Original EBAlfEnv loads only one eval_set into self.dataset (list). + # We reload the full splits.json into a dict {eval_set: [episodes]}. + with open(self.env.data_path) as f: + all_splits = json.load(f) + ds_ratio = self.config.down_sample_ratio + if 0 < ds_ratio < 1: + every = round(1 / ds_ratio) + all_splits = {k: v[::every] for k, v in all_splits.items()} + self.env.dataset = all_splits + self.env._default_eval_set = self.config.eval_set + # Adapter state (reset per episode) self._total_turns: int = 0 self._total_env_steps: int = 0 @@ -184,21 +199,37 @@ async def system_prompt(self) -> Dict[str, Any]: ) return {"obs_str": sys_str + "\n" + fmt_str} - async def reset(self, seed: int) -> Tuple[Dict[str, Any], Dict[str, Any]]: + def _reset_sync(self, eval_set: str, episode_idx: int): + """Synchronous reset that drives the underlying EBAlfEnv. + + We call ``_reset_controller`` directly (instead of ``env.reset()``) + because the upstream EBAlfEnv.reset() only supports sequential + iteration over a single eval_set. This wrapper manages episode + selection externally via (eval_set, episode_idx). + """ + task = self.env.dataset[eval_set][episode_idx] + self.env._reset_controller(task) + self.env._current_step = 0 + self.env._cur_invalid_actions = 0 + self.env._reset = True + self.env.episode_log = [] + self.env._episode_start_time = time.time() + + async def reset(self, seed: int, eval_set: str = None) -> Tuple[Dict[str, Any], Dict[str, Any]]: """ Reset environment for a new episode. The seed selects which episode to load from the dataset - (seed % number_of_episodes). After reset, the observation - includes the task instruction, available actions, and - the initial RGB image from AI2-THOR. + (seed % number_of_episodes_in_eval_set). The eval_set can be + overridden per-reset so a single pooled env can serve any split. """ - # Select episode based on seed - episode_idx = seed % self.env.number_of_episodes - self.env._current_episode_num = episode_idx + es = eval_set or self.config.eval_set + n_episodes = len(self.env.dataset.get(es, [])) + episode_idx = seed % n_episodes await asyncio.wait_for( - asyncio.to_thread(self.env.reset), timeout=300.0 + asyncio.to_thread(self._reset_sync, es, episode_idx), + timeout=300.0, ) # Reset adapter state @@ -218,7 +249,7 @@ async def reset(self, seed: int) -> Tuple[Dict[str, Any], Dict[str, Any]]: info = { "task_instruction": self.env.episode_language_instruction, "num_actions": len(self._action_list), - "eval_set": self.config.eval_set, + "eval_set": es, "episode_idx": episode_idx, } return obs, info diff --git a/vagen/envs/eb_alfred/handler.py b/vagen/envs/eb_alfred/handler.py index cb24cc4aa..4437734ac 100644 --- a/vagen/envs/eb_alfred/handler.py +++ b/vagen/envs/eb_alfred/handler.py @@ -287,9 +287,15 @@ async def _deferred_create(self, ctx: _DeferredSessionContext) -> None: # ---- fast path: reuse from pool ---- if self._env_pool: ctx.env = self._env_pool.pop() + # Update eval_set to match new session (env loads all splits, + # so only the config default needs updating) + new_eval_set = ctx.env_config.get("eval_set") + if new_eval_set and hasattr(ctx.env, "config"): + ctx.env.config.eval_set = new_eval_set LOGGER.info( f"[Handler] Session {ctx.session_id} reused pooled env " - f"(pool: {len(self._env_pool)} remaining)" + f"(pool: {len(self._env_pool)} remaining, " + f"eval_set={new_eval_set})" ) return diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index 96a74d11e..8d664fa35 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -154,9 +154,8 @@ def main( pool_size: int = -1, # Pre-create this many envs on server startup (0 = lazy). # Fills the pool so the first training batch doesn't wait ~90s per env. + # Each env loads all eval_sets, so no split-specific config is needed. preload: int = 0, - # eval_set used for preloaded envs (must match training config). - preload_eval_set: str = "base", # Thread pool for asyncio.to_thread(). Should be >= capacity. thread_pool_size: int = 128, # Session idle timeout before auto-cleanup (seconds). @@ -199,7 +198,7 @@ def main( async def _configure_executor(): asyncio.get_running_loop().set_default_executor(executor) if preload > 0: - await handler.preload(preload, {"eval_set": preload_eval_set}) + await handler.preload(preload, {}) @app.on_event("shutdown") def _shutdown_executor(): From 39f8373010b7324dbb4cff50384e5de526945651 Mon Sep 17 00:00:00 2001 From: march <20082585d@gmail.com> Date: Tue, 24 Mar 2026 21:59:55 -0400 Subject: [PATCH 26/29] fix preload via lifespan, support dynamic eval_set per reset --- vagen/envs/eb_alfred/serve.py | 9 ++------- vagen/envs_remote/service.py | 8 +++++++- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index 8d664fa35..8364f2db4 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -192,17 +192,12 @@ def main( session_timeout=session_timeout, max_sessions=max_sessions, ) - app = GymService(handler, api_key=api_key).build() - - @app.on_event("startup") - async def _configure_executor(): + async def _on_startup(): asyncio.get_running_loop().set_default_executor(executor) if preload > 0: await handler.preload(preload, {}) - @app.on_event("shutdown") - def _shutdown_executor(): - executor.shutdown(wait=True) + app = GymService(handler, api_key=api_key).build(on_startup=_on_startup) uvicorn.run(app, host=host, port=port, workers=workers) diff --git a/vagen/envs_remote/service.py b/vagen/envs_remote/service.py index 6a2ce4738..fe385d22d 100644 --- a/vagen/envs_remote/service.py +++ b/vagen/envs_remote/service.py @@ -241,18 +241,24 @@ def register_routes(self, app: FastAPI) -> None: app.add_api_route("/connect", self.connect, methods=["POST"]) app.add_api_route("/call", self.call, methods=["POST"]) - def build(self) -> FastAPI: + def build(self, on_startup=None) -> FastAPI: """ Build and return the FastAPI application. This is the main entry point. Call once, then run the returned app with uvicorn. + + Args: + on_startup: Optional async callable invoked during lifespan startup + (before the app starts serving requests). """ handler = self.handler @asynccontextmanager async def lifespan(app: FastAPI): try: + if on_startup is not None: + await on_startup() yield finally: await handler.aclose() From 545e1d5484dbb94bd2a315f2aa63f972231eaebd Mon Sep 17 00:00:00 2001 From: Yaning Gao Date: Fri, 27 Mar 2026 03:04:53 +0000 Subject: [PATCH 27/29] delta temp --- .../eb_alfred/train_eb_alfred_vision.yaml | 18 ++++---- .../train_ppo_no_concat_qwen25vl3b.sh | 35 ++++++++-------- .../train/eb_alfred/val_eb_alfred_vision.yaml | 40 +++++++++--------- vagen/envs/eb_alfred/eb_alfred_env.py | 5 ++- vagen/envs/eb_alfred/handler.py | 19 ++++++++- vagen/envs/eb_alfred/serve.py | 5 +++ vagen/envs/eb_alfred/utils/prompt.py | 42 +++++++++---------- vagen/envs/eb_alfred/utils/utils.py | 19 +++++++-- 8 files changed, 110 insertions(+), 73 deletions(-) diff --git a/examples/train/eb_alfred/train_eb_alfred_vision.yaml b/examples/train/eb_alfred/train_eb_alfred_vision.yaml index 2af2d3822..7bac39fb0 100644 --- a/examples/train/eb_alfred/train_eb_alfred_vision.yaml +++ b/examples/train/eb_alfred/train_eb_alfred_vision.yaml @@ -4,15 +4,15 @@ envs: data_source: eb_alfred tag_id: eb_alfred_train_base seed: [0, 200, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: base obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 action_sep: "|" @@ -26,15 +26,15 @@ envs: data_source: eb_alfred tag_id: eb_alfred_train_complex seed: [0, 200, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: complex_instruction obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 action_sep: "|" @@ -48,15 +48,15 @@ envs: data_source: eb_alfred tag_id: eb_alfred_train_visual seed: [0, 200, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: visual_appearance obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 action_sep: "|" diff --git a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh index 6c1b9ded0..fc903851d 100755 --- a/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh +++ b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh @@ -12,11 +12,12 @@ SAVE_CHECKPOINT_DIR=${EXPERIMENT_DIR}/verl_checkpoints DATASET_TRAIN=${SCRIPTDIR}/train_eb_alfred_vision.yaml DATASET_VAL=${SCRIPTDIR}/val_eb_alfred_vision.yaml agent_loop_config_path=${BASEDIR}/vagen/configs/agent_no_concat.yaml -REF_MODEL_PATH=/scratch/bgig/jma6/Qwen2.5-VL-3B-Instruct-EB-ALFRED-vagen-sft +REF_MODEL_PATH=/workspace/.hf_home/hub/models--err00rr--Qwen2.5-VL-3B-Instruct-EB-ALFRED-vagen-sft/snapshots/3ab58797a74bcc4198e140166d935533a82d28f7 mkdir -p ${EXPERIMENT_DIR} export HF_HOME=/workspace/.hf_home export PATH=/venv/vagen/bin:$PATH +export LD_LIBRARY_PATH=/venv/vagen/lib/python3.12/site-packages/nvidia/cuda_runtime/lib:${LD_LIBRARY_PATH} PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ --config-path=${BASEDIR}/vagen/configs \ @@ -24,8 +25,8 @@ PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ data.train_files=${DATASET_TRAIN} \ data.val_files=${DATASET_VAL} \ data.train_batch_size=128 \ - data.max_prompt_length=5000 \ - data.max_response_length=512 \ + data.max_prompt_length=9000 \ + data.max_response_length=2048 \ +data.max_trajectory_length=100000 \ algorithm.adv_estimator=no_concat_gae \ algorithm.kl_ctrl.kl_coef=0.0 \ @@ -35,46 +36,46 @@ PYTHONUNBUFFERED=1 python3 -m vagen.main_ppo \ actor_rollout_ref.model.enable_gradient_checkpointing=True \ actor_rollout_ref.actor.optim.lr=1e-6 \ actor_rollout_ref.actor.ppo_mini_batch_size=32 \ - actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \ actor_rollout_ref.actor.use_kl_loss=False \ actor_rollout_ref.actor.kl_loss_coef=0.0 \ actor_rollout_ref.actor.kl_loss_type=low_var_kl \ actor_rollout_ref.actor.entropy_coeff=0.0 \ actor_rollout_ref.actor.checkpoint.save_contents=['model','hf_model','optimizer','extra'] \ actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \ - actor_rollout_ref.actor.fsdp_config.param_offload=True \ - actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ actor_rollout_ref.actor.freeze_vision_tower=True \ - actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ actor_rollout_ref.rollout.name=sglang \ actor_rollout_ref.rollout.mode=async \ actor_rollout_ref.rollout.n=1 \ - actor_rollout_ref.rollout.max_num_batched_tokens=10000 \ - actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \ + actor_rollout_ref.rollout.max_num_batched_tokens=65536 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ actor_rollout_ref.rollout.enforce_eager=True \ actor_rollout_ref.rollout.free_cache_engine=True \ actor_rollout_ref.rollout.enable_chunked_prefill=True \ actor_rollout_ref.rollout.multi_turn.enable=True \ actor_rollout_ref.rollout.agent.agent_loop_config_path=$agent_loop_config_path \ actor_rollout_ref.rollout.disable_log_stats=False \ - actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ - actor_rollout_ref.ref.fsdp_config.param_offload=True \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.ref.fsdp_config.param_offload=False \ critic.enable=True \ critic.optim.lr=1e-5 \ critic.model.use_remove_padding=True \ critic.model.path=${REF_MODEL_PATH} \ critic.model.enable_gradient_checkpointing=True \ - critic.ppo_micro_batch_size_per_gpu=1 \ - critic.model.fsdp_config.param_offload=True \ - critic.model.fsdp_config.optimizer_offload=True \ + critic.ppo_micro_batch_size_per_gpu=2 \ + critic.model.fsdp_config.param_offload=False \ + critic.model.fsdp_config.optimizer_offload=False \ trainer.critic_warmup=0 \ trainer.logger=['console','wandb'] \ - trainer.val_before_train=True \ - trainer.n_gpus_per_node=8 \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=4 \ trainer.nnodes=1 \ trainer.save_freq=100 \ - trainer.test_freq=20 \ + trainer.test_freq=50 \ trainer.project_name=${PROJECT_NAME} \ trainer.experiment_name=${EXPERIMENT_NAME} \ trainer.default_local_dir=${SAVE_CHECKPOINT_DIR} \ diff --git a/examples/train/eb_alfred/val_eb_alfred_vision.yaml b/examples/train/eb_alfred/val_eb_alfred_vision.yaml index 17deb8e80..e89ab150b 100644 --- a/examples/train/eb_alfred/val_eb_alfred_vision.yaml +++ b/examples/train/eb_alfred/val_eb_alfred_vision.yaml @@ -4,18 +4,18 @@ envs: data_source: eb_alfred tag_id: eb_alfred_val_common seed: [0, 50, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: common_sense obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -26,18 +26,18 @@ envs: data_source: eb_alfred tag_id: eb_alfred_val_spatial seed: [0, 50, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: spatial obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -48,18 +48,18 @@ envs: data_source: eb_alfred tag_id: eb_alfred_train_base seed: [0, 200, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: base obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -70,18 +70,18 @@ envs: data_source: eb_alfred tag_id: eb_alfred_train_complex seed: [0, 200, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: complex_instruction obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 @@ -92,18 +92,18 @@ envs: data_source: eb_alfred tag_id: eb_alfred_train_visual seed: [0, 200, 1] - max_turns: 6 - response_length_per_turn: 512 + max_turns: 20 + response_length_per_turn: 2048 config: base_urls: - "http://localhost:8000" timeout: 600 eval_set: visual_appearance obs_image_size: 500 - max_turns: 6 + max_turns: 20 max_actions_per_step: 20 max_env_steps: 30 - action_sep: "," + action_sep: "|" prompt_format: free_think use_example_in_sys_prompt: true format_reward: 0.1 diff --git a/vagen/envs/eb_alfred/eb_alfred_env.py b/vagen/envs/eb_alfred/eb_alfred_env.py index 06b09b015..c0c19cd2b 100644 --- a/vagen/envs/eb_alfred/eb_alfred_env.py +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -78,7 +78,7 @@ class EbAlfredEnvConfig: max_turns: int = 30 max_actions_per_step: int = 20 max_env_steps: int = 30 # Max total environment actions per episode (matches ERA) - action_sep: str = "," + action_sep: str = "|" image_placeholder: str = "" prompt_format: str = "free_think" use_example_in_sys_prompt: bool = True @@ -190,6 +190,9 @@ async def system_prompt(self) -> Dict[str, Any]: sys_str = system_prompt( task_instruction=self.env.episode_language_instruction, action_list=self._action_list, + max_actions_per_step=self.config.max_actions_per_step, + action_sep=self.config.action_sep, + max_turns=self.config.max_turns, ) fmt_str = format_prompt( max_actions_per_step=self.config.max_actions_per_step, diff --git a/vagen/envs/eb_alfred/handler.py b/vagen/envs/eb_alfred/handler.py index 4437734ac..ae53ec4e6 100644 --- a/vagen/envs/eb_alfred/handler.py +++ b/vagen/envs/eb_alfred/handler.py @@ -103,6 +103,7 @@ def __init__( capacity: int = 16, startup_concurrency: int = 8, pool_size: int = -1, + reset_concurrency: int = 8, **kwargs, ): """ @@ -115,6 +116,9 @@ def __init__( slots open simultaneously. Ignored when capacity = 0. pool_size: Max idle envs kept alive in the pool. -1 (default) = same as capacity (keep every env alive). 0 = disable pooling. + reset_concurrency: Max Unity scene resets happening simultaneously + (0 = unlimited). Prevents I/O and X11 saturation when many + sessions reset at the same time (e.g. start of each rollout). **kwargs: Passed to BaseGymHandler (session_timeout, max_sessions). """ super().__init__(**kwargs) @@ -122,6 +126,7 @@ def __init__( self._pending_counts: Dict[str, int] = {d: 0 for d in self._x_displays} self._capacity = capacity self._startup_concurrency = startup_concurrency + self._reset_concurrency = reset_concurrency # Env pool: idle Unity processes available for immediate reuse. # Each pooled env implicitly holds one capacity-semaphore permit. self._pool_size = pool_size if pool_size >= 0 else max(capacity, 1) @@ -130,11 +135,13 @@ def __init__( # not during __init__ (which runs before uvicorn starts the loop). self._capacity_sem: Optional[asyncio.Semaphore] = None self._startup_sem: Optional[asyncio.Semaphore] = None + self._reset_sem: Optional[asyncio.Semaphore] = None LOGGER.info( f"[Handler] Using X displays: {self._x_displays}, " f"capacity={capacity if capacity > 0 else 'unlimited'}, " f"startup_concurrency={startup_concurrency if startup_concurrency > 0 else 'unlimited'}, " - f"pool_size={self._pool_size}" + f"pool_size={self._pool_size}, " + f"reset_concurrency={reset_concurrency if reset_concurrency > 0 else 'unlimited'}" ) def _ensure_semaphore(self) -> None: @@ -143,6 +150,8 @@ def _ensure_semaphore(self) -> None: self._capacity_sem = asyncio.Semaphore(self._capacity) if self._startup_sem is None and self._startup_concurrency > 0 and self._capacity > 0: self._startup_sem = asyncio.Semaphore(self._startup_concurrency) + if self._reset_sem is None and self._reset_concurrency > 0: + self._reset_sem = asyncio.Semaphore(self._reset_concurrency) async def preload(self, n: int, env_config: Dict[str, Any]) -> None: """Pre-create *n* environments and place them in the idle pool. @@ -383,6 +392,14 @@ async def _release_env(self, ctx: SessionContext) -> None: self._capacity_sem.release() ctx._holds_slot = False + async def _handle_reset(self, ctx: SessionContext, params: Dict[str, Any]) -> HandlerResult: + """Handle reset with concurrency throttling to avoid I/O saturation.""" + self._ensure_semaphore() + if self._reset_sem is not None: + async with self._reset_sem: + return await super()._handle_reset(ctx, params) + return await super()._handle_reset(ctx, params) + async def _handle_close(self, ctx: SessionContext) -> HandlerResult: """Close session: return env to pool or destroy it.""" await self._release_env(ctx) diff --git a/vagen/envs/eb_alfred/serve.py b/vagen/envs/eb_alfred/serve.py index 8364f2db4..ccdf2b66e 100644 --- a/vagen/envs/eb_alfred/serve.py +++ b/vagen/envs/eb_alfred/serve.py @@ -152,6 +152,10 @@ def main( # Max idle envs kept alive in pool for instant reuse. # -1 = same as capacity (default), 0 = disable pooling. pool_size: int = -1, + # Max concurrent Unity scene resets (0 = unlimited). + # Prevents I/O and X11 saturation when all sessions reset simultaneously + # at the start of each rollout/eval batch. + reset_concurrency: int = 8, # Pre-create this many envs on server startup (0 = lazy). # Fills the pool so the first training batch doesn't wait ~90s per env. # Each env loads all eval_sets, so no split-specific config is needed. @@ -189,6 +193,7 @@ def main( capacity=capacity, startup_concurrency=startup_concurrency, pool_size=pool_size, + reset_concurrency=reset_concurrency, session_timeout=session_timeout, max_sessions=max_sessions, ) diff --git a/vagen/envs/eb_alfred/utils/prompt.py b/vagen/envs/eb_alfred/utils/prompt.py index 6fed279b7..565510215 100644 --- a/vagen/envs/eb_alfred/utils/prompt.py +++ b/vagen/envs/eb_alfred/utils/prompt.py @@ -26,7 +26,7 @@ ## The available action id (0 ~ {max_action_id}) and action names are: {available_actions}. ## Guidelines -1. **Output Plan**: Avoid generating empty plan. Each plan should include no more than 20 actions. +1. **Output Plan**: Avoid generating empty plan. Each plan should include no more than {max_actions_per_step} actions. 2. **Visibility**: Always locate a visible object by the 'find' action before interacting with it. 3. **Action Guidelines**: Make sure match the action name and its corresponding action id in the output. Avoid performing actions that do not meet the defined validity criteria. For instance, if you want to put object in a receptacle, use 'put down' rather than 'drop' actions. 4. **Prevent Repeating Action Sequences**: Do not repeatedly execute the same action or sequence of actions. Try to modify the action sequence because previous actions do not lead to success. @@ -34,8 +34,9 @@ 6. **Reflection on History and Feedback**: Use interaction history and feedback from the environment to refine and improve your current plan. If the last action is invalid, reflect on the reason, such as not adhering to action rules or missing preliminary actions, and adjust your plan accordingly. ** Generation Guide ** - - Include the thinking process between <|think_start|> and <|think_end|> - - Include only the target action in <|action_start|> and <|action_end|>, i.e. the content inside <|action_start|> and <|action_end|> should be nothing more than [action_id, 'action_name'], where the action id is an integer and the action name is the corresponding name. Do not include any other thing, such as '"'. + - You have at most {max_turns} turns to complete the task. + - Include the thinking process between <|think_start|> and <|think_end|>. + - Include up to {max_actions_per_step} action(s) in <|action_start|> and <|action_end|>, separated by '{action_sep}'. Each action must be [action_id, 'action_name'], where action_id is an integer and action_name is the corresponding name from the available actions list. """ @@ -43,14 +44,11 @@ def system_prompt( task_instruction: Optional[str] = None, action_list: Optional[List[str]] = None, add_task_examples: bool = True, + max_actions_per_step: int = 20, + action_sep: str = "|", + max_turns: int = 30, ): - """ - Build ERA-aligned system prompt for EB-ALFRED. - - The prompt exactly matches what the ERA EPL-Only model was trained on, - so the model stays in-distribution. Output normalisation (ERA tokens → - VAGEN tags) happens downstream in normalize_era_tokens(). - """ + """Build system prompt for EB-ALFRED.""" if action_list is not None: max_id = len(action_list) - 1 available = ", ".join( @@ -63,6 +61,9 @@ def system_prompt( return ERA_SYSTEM_PROMPT_TEMPLATE.format( max_action_id=max_id, available_actions=available, + max_actions_per_step=max_actions_per_step, + action_sep=action_sep, + max_turns=max_turns, ) @@ -119,10 +120,15 @@ def format_prompt(max_actions_per_step, action_sep, add_example=True, prompt_for def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True): - """Minimal format prompt — the ERA system prompt already has the Generation Guide.""" - # The ERA system prompt already instructs the model on output format. - # Adding extra format instructions can confuse an SFT model, so keep it short. - return "" + """Format prompt for free_think mode with concrete output examples.""" + if not add_example: + return "" + + return f""" +Example output (multiple actions, separated by '{action_sep}'): +I can see the mug nearby and I am not holding anything. I will pick it up and then put it on the table. +[12, 'pick up the Mug']{action_sep}[38, 'put down the object in hand'] +""" def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): @@ -153,13 +159,7 @@ def wm_format_prompt(max_actions_per_step, action_sep, add_example=True): I am close to a Mug on the counter. I am not holding anything. The mug is within reach. The mug is nearby and I'm not holding anything. I should pick it up. [12, 'pick up the Mug'] -I will be holding the mug. The counter will no longer have the mug on it. - -Example 3: -I am holding a Mug. I see a table nearby with an empty spot. -I'm holding the mug and I'm near the table. Let me put it down. -[38, 'put down the object in hand'] -The mug will be placed on the table. I will no longer be holding anything.""" +I will be holding the mug. The counter will no longer have the mug on it.""" return base + "\n" + examples return base diff --git a/vagen/envs/eb_alfred/utils/utils.py b/vagen/envs/eb_alfred/utils/utils.py index 8a030b849..6bd8b0c1d 100644 --- a/vagen/envs/eb_alfred/utils/utils.py +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -27,10 +27,21 @@ def parse_free_think(response: str, action_sep: str = ",", max_actions: int = 1) if max_actions == 1: actions = [action_content.strip()] if action_content.strip() else [] else: - actions = [a.strip() for a in action_content.split(action_sep) if a.strip()] - if len(actions) > max_actions: - actions = actions[:max_actions] - action_content = action_sep.join(actions) + # First try splitting on action_sep + candidates = [a.strip() for a in action_content.split(action_sep) if a.strip()] + # If splitting produced broken bracket fragments (ERA model uses "," inside + # [id, 'action'] notation AND as the multi-action separator), fall back to + # regex extraction of all [id, 'action'] tokens. + has_broken_brackets = any( + re.match(r'^\[?\d+$', c) or re.match(r"^['\"].+$", c) + for c in candidates + ) + if has_broken_brackets or (len(candidates) <= 1 and re.search(r'\]\s*,\s*\[', action_content)): + candidates = re.findall(r'\[\d+,\s*[\'"]?[^\[\]]+?[\'"]?\s*\]', action_content) + if len(candidates) > max_actions: + candidates = candidates[:max_actions] + actions = candidates + action_content = action_sep.join(actions) llm_response = f"{think_content}{action_content}" From 793b952b2a5e81abc5403f39dc1dc2474d250187 Mon Sep 17 00:00:00 2001 From: Jingnan Ma Date: Tue, 7 Apr 2026 04:40:48 -0500 Subject: [PATCH 28/29] tmp: free SGLang memory before critic/actor update on <=2 GPU configs When n_gpus_per_node <= 2, call sleep()/wake_up() on the actor_rollout worker group before and after critic/actor updates to release SGLang's KV cache and weight memory. This prevents OOM during update_critic on memory-constrained setups (e.g. 2x H200 with large prompts). 4+ GPU configs are unaffected. Co-Authored-By: Claude Opus 4.6 (1M context) --- vagen/ray_trainer.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vagen/ray_trainer.py b/vagen/ray_trainer.py index 10ffc2d82..193c477fc 100644 --- a/vagen/ray_trainer.py +++ b/vagen/ray_trainer.py @@ -1534,6 +1534,15 @@ def fit(self): print(f"After filtering: Pad {pad_size} samples to make batch size {batch_size} divisible by {divisor_size} dp_workers") self._balance_batch(batch, metrics=metrics, logging_prefix="filtered_global_seqlen") + # Free SGLang inference memory before training updates (only for small GPU configs) + _free_inference_mem = ( + self.config.trainer.n_gpus_per_node <= 2 + and hasattr(self, 'actor_rollout_wg') + and hasattr(self.actor_rollout_wg, 'sleep') + ) + if _free_inference_mem: + self.actor_rollout_wg.sleep() + # update critic if self.use_critic: with marked_timer("update_critic", timing_raw, color="pink"): @@ -1550,6 +1559,10 @@ def fit(self): actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) metrics.update(actor_output_metrics) + # Resume SGLang inference memory after training updates + if _free_inference_mem: + self.actor_rollout_wg.wake_up() + # Log rollout generations if enabled rollout_data_dir = self.config.trainer.get("rollout_data_dir", None) if rollout_data_dir: From a101871f2594dce0e1c999655271ac53320978fe Mon Sep 17 00:00:00 2001 From: Jingnan Ma Date: Thu, 9 Apr 2026 11:49:28 -0500 Subject: [PATCH 29/29] fix: graceful fallback when env reset fails in no_concat agent loop When env.reset() throws an exception (e.g. server 500, timeout), return a dummy AgentLoopOutput with reward=0 instead of crashing the entire training run. Includes a dummy image to keep batch consistency with image_data fields. Co-Authored-By: Claude Opus 4.6 (1M context) --- vagen/agent_loop/gym_agent_loop_no_concat.py | 27 ++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/vagen/agent_loop/gym_agent_loop_no_concat.py b/vagen/agent_loop/gym_agent_loop_no_concat.py index afd65005a..5793f90f7 100644 --- a/vagen/agent_loop/gym_agent_loop_no_concat.py +++ b/vagen/agent_loop/gym_agent_loop_no_concat.py @@ -114,8 +114,31 @@ async def run(self, sampling_params: Dict[str, Any], **kwargs) -> AgentLoopOutpu env: GymImageEnv = env_cls(env_config=env_config) # Bootstrap: reset -> system_prompt (message order: system, then initial user) - init_obs, info = await env.reset(seed=seed) - sys_obs = await env.system_prompt() + try: + init_obs, info = await env.reset(seed=seed) + sys_obs = await env.system_prompt() + except Exception as exc: + logger.error("Env reset failed in '%s' seed=%s: %s", env_name, seed, exc) + # Return a minimal failed output so training can continue + dummy_ids = [self.tokenizer.eos_token_id or 0] + try: + await env.close() + except Exception: + pass + from PIL import Image as _PILImage + dummy_image = _PILImage.new('RGB', (256, 256), color='white') + return [AgentLoopOutput( + prompt_ids=dummy_ids, + response_ids=dummy_ids, + response_mask=[1], + multi_modal_data={"image": [dummy_image]}, + response_logprobs=None, + reward_score=0.0, + num_turns=0, + metrics=metrics, + extra_fields={"reward_extra_info": {"traj_success": 0.0}, + "last_turn": True}, + )]