fix: replace EnvironmentManagerBase with VAGEN registry-based env integration

The previous implementation incorrectly assumed verl-agent uses an
EnvironmentManagerBase ABC with a hardcoded make_envs() dispatch.
Research reveals VAGEN actually uses:
- GymImageEnv protocol (which WAADesktopEnv already implements)
- YAML-based env registry (vagen/configs/env_registry.yaml)
- GymAgentLoop for training-time rollout orchestration

Changes:
- Replace patch_env_manager() with register_waa_env() (YAML registry)
- Add register_in_vagen() and generate_env_spec() helpers to verl_env.py
- Update launch_training() to generate proper VAGEN training config
- Fix Integration Gap section in decision doc (no EnvironmentManagerBase)
- Update training config YAML with architecture diagram
- Add 5 new tests for registration helpers (40 total, all passing)
- Export new helpers from adapters/__init__.py

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: abrichr

configs/train_waa_vagen.yaml

@@ -1,19 +1,35 @@
 # VAGEN training config for WAA desktop automation
 #
 # This trains a VLM (e.g., Qwen2.5-VL-3B) to automate Windows desktop tasks
-# using GRPO/GiGPO via the verl-agent framework.
+# using GRPO/GiGPO via the VAGEN framework (verl-agent).
 #
 # Prerequisites:
 #   1. WAA server running (via SSH tunnel): ssh -L 5001:localhost:5050 azureuser@<VM_IP>
-#   2. VAGEN installed: pip install vagen
-#   3. Register env: add to vagen's env_registry.yaml:
+#   2. VAGEN installed on GPU VM (see scripts/setup_gpu_training.sh)
+#   3. openadapt-evals installed on GPU VM (pip install openadapt-evals)
+#   4. Register env in VAGEN's env_registry.yaml:
 #        WAADesktop: openadapt_evals.adapters.verl_env.WAADesktopEnv
+#      (automated by: scripts/train_verl_e2e.py or oa-vm gpu-train)
+#
+# Architecture:
+#   GPU VM                              CPU VM
+#   ┌──────────────────────┐            ┌──────────────────┐
+#   │ VAGEN / verl         │            │ Docker           │
+#   │  GymAgentLoop        │  HTTP      │  QEMU (Win 11)   │
+#   │   WAADesktopEnv ─────│───────────>│   WAA Flask API  │
+#   │  GiGPO/GRPO trainer  │            │                  │
+#   │  vLLM inference      │            │                  │
+#   └──────────────────────┘            └──────────────────┘
 #
 # Usage:
-#   python -m vagen.train --config configs/train_waa_vagen.yaml
+#   # Via orchestration script (recommended):
+#   python scripts/train_verl_e2e.py --cloud aws --task-id <UUID>
+#
+#   # Via CLI:
+#   oa-vm gpu-train --cloud aws --task-id <UUID>
 #
 # For mock testing (no VM):
-#   Set server_url to "mock" and use WAAMockAdapter internally
+#   Set server_url to "mock" in env config
 
 # --- Model ---
 model:
@@ -26,11 +42,14 @@ model:
   #   target_modules: [q_proj, k_proj, v_proj, o_proj]
 
 # --- Environment ---
+# VAGEN loads envs from env_registry.yaml using these specs.
+# WAADesktopEnv implements GymImageEnv (async reset/step/close/system_prompt).
+# Each env instance connects to the WAA server independently via HTTP.
 envs:
   - name: WAADesktop
     n_envs: 8                    # Number of parallel environments (= GRPO group size)
     data_source: waa
-    seed: [1, 100, 1]            # [start, end, step] for task selection
+    seed: [1, 100, 1]            # [start, end, step] for deterministic seeding
     max_turns: 15                # Max actions per episode
     response_length_per_turn: 512
     config:
@@ -40,27 +59,27 @@ envs:
       evaluate_at_done: true
       action_type: fractional    # VLM outputs normalized 0-1 coordinates
 
-# --- Training (GRPO) ---
+# --- Training (GRPO/GiGPO) ---
 algorithm:
   name: grpo                     # or "gigpo" for step-level advantages
   kl_coef: 0.0                   # No KL penalty (DAPO/Open-Reasoner-Zero style)
