feat(training): train a bundle's own agents with AReaL's agent workflows

.github/workflows/areal.yml

@@ -0,0 +1,60 @@
+name: AReaL contract
+
+# AgentEnvWorkflow in AReaL's agent workflow (tst/areal): AReaL's rollout engine resolves it from its import path into
+# its proxy workflow in inline mode, grouped, and plays the unit tier's bundle through a stand-in model, on AReaL 2.1.0
+# and on its main branch. AReaL and torch are installed over the dev environment; torch is the CPU build, as nothing
+# generates. It runs on pull requests that touch the helper or the tier, weekly, and on demand; it isn't a required
+# check.
+
+on:
+  pull_request:
+    paths:
+      - "src/agent_env/training/**"
+      - "tst/areal/**"
+      - "tst/unit/training/**"
+      - "tst/util/fake_vllm.py"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/workflows/areal.yml"
+  schedule:
+    - cron: "53 7 * * 1"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: areal-${{ github.event_name == 'pull_request' && github.ref || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+env:
+  PIP_DISABLE_PIP_VERSION_CHECK: "1"
+
+jobs:
+  workflow:
+    name: agent workflow (AReaL ${{ matrix.areal }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        areal: ["v2.1.0", "main"]
+    steps:
+      - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+      - name: Install from uv.lock
+        run: |
+          python -m pip install --quiet -r .github/ci-requirements.txt
+          uv sync --locked --extra dev --python "$(command -v python)"
+      - name: Install AReaL and CPU torch
+        run: |
+          uv pip install --python .venv/bin/python "areal @ git+https://github.com/areal-project/AReaL@${{ matrix.areal }}" torch \
+            --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match
+          .venv/bin/python -c "from importlib.metadata import version; print('areal', version('areal'), 'torch', version('torch'))"
+      - name: AReaL tier
+        run: .venv/bin/python -m pytest tst/areal -rA --timeout=900 --junitxml=junit-areal.xml
+      - name: No skips
+        if: always()
+        run: .venv/bin/python .github/scripts/check_skip_policy.py junit-areal.xml

Makefile

@@ -1,4 +1,4 @@
-.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test clean-install-test test clean
+.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test areal-test clean-install-test test clean
 
 VENV := .venv
 VENV_BIN := $(abspath $(VENV))/bin
@@ -85,6 +85,12 @@ verl-test: ## Play rows through the verl agent loop on verl itself (needs verl a
 skyrl-test: ## Play rows through the SkyRL env and generator on SkyRL itself (needs SkyRL and CPU torch, and the Hub)
 	$(PYTHON) -m pytest tst/skyrl --timeout=900 $(PYTEST_ARGS)
 
+# The AReaL tier plays rows through AgentEnvWorkflow in AReaL's agent workflow, resolved by AReaL's rollout engine.
+# AReaL and torch aren't in the dev extra, and AReaL needs Python 3.11 or 3.12: install them first, as
+# .github/workflows/areal.yml does.
+areal-test: ## Play rows through AgentEnvWorkflow in AReaL's agent workflow (needs AReaL and CPU torch)
+	$(PYTHON) -m pytest tst/areal --timeout=900 $(PYTEST_ARGS)
+
 # The clean-install gate CI runs, on the same Python. The script gives every step an allowlisted
 # environment, so nothing in this shell reaches the install it checks. A failed run keeps its folder.
 clean-install-test: ## Build both distributions, install them clean from public PyPI and run hello twice (needs python3.11, uv, git)

src/agent_env/openenv/app.py

@@ -21,9 +21,9 @@
 from agent_env.bundle.installed import locate_bundle
 from agent_env.store.routing import namespace_routing
 from agent_env.task import EpisodeError, EpisodeSpec, Task
+from agent_env.training.reward import reward_of
 
 from .environment import AgentEnvEnvironment, Served, ServedTask
-from .reward import reward_of
 
 logger = logging.getLogger(__name__)
 

src/agent_env/training/__init__.py

@@ -1,4 +1,5 @@
 """Training on AgentEnv tasks with other projects' trainers, one module per trainer, each behind an extra of its own:
 ``agent_env.training.trl`` for TRL, ``agent_env.training.nemo_gym`` for NeMo Gym and the trainers that take its
 environments, NeMo RL among them, ``agent_env.training.verl`` for verl, and ``agent_env.training.skyrl`` for
-SkyRL."""
+SkyRL. ``agent_env.training.harness`` and ``agent_env.training.areal`` train a bundle's own agents instead, with TRL's
+harness rollouts and AReaL's agent workflows; the second needs no extra."""

src/agent_env/training/_loop.py

@@ -0,0 +1,31 @@
+"""The event loop this process's training sessions share, on a daemon thread of its own: a trainer calls its
+environments from threads, or awaits them on an event loop of its own that a run's blocking calls mustn't stall."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Coroutine
+from typing import Any
+
+_loop: asyncio.AbstractEventLoop | None = None
+_loop_lock = threading.Lock()
+
+
+def on_session_loop(coroutine: Coroutine[Any, Any, Any], timeout: float | None) -> Any:
+    """``coroutine``'s result, run on the session loop; waits at most ``timeout`` seconds."""
+    future = asyncio.run_coroutine_threadsafe(coroutine, session_loop())
+    try:
+        return future.result(timeout)
+    except TimeoutError:
+        future.cancel()
+        raise
+
+
+def session_loop() -> asyncio.AbstractEventLoop:
+    global _loop
+    with _loop_lock:
+        if _loop is None:
+            _loop = asyncio.new_event_loop()
+            threading.Thread(target=_loop.run_forever, name="agentenv-sessions", daemon=True).start()
+        return _loop

src/agent_env/training/_openenv.py

@@ -13,7 +13,7 @@
 import threading
 import time
 import uuid
-from collections.abc import Coroutine, Iterator, Mapping
+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass, field
 from typing import Any, ClassVar
 
@@ -50,8 +50,6 @@
 _waited: set[str] = set()
 #: The servers whose refused connections rollouts have been warned they wait out, once a process.
 _restarting: set[str] = set()
-_loop: asyncio.AbstractEventLoop | None = None
-_loop_lock = threading.Lock()
 
 
 @dataclass(frozen=True)
@@ -426,26 +424,6 @@ def toolsets(name: str, tasks: dict[str, list[dict[str, Any]]]) -> dict[str, Nam
     return named
 
 
-def on_session_loop(coroutine: Coroutine[Any, Any, Any], timeout: float | None) -> Any:
-    """``coroutine``'s result, run on the event loop this process's sessions share, on a daemon thread of its own, for
-    a trainer that calls its environments from threads; waits at most ``timeout`` seconds."""
-    future = asyncio.run_coroutine_threadsafe(coroutine, _session_loop())
-    try:
-        return future.result(timeout)
-    except TimeoutError:
-        future.cancel()
-        raise
-
-
-def _session_loop() -> asyncio.AbstractEventLoop:
-    global _loop
-    with _loop_lock:
-        if _loop is None:
-            _loop = asyncio.new_event_loop()
-            threading.Thread(target=_loop.run_forever, name="agentenv-sessions", daemon=True).start()
-        return _loop
-
-
 def token_in(var: str | None) -> str | None:
     """The token in the environment variable ``var``, which must be set when it's named."""
     if var is None:

src/agent_env/training/_solver.py

@@ -0,0 +1,89 @@
+"""What training a bundle's own agents shares, the harness and AReaL's helpers: the tasks whose solver, the agent a step
+prompts, can train, prepared as ``agent-env run`` prepares them to run on local Docker, which alone reaches the
+trainer's endpoint on this machine, a remote sandbox needing a tunnel; and the overrides that point the solver alone at
+that endpoint, its other agents and judges keeping the configured one; and a finished run's scores, which a run
+``agent-env run`` counts as failed has none of."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+from agent_env.bundle.plan import solver_steps
+from agent_env.bundle.resolve import build_step
+from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params
+from agent_env.store.routing import namespace_routing
+from agent_env.task.task import Task, recorded_scores
+from agent_env.task_step.context import TaskStepContext, sandbox_overrides
+from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep
+from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep
+from agent_env.training.reward import reward_of
+
+logger = logging.getLogger(__name__)
+SANDBOX = "local"
+
+
+def trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str],
+            agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]:
+    """The selected tasks that can train their solver, prepared, each with its solver's deploys and their step params,
+    ``agent``'s if it swaps the solver. The others are left out before they're checked, so what they alone need can't
+    stop the rest."""
+    with namespace_routing():
+        kept = []
+        for entry in planned_tasks(bundle, tasks=tasks, evals=evals, agent=agent):
+            if why := _untrainable(entry.config):
+                logger.warning("%s is left out: %s", entry.entry.name, why)
+            else:
+                kept.append(entry.entry.name)
+        if not kept:
+            raise ValueError("no task of the bundle has a solver to train: an agent a step prompts, deployed by a "
+                             "deploy_agent step, that isn't also its judge")
+        prepared = prepare_bundle(bundle, tasks=kept, sandbox=SANDBOX, agent=agent, solver_endpoint=True)
+    swapped = solver_params(prepared.plan, prepared.materialization)
+    return [(task, {step: swapped.get(entry.id, {}).get(step, {})
+                    for step in solver_steps([step.to_dict() for step in task.steps])})
+            for entry, task in prepared.tasks]
+
+
+def solver_overrides(solver: dict[str, dict[str, Any]], endpoint: dict[str, str]) -> dict[str, Any]:
+    """A run's ``user_overrides``: every deploy on local Docker, and ``endpoint``, a ``litellm_base_url`` and
+    ``litellm_api_key``, on the solver's deploys alone."""
+    return {**sandbox_overrides(SANDBOX),
+            "step_params": {step: {**params, **endpoint} for step, params in solver.items()}}
+
+
+def rewarded(task: Task, reward: str | Callable[[dict[str, float]], float] | None
+             ) -> Callable[[dict[str, float]], float | None]:
+    of, why = reward_of(task, reward)
+    if of is None:
+        raise ValueError(why)
+    return of
+
+
+def run_scores(task: Task, context: TaskStepContext) -> dict[str, float]:
+    """The scores a finished run of ``task`` recorded. Raises when a step failed and no retry recovered it, a run
+    ``agent-env run`` counts as failed, though a step that doesn't fail the task let it finish."""
+    if failed := [failure for failure in context.metadata.get("failed_steps") or [] if not failure.get("retried")]:
+        raise RuntimeError("the run failed at " + "; ".join(f"{failure['step_id']}: {failure['error']}"
+                                                           for failure in failed))
+    return recorded_scores(task.steps, context)
+
+
+def first_prompt(task: Task) -> str:
+    return next((step.prompt for step in task.steps if step.type == "prompt_agent"), "")
+
+
+def _untrainable(config: list[dict[str, Any]]) -> str | None:
+    """Why a task's rollouts can't train its solver alone, or None."""
+    solver = set(solver_steps(config))
+    steps = [build_step(step) for step in config]
+    deploys = [step for step in steps if step.id in solver]
+    if not deploys or not all(isinstance(step, DeployAgentTaskStep) for step in deploys):
+        return "training needs the agent a step prompts deployed by a deploy_agent step"
+    agents = {step.agent_name for step in deploys}
+    if any(isinstance(step, RubricsVerifierTaskStep) and step.use_agent_judge and step.agent_name in agents
+           for step in steps):
+        return "its rubric judge is its solver, whose grading would be trained on"
+    return None

src/agent_env/training/areal.py

@@ -0,0 +1,102 @@
+"""Train a bundle's own agents with AReaL's agent workflows (``rollout.agent.mode: inline``): ``AgentEnvWorkflow`` is
+the agent AReaL's ``OpenAIProxyWorkflow`` runs for each sample, which plays the row's task as ``agent-env run`` does,
+its solver's model calls going through AReaL's proxy, which records the tokens AReaL trains on, and returns the reward
+the task's verifiers score. Needs no extra: AReaL calls it by its methods alone."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+from agent_env.store.routing import namespace_routing
+from agent_env.task.task import Task
+from agent_env.task.teardown import teardown_run
+from agent_env.task_step.context import TaskStepContext
+from agent_env.training._loop import session_loop
+from agent_env.training._solver import first_prompt, rewarded, run_scores, solver_overrides, trained
+
+#: The model the solver names: AReaL's proxy serves its policy whatever a request names, and its own agents name this.
+_MODEL = "default"
+
+
+def task_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (),
+              agent: str | None = None) -> list[dict[str, Any]]:
+    """The trainer's dataset, a row per task ``AgentEnvWorkflow`` plays with the same arguments: the task's id, which
+    is how the workflow finds it, and its first prompt as ``messages``, for AReaL's logs. The model never sees this
+    prompt: the task's own agent builds its own."""
+    return [{"task_id": task.id, "messages": [{"role": "user", "content": first_prompt(task)}]}
+            for task, _ in trained(bundle, tasks, evals, agent)]
+
+
+class AgentEnvWorkflow:
+    """The agent AReaL runs in inline mode, named as ``workflow: agent_env.training.areal.AgentEnvWorkflow`` with the
+    arguments as ``workflow_kwargs``. A process prepares the bundle to run on local Docker once, however often AReaL
+    builds the workflow; ``reward`` is as ``build_app``'s, and ``agent`` the agent the solver deploys, as
+    ``agent-env run --agent``'s."""
+
+    def __init__(self, bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (),
+                 reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None):
+        self._played = {task.id: (task, solver, rewarded(task, reward))
+                        for task, solver in _prepared(str(bundle), tuple(tasks), tuple(evals), agent)}
+
+    async def run(self, data: dict[str, Any], *, base_url: str, api_key: str, **_: Any) -> float:
+        """Plays ``data``'s task, its solver calling the model at ``base_url`` with ``api_key``, the key of the session
+        AReaL records the calls in, and returns its reward. Raises, which has AReaL leave the sample out, when the run
+        fails or records no score its reward reads. Cancelled, it stops the run and passes the cancel on once the run
+        is torn down."""
+        if (played := self._played.get(data.get("task_id"))) is None:
+            raise KeyError(f"the row names no task this workflow plays: {data.get('task_id')!r}; build the dataset "
+                           "with task_rows, from the same bundle, tasks and evals")
+        task, solver, reward = played
+        run = _Run(task, solver_overrides(solver, {"litellm_base_url": base_url, "litellm_api_key": api_key}))
+        done = asyncio.wrap_future(asyncio.run_coroutine_threadsafe(run.scores(), session_loop()))
+        try:
+            scores = await asyncio.shield(done)
+        except asyncio.CancelledError:
+            session_loop().call_soon_threadsafe(run.stop)
+            await asyncio.wait([done])
+            raise
+        if (score := reward(scores)) is None:
+            raise ValueError(f"task {task.id!r} records no score to train on")
+        return float(score)
+
+    def should_reject_failed_sample(self, error: Exception, disposition: str | None) -> bool:
+        """AReaL's question, after 2.1.0, for a sample whose ``run`` raised: it is left out of its group, the group's
+        other samples still training, unless AReaL found the model at fault, which it trains on as a zero reward."""
+        return disposition != "model_failure_zero"
+
+
+class _Run:
+    """One run of a task on the session loop, which ``stop``, called there, cuts short; its teardown always runs to the
+    end."""
+
+    def __init__(self, task: Task, overrides: dict[str, Any]):
+        self._task, self._overrides = task, overrides
+        self._running: asyncio.Task | None = None
+        self._stopped = False
+
+    async def scores(self) -> dict[str, float]:
+        context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides})
+        with namespace_routing():
+            try:
+                if self._stopped:
+                    raise asyncio.CancelledError
+                self._running = asyncio.ensure_future(self._task.run(agent_model=_MODEL, context=context))
+                await self._running
+                return run_scores(self._task, context)
+            finally:
+                await teardown_run(context)
+
+    def stop(self) -> None:
+        self._stopped = True
+        if self._running is not None:
+            self._running.cancel()
+
+
+@functools.cache
+def _prepared(bundle: str, tasks: tuple[str, ...], evals: tuple[str, ...],
+              agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]:
+    return trained(bundle, tasks, evals, agent)