feat(training): train a bundle's own agents with TRL's harness rollouts

.github/workflows/harness.yml

@@ -0,0 +1,55 @@
+name: Harness contract
+
+# TRL's harness rollouts on a bundle's own agents (agent_env.training.harness, tst/unit/training/harness_test.py)
+# through OpenEnv's capture proxy, on the uv.lock pins and on the latest OpenEnv. The harness reads the proxy's internals
+# (its session registry, export and training trace), which OpenEnv can change in any minor release. It runs on pull
+# requests that touch the harness or its tier, weekly, and on demand; it isn't a required check.
+
+on:
+  pull_request:
+    paths:
+      - "src/agent_env/training/**"
+      - "tst/unit/training/**"
+      - "tst/util/fake_vllm.py"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/workflows/harness.yml"
+  schedule:
+    - cron: "53 7 * * 1"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: harness-${{ 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:
+  capture:
+    name: capture (openenv ${{ matrix.openenv }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        openenv: ["locked", "latest"]
+    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, with the latest OpenEnv for latest
+        run: |
+          python -m pip install --quiet -r .github/ci-requirements.txt
+          uv sync --locked --extra dev --python "$(command -v python)"
+          if [ "${{ matrix.openenv }}" = latest ]; then uv pip install --python .venv/bin/python --upgrade openenv; fi
+          .venv/bin/python -c "import importlib.metadata as m; print('openenv', m.version('openenv'))"
+      - name: Harness tier
+        run: .venv/bin/python -m pytest tst/unit/training/harness_test.py -rA --timeout=600 --junitxml=junit-harness.xml
+      - name: No skips
+        if: always()
+        run: .venv/bin/python .github/scripts/check_skip_policy.py junit-harness.xml

pyproject.toml

@@ -55,6 +55,12 @@ dependencies = [
 openenv = [
     "openenv>=0.9",
 ]
+# agent_env.training.harness: TRL's harness rollouts on a bundle's own agents, through OpenEnv's capture proxy (0.10
+# and later). TRL and vLLM are the trainer's to pin.
+harness = [
+    "agentenv-framework[openenv]",
+    "openenv>=0.10",
+]
 # agent_env.training.trl: TRL's trainers on a bundle served as an OpenEnv env. TRL itself is the trainer's to pin.
 trl = [
     "agentenv-framework[openenv]",
@@ -133,7 +139,7 @@ coreweave = [
     "cwsandbox>=1.19.6,<2",  # Reject exec streams that close without an exit status.
 ]
 dev = [
-    "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]",  # exercise portable optional providers and stores in CI
+    "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,harness,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]",  # exercise portable optional providers and stores in CI
     "fastmcp>=3.4,<4",  # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x
     "Pillow>=10.4",  # the protocol's computer extra, whose tests the unit tier runs
     "av>=17",  # the protocol's computer extra, whose tests the unit tier runs and which the recording tests decode with

src/agent_env/bundle/preflight.py

@@ -80,7 +80,7 @@
 from agent_env.utils.docker.build import DEFAULT_BUILD_PLATFORM, docker_plugin_missing, docker_unreachable
 
 from .parse import BundleError, BundleKind
-from .plan import AgentOverride, Plan
+from .plan import AgentOverride, Plan, solver_steps
 from .resolve import BuiltImage, Reference, ReferencedImage, ResolvedEntry, build_step, played_deploys
 
 _SHOWN = 3
@@ -96,13 +96,16 @@ class Preflight:
     built_here: frozenset[str] = frozenset()  # the ids of the images the bundle builds that a deploy here runs
 
 
-def preflight_run(plan: Plan, tasks: Iterable[ResolvedEntry], sandbox: str | None, *, played: bool = False) -> Preflight:
+def preflight_run(plan: Plan, tasks: Iterable[ResolvedEntry], sandbox: str | None, *, played: bool = False,
+                  solver_endpoint: bool = False) -> Preflight:
     """Check where each of ``tasks`` deploys, with ``sandbox`` overriding every deploy's provider, and name the infra
-    envs the run builds; ``played`` skips each task's played agent's deploys (``played_deploys``). Raises BundleError
-    listing every problem found. Reads the store but writes nothing."""
+    envs the run builds; ``played`` skips each task's played agent's deploys (``played_deploys``), and with
+    ``solver_endpoint`` the caller gives the solver's deploys (``solver_steps``) their model endpoint. Raises
+    BundleError listing every problem found. Reads the store but writes nothing."""
     walk = _Walk(plan, sandbox)
     for task in tasks:
-        walk.task(task, skip=played_deploys(task) if played else frozenset())
+        walk.task(task, skip=played_deploys(task) if played else frozenset(),
+                  given=frozenset(solver_steps(task.config)) if solver_endpoint else frozenset())
     return walk.finish()
 
 
@@ -148,7 +151,7 @@ def __init__(self, plan: Plan, sandbox: str | None):
         # Each deploy of an image the bundle builds: (where, what, its provider's links, the images it accepts on each)
         self.uses: dict[str, list[tuple[str, str, list[SandboxProvider], Callable[[SandboxProvider], Accepts]]]] = {}
 
-    def task(self, task: ResolvedEntry, skip: frozenset[str]) -> None:
+    def task(self, task: ResolvedEntry, skip: frozenset[str], given: frozenset[str] = frozenset()) -> None:
         steps = [build_step(config) for config in task.config]
         sandboxes = {step.sandbox_name: step for step in steps if isinstance(step, DeploySandboxTaskStep)}
         solver = self.plan.agent.steps.get(task.entry.id, ()) if self.plan.agent else ()
@@ -161,7 +164,7 @@ def task(self, task: ResolvedEntry, skip: frozenset[str]) -> None:
                 if isinstance(step, DeployEnvTaskStep):
                     self._env(where, step)
                 elif isinstance(step, DeployAgentTaskStep):
-                    self._agent(where, step, sandboxes, override)
+                    self._agent(where, step, sandboxes, override, endpoint_given=step.id in given)
                 elif isinstance(step, DeploySandboxTaskStep):
                     self._sandbox(where, step)
                 elif isinstance(step, RubricsVerifierTaskStep):
@@ -289,7 +292,7 @@ def _planned_version(self, kind: EntityKind, entity_id: str, version: int | None
         return version if version is not None else self.plan.store_latest.get((kind, entity_id))
 
     def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, DeploySandboxTaskStep],
-               override: AgentOverride | None = None) -> None:
+               override: AgentOverride | None = None, endpoint_given: bool = False) -> None:
         """Check a deploy of the step's agent, or of ``override``, the agent ``--agent`` names for the solver."""
         accepts: Callable[[SandboxProvider], Accepts] = A2AAgent.accepts
         if step.sandbox_name:
@@ -303,7 +306,8 @@ def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, Dep
             accepts = functools.partial(A2AAgent.accepts, sandbox_mode=linked.sandbox_mode)  # the VM loads its image
         else:
             provider = _provider(self.sandbox or step.sandbox_type, get_agent_sandbox_provider)
-        given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else [])})
+        given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else []),
+                           *(_MODEL_VARS if endpoint_given else ())})
         agent_id, version = ((override.agent_id, override.version) if override is not None
                              else (step.a2a_agent_id, step.a2a_agent_version))
         self._agent_deploy(where, provider, agent_id, version, "names no agent", accepts, given=given)

src/agent_env/bundle/run.py

@@ -41,7 +41,8 @@
 from .plan import AliasUse, Plan, Write, plan_bundle
 from .preflight import Preflight, preflight_run
 from .requires import Declaring, declare, declaring_for
-from .resolve import BuiltImage, BundleImage, Reference, ReferencedImage, resolve_bundle, without_played_agents
+from .resolve import (BuiltImage, BundleImage, Reference, ReferencedImage, ResolvedEntry, resolve_bundle,
+                      without_played_agents)
 
 logger = logging.getLogger(__name__)
 
@@ -251,7 +252,7 @@ def run_bundle(
         with Interrupts() as interrupts:
             overrides = _overrides(sandbox, user_model, judge_model)
             runs = interrupts.run(_run_all(plan, list(prepared.tasks), model, overrides,
-                                           _solver_params(plan, prepared.materialization), keep, say, interrupts))
+                                           solver_params(plan, prepared.materialization), keep, say, interrupts))
             _mark_cancelled(runs, interrupts.reason)
             result = _bundle_run(plan, prepared.materialization, runs,
                                  _skipped(plan, entries, every=not tasks and not evals))