-  epsilon: 0.2                   # PPO clip range (inactive with single epoch)
-  gamma: 1.0                     # No discounting for episodic tasks
+  epsilon: 0.2                   # PPO clip range
+  gamma: 1.0                     # No discounting for episodic tasks (use 0.95 for gigpo)
 
 trainer:
   total_epochs: 100
   n_gpus_per_node: 2             # Minimum for VLM training
   micro_batch_size: 4
   gradient_accumulation_steps: 2
+  test_freq: 5                   # Evaluate every N epochs
+  experiment_name: grpo_waa_desktop
+  project_name: openadapt-waa-rl
+  logger:
+    - console
+    - wandb
 
 # --- Rollout ---
 rollout:
   temperature: 0.7
   top_p: 0.95
-  mode: async                    # async sglang rollout for throughput
-
-# --- Logging ---
-logging:
-  project: openadapt-waa-rl
-  log_interval: 1
-  save_interval: 10
+  mode: async                    # Async sglang rollout for throughput

docs/verl_agent_decision.md

@@ -250,28 +250,46 @@ By delegating to verl-agent, we avoid building and maintaining:
 
 ---
 
-## Integration Gap: verl-agent Environment Protocol
+## Integration: VAGEN Environment Registry
 
 Our `WAADesktopEnv` implements VAGEN's `GymImageEnv` protocol (async
