feat: vendor GymImageEnv base classes from VAGEN

abrichr · claude · abrichr · commit e2b6d0f9a994 · 2026-03-02T23:09:19.000-05:00
Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/docs/verl_agent_decision.md b/docs/verl_agent_decision.md
@@ -234,7 +234,9 @@ By delegating to verl-agent, we avoid building and maintaining:
 
 3. **Next**: Test end-to-end with verl-agent on a GPU machine. If successful,
    the standalone trainer becomes a reference implementation / fallback, and
-   verl-agent becomes the recommended training path.
+   verl-agent becomes the recommended training path. **Note**: Both backends
+   coexist — see the [Dual Backend Strategy](#dual-backend-strategy) section
+   for the comparison plan and dependency approach.
 
 4. **Future**: If TRL resolves #5120 (multi-turn VLM support), evaluate whether
    to switch. TRL has broader adoption; switching would reduce the dependency
@@ -243,6 +245,50 @@ By delegating to verl-agent, we avoid building and maintaining:
 
 ---
 
+## Dual Backend Strategy
+
+Rather than deprecating the standalone trainer immediately, we maintain both
+backends for comparison:
+
+### Backend 1: Standalone (openadapt-ml)
+
+- **Code**: `openadapt_ml/training/grpo/trainer.py` (~546 lines)
+- **When to use**: Quick experiments, single-GPU, no Ray/vLLM dependency
+- **Limitations**: Episode-level rewards only, no GiGPO, no distributed training
+- **Config**: `GRPOConfig(backend="standalone", ...)`
+
+### Backend 2: verl-agent (openadapt-evals)
+
+- **Code**: `openadapt_evals/adapters/verl_env.py` (~250 lines adapter)
+- **When to use**: Production training, multi-GPU, GiGPO per-step credit
+- **Advantages**: Distributed training, vLLM/sglang, step-level advantages
+- **Config**: `configs/train_waa_vagen.yaml`
+
+### Dependency Strategy
+
+The `GymImageEnv` and `GymBaseEnv` abstract base classes (~150 lines) are
+**vendored** into `openadapt_evals/adapters/_vendored/` to avoid requiring the
+full VAGEN installation. The vendored classes are pure interfaces with only a
+`Pillow` dependency. Import priority:
+
+1. `from vagen.envs.gym_image_env import GymImageEnv` (if VAGEN installed)
+2. `from openadapt_evals.adapters._vendored.gym_image_env import GymImageEnv` (fallback)
+
+The full VAGEN/verl-agent stack (Ray, vLLM, etc.) is only needed when actually
+running distributed training, not for defining or testing environments.
+
+### Comparison Plan
+
+To validate the verl-agent integration provides real value over standalone:
+
+1. Train on the same WAA task with both backends
+2. Compare: final reward, training wall time, GPU memory usage
+3. Specifically measure whether GiGPO's per-step credit improves sample
+   efficiency on long-horizon tasks (15+ steps)
+4. Document results in a comparison report
+
+---
+
 ## References
 
 - [verl-agent](https://github.com/langfengQ/verl-agent) — GiGPO paper implementation
diff --git a/openadapt_evals/adapters/_vendored/__init__.py b/openadapt_evals/adapters/_vendored/__init__.py
@@ -0,0 +1,10 @@
+"""Vendored pure-abstract base classes from VAGEN.
+
+These are copied from https://github.com/mll-lab-nu/VAGEN so that
+openadapt-evals can implement the GymImageEnv protocol without
+requiring the full VAGEN package (and its heavy transitive
+dependencies) to be installed.
+
+Only the abstract interface definitions are vendored here -- no
+concrete implementations or utilities.
+"""
diff --git a/openadapt_evals/adapters/_vendored/gym_base_env.py b/openadapt_evals/adapters/_vendored/gym_base_env.py
@@ -0,0 +1,40 @@
+# Vendored from https://github.com/mll-lab-nu/VAGEN
+# These are pure abstract base classes with no heavy dependencies.
+# Vendored to avoid requiring the full VAGEN installation.
+# Last synced: 2026-03-02
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Tuple
+
+
+class GymBaseEnv(ABC):
+    """
+    Abstract async environment API.
+    The handler does not assume any obs/data schema beyond what you return.
+
+    Contract:
+      - reset(seed) -> (obs, info)
+      - step(action_str) -> (obs, reward, done, info)
+    """
+
+    def __init__(self, env_config: Dict[str, Any]):
+        self.config = env_config
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Async teardown."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def reset(self, seed: int):
+        raise NotImplementedError
+
+    @abstractmethod
+    async def step(self, action_str: str):
+        raise NotImplementedError
+
+    @abstractmethod
+    async def system_prompt(self):
+        raise NotImplementedError
diff --git a/openadapt_evals/adapters/_vendored/gym_image_env.py b/openadapt_evals/adapters/_vendored/gym_image_env.py
@@ -0,0 +1,212 @@
+# Vendored from https://github.com/mll-lab-nu/VAGEN
+# These are pure abstract base classes with no heavy dependencies.
+# Vendored to avoid requiring the full VAGEN installation.
+# Last synced: 2026-03-02
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, Dict, Tuple
+
+from .gym_base_env import GymBaseEnv
+
+
+class GymImageEnv(GymBaseEnv):
+    """
+    GymImageEnv is a base environment class that supports optional
+    **image-based multi-modal observations**, while keeping the same API
+    as GymBaseEnv.
+
+    --------------------------------------------------------------------
+    Observation Protocol
+    --------------------------------------------------------------------
+
+    WITH images
+    -------------------------------
+    If the environment returns images, the observation should follow:
+
+        obs = {
+            "obs_str": "... <image> ...",
+            "multi_modal_input": {
+                "<image>": [PIL.Image.Image, ...]
+            }
+        }
+
+    - Images are stored under obs["multi_modal_input"]["<image>"].
+    - "<image>" in obs_str is a placeholder indicating where each image
+      should appear in the prompt.
+    - The number of "<image>" in obs_str should match the number of
+      images in the list.
+
+    WITHOUT images:
+    ----------------------------------
+    Can simply use:
+
+        obs = {
+            "obs_str": "..."
+        }
+
+    - "multi_modal_input" is optional and may be omitted.
+    - obs_str should NOT contain "<image>" placeholders.
+
+
+    --------------------------------------------------------------------
+    Agent-Loop Rollout
+    --------------------------------------------------------------------
+    - sys      : system prompt (from system_prompt()).
+    - init_obs : observation from reset().
+    - step_obs : observation from step().
+    - res_i    : agent response at step i.
+
+    Concat mode (single growing context):
+        sys + init_obs + res_0 + step_obs_1 + res_1 + ...
+
+    Non-concat mode (step-wise independent contexts):
+        Step 0: sys + init_obs   + res_0
+        Step 1: sys + step_obs_1 + res_1
+        Step 2: sys + step_obs_2 + res_2
+
+    --------------------------------------------------------------------
+    Info
+    --------------------------------------------------------------------
+    The `info` dict returned by reset() and step() may include:
+        - success (bool): whether the task/episode is considered
+          successful, this will be used for wandb logging.
+    """
+
+    def __init__(self, env_config: Dict[str, Any]):
+        """
+        Initialize the environment.
+
+        Args:
+            env_config (Dict[str, Any]):
+                Environment configuration. The exact schema is defined by
+                the concrete environment implementation and/or GymBaseEnv.
+
+        Side effects:
+            - Calls GymBaseEnv.__init__(env_config).
+        """
+        super().__init__(env_config)
+
+    @abstractmethod
+    async def close(self) -> None:
+        """
+        Close the environment and release all resources.
+
+        This should clean up anything created by the environment, e.g.:
+        - windows / renderers
+        - subprocesses
+        - file handles
+        - GPU memory / models
+
+        Returns:
+            None
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def system_prompt(self) -> Dict[str, Any]:
+        """
+        Return the system-level prompt/observation for the environment.
+
+        Returns:
+            obs (Dict[str, Any]):
+                A dict representing the system prompt observation.
+
+                If returning images, it must follow:
+
+                    obs = {
+                        "obs_str": "... <image> ...",
+                        "multi_modal_input": {
+                            "<image>": [PIL.Image.Image, ...]
+                        }
+                    }
+
+                If returning no images, it must follow:
+
+                    obs = {
+                        "obs_str": "..."
+                    }
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def reset(self, seed: int) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Reset the environment to the initial state.
+
+        Args:
+            seed (int):
+                Random seed used to initialize the environment
+
+        Returns:
+            obs (Dict[str, Any]):
+                The initial observation after reset.
+
+                If returning images, it must follow:
+
+                    obs = {
+                        "obs_str": "... <image> ...",
+                        "multi_modal_input": {
+                            "<image>": [PIL.Image.Image, ...]
+                        }
+                    }
+
+                If returning no images, it must follow:
+
+                    obs = {
+                        "obs_str": "..."
+                    }
+
+            info (Dict[str, Any]):
+                A dict containing any additional metadata about the reset,
+                e.g. debug information, episode identifiers, etc.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def step(
+        self, action_str: str
+    ) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
+        """
+        Execute one environment step using an agent-provided action.
+
+        Args:
+            action_str (str):
+                The action produced by the agent, in text form.
+
+        Returns:
+            obs (Dict[str, Any]):
+                The next observation after applying the action.
+
+                If returning images, it must follow:
+
+                    obs = {
+                        "obs_str": "... <image> ...",
+                        "multi_modal_input": {
+                            "<image>": [PIL.Image.Image, ...]
+                        }
+                    }
+
+                If returning no images, it must follow:
+
+                    obs = {
+                        "obs_str": "..."
+                    }
+
+            reward (float):
+                Scalar reward for the current step.
+
+            done (bool):
+                Whether the current episode has terminated after this
+                step.
+
+            info (Dict[str, Any]):
+                Additional step-level metadata.
+
+                Common optional keys:
+                    - success (bool): whether the task/episode is
+                      considered successful, typically used for logging
+                      (e.g. wandb).
+        """
+        raise NotImplementedError
diff --git a/openadapt_evals/adapters/verl_env.py b/openadapt_evals/adapters/verl_env.py
@@ -10,8 +10,8 @@
 
 Dependencies:
     - openadapt-evals (always required)
-    - vagen (optional; without it, the class implements the same protocol
-      but doesn't inherit from GymImageEnv)
+    - vagen (optional; a vendored copy of the GymImageEnv ABC is used
+      as fallback when the full vagen package is not installed)
 
 Usage with VAGEN training:
     Register in env_registry.yaml:
@@ -57,13 +57,14 @@
 except ImportError:
     Image = None  # type: ignore[misc, assignment]
 
-# Try importing VAGEN's base class for proper inheritance
+# Import VAGEN's GymImageEnv base class.
+# Prefer the real vagen package; fall back to our vendored copy.
 try:
     from vagen.envs.gym_image_env import GymImageEnv as _GymImageEnvBase
-
-    _HAS_VAGEN = True
 except ImportError:
-    _HAS_VAGEN = False
+    from openadapt_evals.adapters._vendored.gym_image_env import (
+        GymImageEnv as _GymImageEnvBase,
+    )
 
 # --- Action parsing (matches openadapt-ml trainer DSL) ---
 
@@ -171,19 +172,8 @@ def _build_obs_dict(
 
 # --- Main environment class ---
 
-# Dynamically choose base class
-if _HAS_VAGEN:
-    _Base = _GymImageEnvBase
-else:
-
-    class _Base:  # type: ignore[no-redef]
-        """Stub base when VAGEN is not installed."""
-
-        def __init__(self, env_config: dict[str, Any]) -> None:
-            pass
-
 
-class WAADesktopEnv(_Base):
+class WAADesktopEnv(_GymImageEnvBase):
     """VAGEN-compatible environment for WAA desktop automation.
 
     Implements the GymImageEnv protocol (async reset/step/close/system_prompt)
