refactor: make openadapt-ml a leaf; break ml<->evals import cycle

openadapt_ml/benchmarks/agent.py

@@ -32,8 +32,10 @@
 import re
 from typing import TYPE_CHECKING, Any
 
-# Import base classes from openadapt-evals (canonical location)
-from openadapt_evals import (
+# Import base classes from openadapt-types (canonical schema package).
+# These used to live in openadapt-evals; importing them from the schema
+# package keeps openadapt-ml a leaf (no module-level ml -> evals import).
+from openadapt_types import (
     BenchmarkAction,
     BenchmarkAgent,
     BenchmarkObservation,

openadapt_ml/experiments/waa_demo/runner.py

@@ -38,7 +38,7 @@
 )
 
 if TYPE_CHECKING:
-    from openadapt_evals import (
+    from openadapt_types import (
         BenchmarkAction,
         BenchmarkObservation,
         BenchmarkTask,
@@ -305,7 +305,7 @@ def act(
         Returns:
             BenchmarkAction parsed from VLM response
         """
-        from openadapt_evals import BenchmarkAction
+        from openadapt_types import BenchmarkAction
 
         adapter = self._get_adapter()
 
@@ -447,7 +447,7 @@ def _parse_response(
         Uses the same parsing logic as APIBenchmarkAgent.
         """
         import re
-        from openadapt_evals import BenchmarkAction
+        from openadapt_types import BenchmarkAction
 
         raw_action = {"response": response}
 

openadapt_ml/training/grpo/rollout_collector.py

@@ -14,28 +14,20 @@
 import logging
 import random
 from dataclasses import dataclass, field
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 from openadapt_ml.training.grpo.config import GRPOConfig
 from openadapt_ml.training.grpo.reward import binary_task_success
+from openadapt_ml.training.grpo.rollout_env import RolloutEnv
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    # The concrete environment/adapter live in openadapt-evals and implement
+    # the RolloutEnv Protocol defined in openadapt-ml. They are imported
+    # lazily (inside __init__) at runtime to keep openadapt-ml a leaf with no
+    # module-level openadapt-evals import.
+    from openadapt_evals.adapters import WAALiveAdapter
 
-# Deferred imports for openadapt-evals dependencies (optional at install time)
-try:
-    from openadapt_evals.adapters import (
-        RLEnvironment,
-        RolloutStep,
-        WAALiveAdapter,
-        WAALiveConfig,
-    )
-    from openadapt_evals.adapters.rl_env import ResetConfig
-except ImportError:
-    RLEnvironment = None  # type: ignore[assignment, misc]
-    RolloutStep = None  # type: ignore[assignment, misc]
-    WAALiveAdapter = None  # type: ignore[assignment, misc]
-    WAALiveConfig = None  # type: ignore[assignment, misc]
-    ResetConfig = None  # type: ignore[assignment, misc]
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -81,25 +73,35 @@ def __init__(
         config: GRPOConfig,
         task_configs: dict[str, Any] | None = None,
     ) -> None:
-        if RLEnvironment is None:
+        # Lazy import: the concrete WAA adapter + RLEnvironment live in
+        # openadapt-evals and implement the RolloutEnv Protocol. Importing
+        # them here (not at module scope) keeps openadapt-ml a leaf.
+        try:
+            from openadapt_evals.adapters import (
+                RLEnvironment,
+                WAALiveAdapter,
+                WAALiveConfig,
+            )
+        except ImportError as exc:
             raise ImportError(
                 "openadapt-evals is required for rollout collection. "
                 "Install it with: uv add openadapt-evals"
-            )
+            ) from exc
 
         self._config = config
         self._task_configs = task_configs or {}
-        self._adapter = WAALiveAdapter(
+        self._adapter: WAALiveAdapter = WAALiveAdapter(
             WAALiveConfig(
                 server_url=config.server_url,
                 evaluate_url=config.evaluate_url,
             )
         )
-        self._env = RLEnvironment(self._adapter)
+        # RLEnvironment structurally satisfies the RolloutEnv Protocol.
+        self._env: RolloutEnv = RLEnvironment(self._adapter)
 
     @property
-    def env(self) -> Any:
-        """The underlying RLEnvironment instance."""
+    def env(self) -> RolloutEnv:
+        """The underlying environment (implements the RolloutEnv Protocol)."""
         return self._env
 
     def collect_group(

openadapt_ml/training/grpo/rollout_env.py

@@ -0,0 +1,64 @@
+"""Minimal environment interface driven by GRPO training.
+
+RL training in ``openadapt-ml`` needs to *drive* an environment (reset it,
+step actions, observe, and collect whole rollouts), but the *concrete*
+environment is an evaluation-harness concern that lives in
+``openadapt-evals`` (e.g. ``RLEnvironment``, ``WAALiveAdapter``,
+``WAADesktopEnv``).
+
+To keep ``openadapt-ml`` a dependency **leaf** (no module-level import of
+``openadapt-evals``), the trainer types against this thin ``RolloutEnv``
+Protocol instead of a concrete class. The concrete adapters in
+``openadapt-evals`` implement it structurally.
+
+Only the surface the GRPO trainer/collector actually uses is declared here;
+signatures are intentionally loose (``Any``) so the concrete evals
+implementations conform structurally without importing evals-specific
+config types into ml.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+from openadapt_types import BenchmarkAction, BenchmarkObservation
+
+
+@runtime_checkable
+class RolloutEnv(Protocol):
+    """Structural interface for an environment the GRPO trainer can drive.
+
+    Implemented by ``openadapt_evals.adapters.rl_env.RLEnvironment`` (and
+    compatible desktop environments). ml depends on this interface; evals
+    provides the implementation.
+    """
+
+    @property
+    def screen_size(self) -> tuple[int, int]:
+        """Current environment screen size as (width, height)."""
+        ...
+
+    def reset(self, config: Any = None) -> BenchmarkObservation:
+        """Reset the environment and return the initial observation."""
+        ...
+
+    def step(self, action: BenchmarkAction) -> Any:
+        """Execute an action, returning a step result (obs/reward/done/info)."""
+        ...
+
+    def observe(self) -> BenchmarkObservation:
+        """Return the current observation without stepping."""
+        ...
+
+    def collect_rollout(
+        self,
+        agent_fn: Any,
+        max_steps: int = ...,
+        stuck_window: int = ...,
+        task_id: Any = None,
+    ) -> list[Any]:
+        """Run ``agent_fn`` to completion, returning the list of rollout steps."""
+        ...
+
+
+__all__ = ["RolloutEnv"]

openadapt_ml/training/grpo/trainer.py

@@ -54,15 +54,6 @@
     Rollout,
 )
 
-# Optional import for TaskConfig (openadapt-evals may not be installed)
-try:
-    from openadapt_evals.task_config import TaskConfig
-
-    _HAS_TASK_CONFIG = True
-except ImportError:
-    TaskConfig = None  # type: ignore[assignment, misc]
-    _HAS_TASK_CONFIG = False
-
 logger = logging.getLogger(__name__)
 
 DEFAULT_SCREEN_SIZE: tuple[int, int] = (1920, 1080)
@@ -166,18 +157,7 @@ def _parse_vlm_output_to_action(
 
     Supports: CLICK(x=0.XX, y=0.XX), TYPE(text="..."), WAIT(), DONE().
     """
-    try:
-        from openadapt_evals.adapters.base import BenchmarkAction
-    except ImportError:
-        from dataclasses import dataclass as _dc
-
-        @_dc
-        class BenchmarkAction:  # type: ignore[no-redef]
-            type: str = "done"
-            x: float | None = None
-            y: float | None = None
-            text: str | None = None
-            key: str | None = None
+    from openadapt_types import BenchmarkAction
 
     text = text.strip()
     width, height = screen_size
@@ -371,11 +351,16 @@ def _load_task_configs(self, task_dir: str) -> None:
             ImportError: If openadapt-evals is not installed.
             FileNotFoundError: If the directory does not exist.
         """
-        if not _HAS_TASK_CONFIG:
+        # Lazy import: the concrete TaskConfig is an eval-harness concept.
+        # Keeping it out of module scope keeps openadapt-ml a leaf (no
+        # module-level openadapt-evals import).
+        try:
+            from openadapt_evals.task_config import TaskConfig
+        except ImportError as exc:
             raise ImportError(
                 "openadapt-evals is required for --task-dir support. "
                 "Install with: pip install openadapt-evals"
-            )
+            ) from exc
 
         task_dir_path = Path(task_dir)
         if not task_dir_path.is_dir():
@@ -459,7 +444,7 @@ def _make_agent_fn(self) -> Callable:
         Captures model/processor by reference so weight updates during
         training are reflected in subsequent rollouts.
         """
-        from openadapt_evals.adapters.base import BenchmarkAction
+        from openadapt_types import BenchmarkAction
 
         model = self._model
         processor = self._processor

openadapt_ml/training/grpo/verl_backend.py

@@ -33,11 +33,19 @@
 
 logger = logging.getLogger(__name__)
 
-# Deferred import for openadapt-evals WAADesktopEnv (optional dependency)
-try:
-    from openadapt_evals.adapters.verl_env import WAADesktopEnv
-except ImportError:
-    WAADesktopEnv = None  # type: ignore[assignment, misc]
+
+def _load_waa_desktop_env() -> Any | None:
+    """Lazily import the concrete WAADesktopEnv from openadapt-evals.
+
+    Kept out of module scope so openadapt-ml has no module-level
+    openadapt-evals import (ml stays a dependency leaf). Returns the class,
+    or ``None`` if openadapt-evals is not installed.
+    """
+    try:
+        from openadapt_evals.adapters.verl_env import WAADesktopEnv
+    except ImportError:
+        return None
+    return WAADesktopEnv
 
 
 def build_vagen_config(config: GRPOConfig) -> dict[str, Any]:
@@ -96,7 +104,7 @@ def train_with_verl(config: GRPOConfig) -> None:
     """
     vagen_config = build_vagen_config(config)
 
-    if WAADesktopEnv is not None:
+    if _load_waa_desktop_env() is not None:
         logger.info(
             "WAADesktopEnv is available. verl-agent can use it for "
             "desktop environment interaction."

pyproject.toml

@@ -28,6 +28,10 @@ dependencies = [
     "matplotlib>=3.10.7",
     "modal>=1.3.4",
     "openadapt-capture>=0.3.0",
+    # Canonical schema package. Benchmark* types (Task/Observation/Action/Agent)
+    # live here so ml and evals don't import each other (breaks ml<->evals cycle).
+    # NOTE: requires the openadapt-types release that adds openadapt_types.benchmark.
+    "openadapt-types>=0.3.0",
     "pillow>=12.0.0",
     "pyautogui>=0.9.54",
     "pydantic-settings>=2.0.0",

tests/test_grpo.py

@@ -186,16 +186,18 @@ def test_rollout_defaults():
 
 def test_rollout_collector_requires_evals():
     """GRPORolloutCollector raises ImportError without openadapt-evals."""
-    from openadapt_ml.training.grpo.rollout_collector import (
-        GRPORolloutCollector,
-        RLEnvironment,
-    )
+    import importlib.util
+
+    from openadapt_ml.training.grpo.rollout_collector import GRPORolloutCollector
     from openadapt_ml.training.grpo import GRPOConfig
 
-    # If openadapt-evals is not installed, this should raise ImportError.
-    # If it IS installed, the constructor will try to connect (which we skip).
+    # The concrete env/adapter are imported lazily inside __init__ (openadapt-ml
+    # is a leaf and has no module-level openadapt-evals import). If evals is not
+    # installed, constructing the collector should raise ImportError. If it IS
+    # installed, the constructor will try to connect (which we skip).
     config = GRPOConfig(server_url="http://localhost:99999")
-    if RLEnvironment is None:
+    evals_installed = importlib.util.find_spec("openadapt_evals") is not None
+    if not evals_installed:
         with pytest.raises(ImportError, match="openadapt-evals"):
             GRPORolloutCollector(config)