feat(training): train on a bundle with Miles through an agent function over its OpenEnv server

.github/workflows/miles.yml

@@ -0,0 +1,78 @@
+name: Miles contract
+
+# agent_env.training.miles on Miles itself (tst/miles): Miles' rollout plays rows the writer wrote through its agentic
+# generate path, session server, router and mock SGLang server, scores them with the helper's reward hook and groups
+# them through its filter, on Miles v0.1.3 and on its main branch. Miles isn't on PyPI: as Miles' own CPU CI does, its
+# checkout, SGLang's and Megatron-LM's go on PYTHONPATH, SGLang and Megatron-LM at the commits a release's
+# release-lock.json names or at the branches Miles' main tests against, and Miles' requirements and its CPU CI's are
+# installed over the dev environment, torch the CPU build. It runs on pull requests that touch the helper, the OpenEnv
+# server it plays or the tier, weekly, and on demand; it isn't a required check.
+
+on:
+  pull_request:
+    paths:
+      - "src/agent_env/training/**"
+      - "src/agent_env/openenv/**"
+      - "tst/miles/**"
+      - "tst/unit/training/**"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/workflows/miles.yml"
+  schedule:
+    - cron: "17 8 * * 1"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: miles-${{ 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:
+  rollout:
+    name: rollout (Miles ${{ matrix.miles }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        miles: ["v0.1.3", "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 with the miles extra
+        run: |
+          python -m pip install --quiet -r .github/ci-requirements.txt
+          uv sync --locked --extra dev --extra miles --python "$(command -v python)"
+      - name: Check out Miles, SGLang and Megatron-LM
+        env:
+          MILES_REF: ${{ matrix.miles }}
+        run: |
+          checkout() { git init --quiet "$RUNNER_TEMP/$1" && git -C "$RUNNER_TEMP/$1" fetch --quiet --depth 1 "$2" "$3" \
+            && git -C "$RUNNER_TEMP/$1" checkout --quiet FETCH_HEAD && echo "$1 $3 $(git -C "$RUNNER_TEMP/$1" rev-parse HEAD)"; }
+          checkout miles https://github.com/radixark/miles "$MILES_REF"
+          sglang=sglang-miles megatron=miles-main
+          if [ -f "$RUNNER_TEMP/miles/release-lock.json" ]; then
+            sglang=$(jq -er .sglang_commit "$RUNNER_TEMP/miles/release-lock.json")
+            megatron=$(jq -er .megatron_commit "$RUNNER_TEMP/miles/release-lock.json")
+          fi
+          checkout sglang https://github.com/sgl-project/sglang "$sglang"
+          checkout Megatron-LM https://github.com/radixark/Megatron-LM "$megatron"
+          echo "PYTHONPATH=$RUNNER_TEMP/miles:$RUNNER_TEMP/sglang/python:$RUNNER_TEMP/Megatron-LM" >> "$GITHUB_ENV"
+      - name: Install Miles' requirements and CPU torch
+        run: |
+          uv pip install --python .venv/bin/python -r "$RUNNER_TEMP/miles/requirements.txt" \
+            -r "$RUNNER_TEMP/miles/tests/ci/requirements-ci-cpu.txt" \
+            --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match
+          .venv/bin/python -c "from importlib.metadata import version; print('torch', version('torch'), 'transformers', version('transformers'))"
+      - name: Miles tier
+        run: .venv/bin/python -m pytest tst/miles -rA --timeout=900 --junitxml=junit-miles.xml
+      - name: No skips
+        if: always()
+        run: .venv/bin/python .github/scripts/check_skip_policy.py junit-miles.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 miles-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 Miles tier plays rows through Miles' own rollout: its agentic generate path, session server, router and mock
+# SGLang server, with Qwen3's tokenizer from the Hub. Miles isn't on PyPI: put checkouts of Miles, SGLang and Megatron-LM
+# on PYTHONPATH and install Miles' requirements and CPU torch first, as .github/workflows/miles.yml does.
+miles-test: ## Play rows through Miles' rollout and session server (needs Miles, SGLang and Megatron-LM on PYTHONPATH, CPU torch, and the Hub)
+	$(PYTHON) -m pytest tst/miles --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)

pyproject.toml

@@ -91,6 +91,11 @@ skyrl = [
     "pyarrow>=19",
     "transformers>=5",
 ]
+# agent_env.training.miles: Miles' trainers on a bundle served as an OpenEnv env; installed into Miles' own environment,
+# whose Miles, SGLang and torch are the trainer's to pin. It needs nothing beyond the openenv extra.
+miles = [
+    "agentenv-framework[openenv]",
+]
 # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml.
 aws = [
     "boto3>=1.35.0",

src/agent_env/training/__init__.py

@@ -1,4 +1,4 @@
 """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."""
+environments, NeMo RL among them, ``agent_env.training.verl`` for verl, ``agent_env.training.skyrl`` for SkyRL, and
+``agent_env.training.miles`` for Miles."""

src/agent_env/training/miles/__init__.py

@@ -0,0 +1,5 @@
+"""Miles' trainers (https://github.com/radixark/miles) on a bundle served as an OpenEnv env (``agent_env.openenv.build_app``,
+or the package ``openenv import`` writes): ``python -m agent_env.training.miles URL OUT`` writes the server's tasks as
+Miles' prompt data, and ``agent_env.training.miles.agent`` holds the agent function Miles calls for each sample, which
+plays the row's task as one episode on the server through Miles' session server, with the reward hook and the
+dynamic-sampling filter beside it."""

src/agent_env/training/miles/__main__.py

@@ -0,0 +1,33 @@
+"""``python -m agent_env.training.miles URL OUT``: writes a bundle's tasks as Miles' prompt data."""
+
+import click
+
+from agent_env.training.miles.prepare import write_rows
+
+
+@click.command()
+@click.argument("url")
+@click.argument("out", type=click.Path(file_okay=False))
+@click.option("--tool-timeout", type=float, default=60.0, show_default=True,
+              help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.")
+@click.option("--reset-timeout", type=float, default=600.0, show_default=True,
+              help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a "
+                   "rollout waits for room on the server.")
+@click.option("--model-timeout", type=float, default=600.0, show_default=True,
+              help="How long the agent function waits for each of the model's replies, in seconds; an episode whose "
+                   "model misses it fails, and its group is left out.")
+@click.option("--max-turns", type=int, default=4, show_default=True,
+              help="The model turns a rollout may take; the episode is ended after the last.")
+@click.option("--token-var", help="The environment variable that holds the token the server asks for, read now and in "
+                                  "each of Miles' rollout processes.")
+def main(url: str, out: str, tool_timeout: float, reset_timeout: float, model_timeout: float, max_turns: int,
+         token_var: str | None) -> None:
+    """Writes the tasks of the OpenEnv server at URL as OUT/<split>.jsonl, a row per task, for Miles' --prompt-data
+    with --input-key messages and --metadata-key metadata."""
+    folder = write_rows(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, model_timeout=model_timeout,
+                        max_turns=max_turns, token_var=token_var)
+    click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}")
+
+
+if __name__ == "__main__":
+    main()

src/agent_env/training/miles/agent.py

@@ -0,0 +1,113 @@
+"""What Miles calls for each sample of the rows ``python -m agent_env.training.miles`` writes: the agent function
+(``--custom-agent-function-path agent_env.training.miles.agent.run``), the reward hook
+(``--custom-rm-path agent_env.training.miles.agent.reward``) and the dynamic-sampling filter
+(``--dynamic-sampling-filter-path agent_env.training.miles.agent.keep_scored_groups``). Imports nothing of Miles: it
+calls them by their arguments."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from typing import Any
+
+import httpx
+
+from agent_env.training._openenv import Session, Shared, token_in
+
+
+class _Episode(Session):
+    """One sample's episode on the server its row names."""
+
+    offered_by = "its row lists"
+
+    def __init__(self, shared: Shared):
+        super().__init__()
+        self._shared = shared
+
+
+async def run(base_url: str, prompt: list[dict[str, Any]], request_kwargs: dict[str, Any] | None = None,
+              metadata: dict[str, Any] | None = None, **kwargs: Any) -> dict[str, Any]:
+    """Plays the row's task as one episode: the model chats at ``base_url``, the sample's session, which Miles records,
+    with the row's tools; each call is a call on the episode, and the model's last reply, or the turn cap, ends the
+    turn. A call whose arguments aren't a JSON object is answered with why, as the model's error, and the episode goes
+    on. Returns what Miles merges into the sample's metadata: the episode's reward and scores, or, when the env or the
+    session failed it, the model missed the row's ``model_timeout`` or the episode ended without a reward, a reward of 0
+    and why, which ``keep_scored_groups`` leaves out of training; and the calls made and refused. A failed episode
+    isn't played again: the session already holds its turns."""
+    row = metadata or {}
+    server, tools = row["agentenv"], row["tools"]
+    episode = _Episode(Shared(server["url"], frozenset(tool["function"]["name"] for tool in tools),
+                              server["tool_timeout"], server["reset_timeout"], patience=server["reset_timeout"],
+                              waits_for_server=True, token=token_in(server.get("token_var"))))
+    counts = {"agentenv_tool_calls": 0, "agentenv_invalid_calls": 0}
+    try:
+        await episode._start(row["split"], int(row["index"]), row["task_id"])
+        reply = await _converse(episode, base_url, list(prompt), tools, request_kwargs or {}, server, counts)
+        rollout = episode._rollout
+        if not (rollout.done or rollout.failure):
+            await episode._finish(rollout, reply)
+        if rollout.failure is None and rollout.reward is None:
+            rollout.failure = "the episode ended without a reward"
+        if rollout.failure is not None:
+            return {"reward": 0.0, "agentenv_failure": rollout.failure, **counts}
+        return {"reward": float(rollout.reward), "agentenv_scores": rollout.scores or {}, **counts}
+    except Exception as e:
+        return {"reward": 0.0, "agentenv_failure": f"{type(e).__name__}: {e}", **counts}
+    finally:
+        await episode._close()
+
+
+async def reward(args: Any, samples: Any, **kwargs: Any) -> float | list[float]:
+    """Each sample's episode reward, which ``run`` put in its metadata."""
+    if isinstance(samples, list):
+        return [float(sample.metadata.get("reward", 0.0)) for sample in samples]
+    return float(samples.metadata.get("reward", 0.0))
+
+
+def keep_scored_groups(args: Any, samples: list[Any], **kwargs: Any) -> bool:
+    """Whether to train on a group: not when an episode of it failed in the env, its agent function never returned
+    (a child process that died mid-episode leaves the session's turns with no reward), or Miles aborted a sample of
+    it, none of them the policy's doing."""
+    return not any(sample.metadata.get("agentenv_failure") or "reward" not in sample.metadata
+                   or getattr(sample.status, "name", None) == "ABORTED" for sample in _samples(samples))
+
+
+async def _converse(episode: _Episode, base_url: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]],
+                    request_kwargs: dict[str, Any], server: dict[str, Any], counts: dict[str, int]) -> str:
+    """The model's turns, until it replies without a call, the episode ends, or the server's ``max_turns`` pass; its
+    last reply."""
+    async with httpx.AsyncClient(timeout=server["model_timeout"]) as client:
+        for _ in range(server["max_turns"]):
+            response = await client.post(f"{base_url}/v1/chat/completions", json={
+                "model": "default", "messages": messages, "tools": tools, **request_kwargs})
+            response.raise_for_status()
+            message = response.json()["choices"][0]["message"]
+            # Verbatim: the session server matches each message a request replays against the one it recorded.
+            messages.append(message)
+            if not (calls := message.get("tool_calls")):
+                return message.get("content") or ""
+            for call in calls:
+                messages.append({"role": "tool", "tool_call_id": call.get("id", ""),
+                                 "content": await _answer(episode, call, counts)})
+            if episode._rollout.done or episode._rollout.failure:
+                return ""
+    return ""
+
+
+async def _answer(episode: _Episode, call: dict[str, Any], counts: dict[str, int]) -> str:
+    name, arguments = call["function"]["name"], call["function"].get("arguments")
+    try:
+        arguments = json.loads(arguments) if isinstance(arguments, str) else arguments
+    except json.JSONDecodeError as e:
+        counts["agentenv_invalid_calls"] += 1
+        return f"{name} wasn't called: its arguments aren't JSON ({e})"
+    if not isinstance(arguments, dict):
+        counts["agentenv_invalid_calls"] += 1
+        return f"{name} wasn't called: its arguments must be a JSON object, not {type(arguments).__name__}"
+    counts["agentenv_tool_calls"] += 1
+    return await episode._step(name, arguments)
+
+
+def _samples(group: list[Any]) -> Iterator[Any]:
+    for sample in group:
+        yield from sample if isinstance(sample, list) else [sample]

src/agent_env/training/miles/prepare.py

@@ -0,0 +1,33 @@
+"""The server's tasks as Miles' prompt data: a JSON line per task of each split, its chat ``messages`` and the
+``metadata`` the agent function plays it from."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from agent_env.training._openenv import function_tool, key, messages, tasks_and_tools, token_in
+
+
+def write_rows(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0,
+               model_timeout: float = 600.0, max_turns: int = 4, token_var: str | None = None) -> Path:
+    """The tasks of the OpenEnv server at ``url`` as ``out/<split>.jsonl``, a row per task of each split: its system
+    prompt and first prompt as ``messages``, and as ``metadata`` its split, index and task id, the tools its toolset
+    offers, and where and how the agent function plays it: the server, its deadlines and the model's in seconds, the
+    model turns a rollout may take, and the environment variable holding the token a server given one asks for, read in each of
+    Miles' rollout processes."""
+    url = url.rstrip("/")
+    _, tasks, probed = tasks_and_tools(url, local="serve the bundle with agent_env.openenv.build_app first",
+                                       timeout=reset_timeout, token=token_in(token_var))
+    tools = {each: [function_tool(tool) for tool in found] for each, found in probed.items()}
+    server = {"url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, "model_timeout": model_timeout,
+              "max_turns": max_turns, **({"token_var": token_var} if token_var else {})}
+    out = Path(out)
+    out.mkdir(parents=True, exist_ok=True)
+    for split, described in tasks.items():
+        (out / f"{split}.jsonl").write_text("".join(json.dumps({
+            "messages": messages(task["system_prompt"], task["prompt"]),
+            "metadata": {"split": split, "index": task["index"], "task_id": task["id"], "tools": tools[key(task)],
+                         "agentenv": server},
+        }) + "\n" for task in described))
+    return out

tst/miles/conftest.py

@@ -0,0 +1,7 @@
+"""The unit tier's bundle served on loopback over the Episode API's fake steps, which these tests play through Miles'
+own rollout."""
+
+from tst.unit.conftest import local_stores  # noqa: F401
+from tst.unit.openenv.environment_test import tools  # noqa: F401
+from tst.unit.task.episode.run_test import teardowns  # noqa: F401
+from tst.unit.training._support import serve  # noqa: F401