@@ -274,17 +275,20 @@ def prepare_bundle(
     on_progress: Callable[[str], None] | None = None,
     id_root: str | None = None,
     played: bool = False,
+    agent: str | None = None,
+    solver_endpoint: bool = False,
 ) -> PreparedBundle:
     """Do what ``run_bundle`` with these arguments does before its first task runs, and stop there. ``played``
     leaves out what each task's only prompted agent would deploy, for a caller that plays that agent itself
-    (``EpisodeSpec``), so the agent needn't exist.
+    (``EpisodeSpec``), so the agent needn't exist. With ``solver_endpoint`` the caller gives the solver's deploys their
+    model endpoint in each run's ``step_params``, so the config needn't set one.
 
     It doesn't turn namespace routing on: the caller keeps it on for as long as it runs the tasks
     (``agent_env.store.routing.namespace_routing``), which suits a service that runs nothing else. Raises what
     ``run_bundle`` raises before its first task runs, except for a running event loop, which it doesn't need."""
     composing = _composing(sandbox)
-    return _prepare(root, tasks, evals, sandbox, id_root, None, played, _progress(on_progress), _hint(id_root),
-                    composing)
+    return _prepare(root, tasks, evals, sandbox, id_root, agent, played, _progress(on_progress), _hint(id_root),
+                    composing, solver_endpoint=solver_endpoint)
 
 
 def _write_bundle(
@@ -373,24 +377,26 @@ def _bundle_run(plan: Plan, materialization: Materialization, runs: tuple[TaskRu
 
 def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandbox: str | None, id_root: str | None,
              agent: str | None, played: bool, say: Callable[[str], None], declaring: Declaring,
-             composing: Composing = Composing()) -> PreparedBundle:
-    plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played)
+             composing: Composing = Composing(), *, solver_endpoint: bool = False) -> PreparedBundle:
+    plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played,
+                               solver_endpoint=solver_endpoint)
     _say_plan(plan, say, declaring)
     materialization = _written_with_infra(plan, preflight, say, composing)
     to_run = tuple((entry, Task.get(entry.id, materialization.version_of("task", entry.id))) for entry in _to_run(plan))
     return PreparedBundle(plan, materialization, to_run)
 
 
 def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_root: str | None,
