diff --git a/examples/evaluate/eb_alfred/config.yaml b/examples/evaluate/eb_alfred/config.yaml new file mode 100644 index 000000000..5be4b0ad8 --- /dev/null +++ b/examples/evaluate/eb_alfred/config.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: "sglang" + base_seed: 0 + max_concurrent_jobs: 1 + resume: skip_completed + live_summary: true + +backends: + sglang: + 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 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 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/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_eb_alfred_vision.yaml b/examples/train/eb_alfred/train_eb_alfred_vision.yaml new file mode 100644 index 000000000..7bac39fb0 --- /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: 20 + response_length_per_turn: 2048 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: base + obs_image_size: 500 + max_turns: 20 + 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: 20 + response_length_per_turn: 2048 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: complex_instruction + obs_image_size: 500 + max_turns: 20 + 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: 20 + response_length_per_turn: 2048 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: visual_appearance + obs_image_size: 500 + max_turns: 20 + 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..fc903851d --- /dev/null +++ b/examples/train/eb_alfred/train_ppo_no_concat_qwen25vl3b.sh @@ -0,0 +1,87 @@ +#!/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=/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 \ + --config-name='vagen_multiturn' \ + data.train_files=${DATASET_TRAIN} \ + data.val_files=${DATASET_VAL} \ + data.train_batch_size=128 \ + 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 \ + 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=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=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=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=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=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=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=False \ + trainer.n_gpus_per_node=4 \ + trainer.nnodes=1 \ + trainer.save_freq=100 \ + trainer.test_freq=50 \ + 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..e89ab150b --- /dev/null +++ b/examples/train/eb_alfred/val_eb_alfred_vision.yaml @@ -0,0 +1,110 @@ +envs: + - name: RemoteEnv + n_envs: 10 + data_source: eb_alfred + tag_id: eb_alfred_val_common + seed: [0, 50, 1] + 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: 20 + 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_val_spatial + seed: [0, 50, 1] + 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: 20 + 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_base + seed: [0, 200, 1] + 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: 20 + 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: 20 + response_length_per_turn: 2048 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: complex_instruction + obs_image_size: 500 + max_turns: 20 + 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: 20 + response_length_per_turn: 2048 + config: + base_urls: + - "http://localhost:8000" + timeout: 600 + eval_set: visual_appearance + obs_image_size: 500 + max_turns: 20 + 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/tests/test_env_pool.py b/tests/test_env_pool.py new file mode 100644 index 000000000..f13337acb --- /dev/null +++ b/tests/test_env_pool.py @@ -0,0 +1,336 @@ +""" +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 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!") + + +if __name__ == "__main__": + asyncio.run(main()) 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}, + )] diff --git a/vagen/envs/eb_alfred/README.md b/vagen/envs/eb_alfred/README.md new file mode 100644 index 000000000..e5c8af0be --- /dev/null +++ b/vagen/envs/eb_alfred/README.md @@ -0,0 +1,108 @@ +# EB-ALFRED Environment + +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 + +pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126 +pip install "ai2thor==2.1.0" "gym==0.23.0" "numpy<2.0" \ + scipy Pillow networkx revtok vocab h5py tqdm natsort pyquaternion +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): + +```bash +git clone https://huggingface.co/datasets/EmbodiedBench/EB-ALFRED +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): + +```bash +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) +- `startup_concurrency`: max Unity processes starting simultaneously, prevents CPU spikes (default: 8) +- `session_timeout`: idle session cleanup in seconds (default: 3600) + +```bash +# Example: 2 GPUs, higher capacity +python -m vagen.envs.eb_alfred.serve --devices='[0,1]' --capacity=90 --startup_concurrency=6 --port=8000 +``` + +**SSH tunnel** (if training machine is remote): + +```bash +# Run on the env server — forwards port 8000 to training machine +ssh -p -R 8000:localhost:8000 \ + -o ServerAliveInterval=30 -o ServerAliveCountMax=5 \ + -N -f user@training-machine-ip + +# 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/` + 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..c0c19cd2b --- /dev/null +++ b/vagen/envs/eb_alfred/eb_alfred_env.py @@ -0,0 +1,484 @@ +""" +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 json +import os +import signal +import time +import threading +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 + +# 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: + """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) + + # 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 + + _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 + + # 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 + 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 + + # ------------------------------------------------------------------ + # GymImageEnv abstract methods + # ------------------------------------------------------------------ + + async def close(self) -> None: + """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]: + """ + 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, + 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, + 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} + + 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_in_eval_set). The eval_set can be + overridden per-reset so a single pooled env can serve any split. + """ + 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._reset_sync, es, episode_idx), + timeout=300.0, + ) + + # Reset adapter state + self._total_turns = 0 + 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) + 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": es, + "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(parsed) + + actions = parsed.get("actions", []) + format_correct = parsed.get("format_correct", False) + self._last_thinking = parsed.get("think_content", "") + + 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.wait_for( + asyncio.to_thread(self.env.step, matched), timeout=60.0 + ) + ) + + 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) + 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, + task_instruction=self.env.episode_language_instruction, + ) + else: + obs_str = action_template( + last_action=self._last_action, + 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 { + "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..ae53ec4e6 --- /dev/null +++ b/vagen/envs/eb_alfred/handler.py @@ -0,0 +1,490 @@ +""" +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. + +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, 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. + + 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 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) + - 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 + + 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, + pool_size: int = -1, + reset_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. + 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) + self._x_displays = x_displays if x_displays is not None else detect_gpu_displays() + 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) + 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 + 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"reset_concurrency={reset_concurrency if reset_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) + 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. + + 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. + + 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 + 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 + + 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() + self._pending_counts[display] = self._pending_counts.get(display, 0) + 1 + env_config = {**env_config, "x_display": display} + + 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 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() + # 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"eval_set={new_eval_set})" + ) + 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: + 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 _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 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 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 + + 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) + 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}, " + 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/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) + 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 or pooling envs.""" + 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 + 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 and pooled envs, 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() + + # 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 new file mode 100644 index 000000000..ccdf2b66e --- /dev/null +++ b/vagen/envs/eb_alfred/serve.py @@ -0,0 +1,211 @@ +""" +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. + +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 all GPUs, start Xorg automatically: + python -m vagen.envs.eb_alfred.serve + + # 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"] + # eval_set: "base" + # resolution: 500 +""" + +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 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 _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") + or os.path.exists(f"/tmp/.X11-unix/X{display}")) + + +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, + ) + + 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, + # 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. + preload: int = 0, + # 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"pool_size: {pool_size} | threads: {thread_pool_size}" + ) + + handler = EbAlfredHandler( + x_displays=x_displays, + capacity=capacity, + startup_concurrency=startup_concurrency, + pool_size=pool_size, + reset_concurrency=reset_concurrency, + session_timeout=session_timeout, + max_sessions=max_sessions, + ) + async def _on_startup(): + asyncio.get_running_loop().set_default_executor(executor) + if preload > 0: + await handler.preload(preload, {}) + + app = GymService(handler, api_key=api_key).build(on_startup=_on_startup) + + uvicorn.run(app, host=host, port=port, workers=workers) + + +if __name__ == "__main__": + fire.Fire(main) diff --git a/vagen/envs/eb_alfred/start_server.sh b/vagen/envs/eb_alfred/start_server.sh new file mode 100755 index 000000000..0a8be6431 --- /dev/null +++ b/vagen/envs/eb_alfred/start_server.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# 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 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..565510215 --- /dev/null +++ b/vagen/envs/eb_alfred/utils/prompt.py @@ -0,0 +1,165 @@ +from typing import List, Optional + + +# ────────────────────────────────────────────────────────────────────── +# 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. +# ────────────────────────────────────────────────────────────────────── + +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. 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 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. +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 ** + - 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. + """ + + +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 system prompt for EB-ALFRED.""" + 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)" + + 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, + ) + + +def init_observation_template(img_str, task_instruction=None): + """ERA-style initial user message. + + ERA format: \\n instruction: {task} \\n interaction_history: [] ... + """ + 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. + + 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": + 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): + """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): + """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. +Your response must be in the format of: +......[N, 'action name'].... + +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 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'] +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'] +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 new file mode 100644 index 000000000..6bd8b0c1d --- /dev/null +++ b/vagen/envs/eb_alfred/utils/utils.py @@ -0,0 +1,235 @@ +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: + # 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}" + + 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 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. + + 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'] + """ + # 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 + + +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.""" + response = normalize_era_tokens(response) + 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 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" + + Returns the original action string if matched, None otherwise. + """ + name = action_name.strip() + + # 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): + 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) + 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).") 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() 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 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: