diff --git a/openadapt_ml/benchmarks/agent.py b/openadapt_ml/benchmarks/agent.py index 926d5cc..dec64a9 100644 --- a/openadapt_ml/benchmarks/agent.py +++ b/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, diff --git a/openadapt_ml/experiments/waa_demo/runner.py b/openadapt_ml/experiments/waa_demo/runner.py index a11f895..523d796 100644 --- a/openadapt_ml/experiments/waa_demo/runner.py +++ b/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} diff --git a/openadapt_ml/training/grpo/rollout_collector.py b/openadapt_ml/training/grpo/rollout_collector.py index 55c9be1..56e5af1 100644 --- a/openadapt_ml/training/grpo/rollout_collector.py +++ b/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( diff --git a/openadapt_ml/training/grpo/rollout_env.py b/openadapt_ml/training/grpo/rollout_env.py new file mode 100644 index 0000000..00c59b1 --- /dev/null +++ b/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"] diff --git a/openadapt_ml/training/grpo/trainer.py b/openadapt_ml/training/grpo/trainer.py index b13656f..852df5c 100644 --- a/openadapt_ml/training/grpo/trainer.py +++ b/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 diff --git a/openadapt_ml/training/grpo/verl_backend.py b/openadapt_ml/training/grpo/verl_backend.py index 3904a5c..546edf9 100644 --- a/openadapt_ml/training/grpo/verl_backend.py +++ b/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." diff --git a/pyproject.toml b/pyproject.toml index 1070453..e4a90bd 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/test_grpo.py b/tests/test_grpo.py index dee3ebb..8035e8c 100644 --- a/tests/test_grpo.py +++ b/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)