-             sandbox: str | None, agent: str | None = None, *, played: bool = False) -> tuple[Plan, Preflight]:
+             sandbox: str | None, agent: str | None = None, *, played: bool = False,
+             solver_endpoint: bool = False) -> tuple[Plan, Preflight]:
     resolved = resolve_bundle(parse_bundle(Path(root), id_root=id_root))
     plan = plan_bundle(without_played_agents(resolved) if played else resolved, tasks=tasks, evals=evals,
                        installed=installed_bundles, agent=agent)
     refuse_store_tasks(plan)
     to_run = set(_to_run(plan))
     return plan, preflight_run(plan, [write.source for write in plan.writes
                                       if write.kind is BundleKind.TASK and write.source.entry in to_run], sandbox,
-                               played=played)
+                               played=played, solver_endpoint=solver_endpoint)
 
 
 def _composing(sandbox: str | None, emulate: bool = True) -> Composing:
@@ -399,6 +405,16 @@ def _composing(sandbox: str | None, emulate: bool = True) -> Composing:
     return Composing(build_sandbox_provider(sandbox) if sandbox else None, emulate)
 
 
+def planned_tasks(root: Path | str, *, tasks: Sequence[str] = (), evals: Sequence[str] = (),
+                  agent: str | None = None) -> tuple[ResolvedEntry, ...]:
+    """The tasks ``run_bundle`` with these arguments would run, as resolved, without checking or writing anything.
+    Namespace routing must be on, as for ``plan_bundle``."""
+    plan = plan_bundle(resolve_bundle(parse_bundle(Path(root))), tasks=tasks, evals=evals, installed=installed_bundles,
+                       agent=agent)
+    selected = set(_to_run(plan))
+    return tuple(write.source for write in plan.writes if write.kind is BundleKind.TASK and write.source.entry in selected)
+
+
 def _to_run(plan: Plan) -> tuple[BundleEntry, ...]:
     """The tasks selected on their own and every task a selected eval names, each once, in the plan's order."""
     wanted = {entry.entry.id for entry in plan.tasks}