-`reset`/`step`/`close`). However, verl-agent uses a **different, synchronous
-environment protocol** (`EnvironmentManagerBase`) with a **hardcoded dispatch**
-in `agent_system/environments/env_manager.py` — you cannot pass a Python class
-path as `env.env_name`.
-
-To integrate with verl-agent, we need to:
-
-1. **Patch `make_envs()`** — add an `elif "waa" in config.env.env_name.lower()`
-   branch (automated by `scripts/train_verl_e2e.py`)
-2. **Implement `EnvironmentManagerBase`** — wraps our async `WAADesktopEnv` in
-   verl-agent's sync vectorized env interface (`reset`, `step`, `build_text_obs`,
-   `success_evaluator`)
-3. **Prepare parquet data** — verl-agent requires `data.train_files` and
-   `data.val_files` even for env-based training
-4. **Use env-specific config** — `env.waa.server_url` instead of `env.env_kwargs`
-
-The `GymImageEnv` protocol remains our **portable interface**. The verl-agent
-`EnvironmentManagerBase` adapter is a thin sync wrapper around it. If we switch
-to a different framework, only the wrapper changes.
+`reset`/`step`/`close`/`system_prompt`), which is the **native environment
+interface** for VAGEN. No additional adapter is needed.
+
+**Note**: Earlier analysis referenced an `EnvironmentManagerBase` ABC and a
+`make_envs()` dispatch function. These do not exist in the current VAGEN
+codebase. The actual architecture uses:
+
+- `GymBaseEnv` → `GymImageEnv` — the environment ABC (which we implement)
+- `vagen/envs/registry.py` — YAML-based env registry for dispatch
+- `GymAgentLoop` — training-time rollout orchestrator that instantiates envs
+
+Integration steps (automated by `scripts/train_verl_e2e.py`):
+
+1. **Register in VAGEN's env registry** — add `WAADesktop:
+   openadapt_evals.adapters.verl_env.WAADesktopEnv` to
+   `vagen/configs/env_registry.yaml`. This is the only configuration needed.
+2. **Prepare parquet data** — VAGEN's `AgenticDataset` requires train/val
+   parquet files even for env-based training
+3. **Configure training** — provide env spec (server URL, task ID, max turns)
+   via the VAGEN training YAML (see `configs/train_waa_vagen.yaml`)
+
+The `GymImageEnv` protocol is our **portable interface**. If we switch to a
+different framework, only the ~250-line `WAADesktopEnv` adapter changes. The
+environment, evaluation, and infrastructure code remain untouched.
+
+### VAGEN Remote Env Pattern (Optional)
+
+For production deployments where the WAA VM and GPU VM have poor connectivity,
+VAGEN provides a remote env service pattern:
+
+- **Server** (WAA VM): `BaseGymHandler` + `build_gym_service()` → FastAPI
+- **Client** (GPU VM): `GymImageEnvClient` (registered as `RemoteEnv`)
+
+This adds HTTP session management, multipart encoding (JSON + images), and
+automatic retry/failover. Currently unnecessary since `WAADesktopEnv` already
+handles remote connectivity via the WAA Flask API, but documented for future
+scaling to multi-VM env pools.
 
 ---
 

openadapt_evals/adapters/__init__.py

@@ -39,7 +39,11 @@
     RLEnvironment,
     RolloutStep,
 )
-from openadapt_evals.adapters.verl_env import WAADesktopEnv
+from openadapt_evals.adapters.verl_env import (
+    WAADesktopEnv,
+    generate_env_spec,
+    register_in_vagen,
+)
 from openadapt_evals.adapters.waa import (
     WAAAdapter,
     WAAConfig,
@@ -72,6 +76,8 @@
     "WAALiveConfig",
     # verl-agent / VAGEN integration
     "WAADesktopEnv",
+    "register_in_vagen",
+    "generate_env_spec",
     # Task ID validation
     "SyntheticTaskError",
     "is_real_waa_task_id",

openadapt_evals/adapters/verl_env.py

@@ -14,11 +14,15 @@
       as fallback when the full vagen package is not installed)
 
 Usage with VAGEN training:
-    Register in env_registry.yaml:
+    1. Register the env on the GPU VM:
+        from openadapt_evals.adapters.verl_env import register_in_vagen
+        register_in_vagen("/path/to/vagen/configs/env_registry.yaml")
+
+    2. Or manually add to vagen/configs/env_registry.yaml:
         env_registry:
             WAADesktop: openadapt_evals.adapters.verl_env.WAADesktopEnv
 
-    Training config:
+    3. Training config (envs section):
         envs:
             - name: WAADesktop
               n_envs: 8
@@ -44,6 +48,7 @@
 import io
 import logging
 import re
+from pathlib import Path
 from typing import Any
 
 from openadapt_evals.adapters.base import BenchmarkAction, BenchmarkObservation
@@ -354,3 +359,110 @@ async def step(
         )
 
         return obs_dict, reward, done, info
+
+
+# --- VAGEN registration helpers ---
+
+ENV_REGISTRY_KEY = "WAADesktop"
+ENV_CLASS_PATH = "openadapt_evals.adapters.verl_env.WAADesktopEnv"
+
+
+def register_in_vagen(registry_yaml_path: str | Path | None = None) -> bool:
+    """Register WAADesktopEnv in VAGEN's env_registry.yaml.
+
+    If VAGEN is installed and importable, also registers via the Python API.
+    Otherwise, edits the YAML file directly.
+
+    Args:
+        registry_yaml_path: Path to vagen/configs/env_registry.yaml.
+            If None, attempts to find it from the vagen package location.
+
+    Returns:
+        True if registration succeeded.
+    """
+    # Try programmatic registration first (if vagen is installed)
+    try:
+        from vagen.envs.registry import register_env
+
+        register_env(ENV_REGISTRY_KEY, WAADesktopEnv)
+        logger.info("Registered %s in VAGEN env registry (Python API)", ENV_REGISTRY_KEY)
+        return True
+    except ImportError:
+        pass
+
+    # Fall back to YAML file editing
+    if registry_yaml_path is None:
+        # Try to find it from common locations
+        candidates = [
+            Path.home() / "verl-agent" / "vagen" / "configs" / "env_registry.yaml",
+            Path.home() / "VAGEN" / "vagen" / "configs" / "env_registry.yaml",
+        ]
+        for candidate in candidates:
+            if candidate.exists():
+                registry_yaml_path = candidate
+                break
+
+    if registry_yaml_path is not None:
+        registry_yaml_path = Path(registry_yaml_path)
+        if not registry_yaml_path.exists():
+            logger.warning("Registry file not found: %s", registry_yaml_path)
+            return False
+
+    if registry_yaml_path is None:
+        logger.warning(
+            "Cannot find env_registry.yaml. Register manually by adding:\n"
+            "  %s: %s\n"
+            "to vagen/configs/env_registry.yaml",
+            ENV_REGISTRY_KEY,
+            ENV_CLASS_PATH,
+        )
+        return False
+
+    content = registry_yaml_path.read_text()
+
+    if ENV_REGISTRY_KEY in content:
+        logger.info("%s already registered in %s", ENV_REGISTRY_KEY, registry_yaml_path)
+        return True
+
+    # Add entry to the env_registry section
+    entry = f"  {ENV_REGISTRY_KEY}: {ENV_CLASS_PATH}\n"
+    if "env_registry:" in content:
+        content = content.replace("env_registry:\n", f"env_registry:\n{entry}", 1)
+    else:
+        content += f"\nenv_registry:\n{entry}"
+
+    registry_yaml_path.write_text(content)
+    logger.info("Registered %s in %s", ENV_REGISTRY_KEY, registry_yaml_path)
+    return True
+
+
+def generate_env_spec(
+    server_url: str = "http://localhost:5001",
+    task_id: str = "REPLACE_WITH_WAA_TASK_UUID",
+    n_envs: int = 8,
+    max_turns: int = 15,
+) -> dict[str, Any]:
+    """Generate a VAGEN EnvSpec dict for WAA desktop training.
+
+    Returns a dict suitable for inclusion in a VAGEN training config YAML
+    under the ``envs`` key.
+
+    Example:
+        spec = generate_env_spec(server_url="http://10.0.0.5:5001", task_id="abc-123")
+        # Write to YAML or pass to AgenticDataset
+    """
+    return {
+        "name": ENV_REGISTRY_KEY,
+        "n_envs": n_envs,
+        "data_source": "waa",
+        "seed": [1, 100, 1],
+        "max_turns": max_turns,
+        "response_length_per_turn": 512,
+        "config": {
+            "server_url": server_url,
+            "task_id": task_id,
+            "max_steps": max_turns,
+            "evaluate_at_done": True,
+            "action_type": "fractional",
+        },
+    }

openadapt_evals/benchmarks/vm_cli.py

@@ -7919,36 +7919,33 @@ def cmd_gpu_train(args):
                 print(f"ERROR: Setup failed")
                 return 1
 
-    # Launch training via the E2E script (handles data prep, env patching, config)
+    # Launch training via the E2E script (handles data prep, env registration, config)
     # Import here to avoid circular deps
     import sys
     sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
-    from train_verl_e2e import prepare_training_data, patch_env_manager
+    from train_verl_e2e import prepare_training_data, register_waa_env
+
+    print("Registering WAADesktopEnv in VAGEN env registry...")
+    register_waa_env(ip, args.waa_server, args.task_id, username=username)
 
     print("Preparing training data...")
     prepare_training_data(ip, group_size=8, username=username)
 
-    print("Patching verl-agent for WAA environment...")
-    patch_env_manager(ip, args.waa_server, args.task_id, username=username)
-
-    # Validated Hydra config (see docs/verl_agent_decision.md)
+    # Hydra overrides for verl training loop + VAGEN env config.
+    # WAADesktopEnv is registered in VAGEN's env_registry.yaml as 'WAADesktop'.
+    # The env connects to WAA server via HTTP (GymImageEnv protocol).
     train_cmd = (
         f"cd ~/verl-agent && "
         f"conda run -n verl-agent python3 -m verl.trainer.main_ppo "
         f"algorithm.adv_estimator={args.algorithm} "
-        f"algorithm.gamma=0.95 "
+        f"algorithm.gamma={'0.95' if args.algorithm == 'gigpo' else '1.0'} "
         f"actor_rollout_ref.model.path={args.model} "
         f"actor_rollout_ref.rollout.name=vllm "
         f"actor_rollout_ref.rollout.tensor_model_parallel_size={args.n_gpus} "
         f"actor_rollout_ref.rollout.gpu_memory_utilization=0.6 "
         f"actor_rollout_ref.rollout.enable_chunked_prefill=False "
         f"actor_rollout_ref.actor.ppo_mini_batch_size=64 "
         f"actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 "
-        f"env.env_name=waa_desktop "
-        f"env.max_steps=15 "
-        f"env.rollout.n=8 "
-        f"env.waa.server_url={args.waa_server} "
-        f"env.waa.task_id={args.task_id} "
         f"data.train_files=$HOME/data/verl-agent/visual/train.parquet "
         f"data.val_files=$HOME/data/verl-agent/visual/test.parquet "
         f"data.train_batch_size=8 "