@@ -445,7 +461,7 @@ def _hint(id_root: str | None) -> Declaring:
     return Declaring.QUIET if id_root is not None else Declaring.HINT
 
 
-def _solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]:
+def solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]:
     """Each task's ``step_params`` that deploy ``Plan.agent`` at its solver steps: at the version this run wrote, for
     an agent it writes, or the store's version the plan read."""
     if plan.agent is None:

src/agent_env/openenv/app.py

@@ -23,6 +23,7 @@
 from agent_env.task import EpisodeError, EpisodeSpec, Task
 
 from .environment import AgentEnvEnvironment, Served, ServedTask
+from .reward import reward_of
 
 logger = logging.getLogger(__name__)
 
@@ -142,11 +143,11 @@ def _served(splits: dict[str, list[tuple[str, Task]]],
             except EpisodeError as e:
                 refused[task.id] = str(e)
                 continue
-            reward_of, problem = _reward_of(task, reward)
+            rewarded, problem = reward_of(task, reward)
             if problem:
                 unscorable[task.id] = problem
                 continue
-            kept.append(ServedTask(task, name, spec, reward_of))
+            kept.append(ServedTask(task, name, spec, rewarded))
         served[split] = tuple(kept)
     for why in refused.values():
         logger.warning("left out: %s", why)
@@ -156,32 +157,3 @@ def _served(splits: dict[str, list[tuple[str, Task]]],
         raise BundleError([f"no task can be served as an OpenEnv env: {why}" for why in refused.values()]
                           or ["no task can be served as an OpenEnv env: the bundle has none"])
     return served
-
-
-def _reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None
-               ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]:
-    """``task``'s reward as a function of its scores, or why ``reward`` can't make one."""
-    if callable(reward):
-        return reward, None
-    scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)]
-    if reward is None:
-        if len(scoring) > 1:
-            return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is "
-                          "its reward with reward=<step id>")
-        return (_score_at(scoring[0]) if scoring else _unscored), None
-    if reward not in scoring:
-        return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}"
-    return _score_at(reward), None
-
-
-def _score_at(step: str) -> Callable[[dict[str, float]], float]:
-    def score(scores: dict[str, float]) -> float:
-        if step not in scores:
-            raise LookupError(f"the run recorded no score at step {step!r}")
-        return scores[step]
-
-    return score
-
-
-def _unscored(scores: dict[str, float]) -> None:
-    return None

src/agent_env/openenv/reward.py

@@ -0,0 +1,36 @@
+"""How a task's recorded scores make its reward: the rules ``build_app`` applies, which harness training shares."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from agent_env.task import Task
+
+
+def reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None
+               ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]:
+    """``task``'s reward as a function of its scores, or why ``reward`` can't make one."""
+    if callable(reward):
+        return reward, None
+    scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)]
+    if reward is None:
+        if len(scoring) > 1:
+            return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is "
+                          "its reward with reward=<step id>")
+        return (_score_at(scoring[0]) if scoring else _unscored), None
+    if reward not in scoring:
+        return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}"
+    return _score_at(reward), None
+
+
+def _score_at(step: str) -> Callable[[dict[str, float]], float]:
+    def score(scores: dict[str, float]) -> float:
+        if step not in scores:
+            raise LookupError(f"the run recorded no score at step {step!r}")
+        return scores[step]
+
+    return score
+
+
+def _unscored(scores: dict[str, float]) -> None:
+    return None