From 50e407a65b8c07433f6297715eb7aaa8f9df0d86 Mon Sep 17 00:00:00 2001 From: Tejas Polakam Date: Mon, 2 Nov 2026 19:20:40 -0500 Subject: [PATCH] feat(training): train with Tinker on a bundle's tasks `agent_env.training.tinker.AgentEnvDatasetBuilder` is tinker-cookbook's `rl.train` dataset builder: each batch is groups of episodes of one task, and every rollout is an episode played in process on the Episode API, not through the OpenEnv server, since Tinker samples itself. The turn's tools are rendered by the cookbook's renderer, each sampled message is recorded on the turn before its calls run, a reply without a call ends the turn, and the reward is made by the rules `build_app` applies (`openenv/reward.py`). Tool-call arguments are typed by the tool's schema, since a renderer that writes parameters as untyped text reads digits as numbers. A rollout the cookbook cuts is scored once its group is done; a failed episode raises for the cookbook's retry. Each group is filed in the run index as `/train|eval/step-` (`_episodes.RunGroup`). The `tinker` extra is `openenv`; tinker-cookbook, Tinker and torch are the trainer's to pin. The `tst/tinker` tier plays episodes on tinker-cookbook with a stand-in renderer, run by `tinker.yml` and `make tinker-test`. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/tinker.yml | 65 +++ AGENTS.md | 2 +- Makefile | 8 +- pyproject.toml | 7 +- src/agent_env/training/_episodes.py | 76 +++ src/agent_env/training/tinker.py | 443 ++++++++++++++++++ tst/tinker/__init__.py | 0 tst/tinker/conftest.py | 6 + tst/tinker/episode_env_test.py | 133 ++++++ .../test_optional_extras_stay_optional.py | 2 +- tst/unit/training/episodes_test.py | 44 ++ tst/util/capabilities.py | 4 + uv.lock | 8 +- 13 files changed, 792 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/tinker.yml create mode 100644 src/agent_env/training/_episodes.py create mode 100644 src/agent_env/training/tinker.py create mode 100644 tst/tinker/__init__.py create mode 100644 tst/tinker/conftest.py create mode 100644 tst/tinker/episode_env_test.py create mode 100644 tst/unit/training/episodes_test.py diff --git a/.github/workflows/tinker.yml b/.github/workflows/tinker.yml new file mode 100644 index 000000000..d6babb8e8 --- /dev/null +++ b/.github/workflows/tinker.yml @@ -0,0 +1,65 @@ +name: Tinker contract + +# The Tinker helper's episodes on tinker-cookbook itself (tst/tinker): EpisodeEnv played as the cookbook's rollouts step +# a MessageEnv, on the Episode API's fake steps with a stand-in renderer, so nothing samples and no Tinker key is needed, +# on tinker-cookbook as released and on its main branch. tinker-cookbook and CPU torch are installed over the dev +# environment. It runs on pull requests that touch the helper, the Episode API or the tier, weekly, and on demand; it +# isn't a required check. + +on: + pull_request: + paths: + - "src/agent_env/training/**" + - "src/agent_env/task/episode/**" + - "src/agent_env/openenv/reward.py" + - "tst/tinker/**" + - "tst/unit/training/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/tinker.yml" + schedule: + - cron: "17 8 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: tinker-${{ 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: + episodes: + name: episodes (tinker-cookbook ${{ matrix.cookbook }}) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + include: + - cookbook: "released" + requirement: "tinker-cookbook" + - cookbook: "main" + requirement: "tinker-cookbook@git+https://github.com/thinking-machines-lab/tinker-cookbook@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 tinker extra + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --extra tinker --python "$(command -v python)" + - name: Install tinker-cookbook and CPU torch + run: | + uv pip install --python .venv/bin/python ${{ matrix.requirement }} 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('tinker-cookbook', version('tinker-cookbook'), 'tinker', version('tinker'), 'torch', version('torch'))" + - name: Tinker tier + run: .venv/bin/python -m pytest tst/tinker -rA --timeout=900 --junitxml=junit-tinker.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-tinker.xml diff --git a/AGENTS.md b/AGENTS.md index b2a31d71d..bd6243cd7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -123,7 +123,7 @@ and deserialized through a registry. | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap; this package and `training/` are the only ones that import it. Each task's description in the Task API names its `toolset` (the played agent's envs, role, and whether it gets `bash`). With `token`, every request and WebSocket but `/health` must carry `Authorization: Bearer `, for a server others reach. | -| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_openenv.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. | +| `training/` | One module per outside trainer, each behind its own extra; `_openenv.py`, the session on a server `openenv/` makes that each helper plays an episode through (`Session`: a background start, tool calls, the turns left open ended until scored), with the Task API reads, the per-toolset probe, `tasks_and_tools` (what the row writers read: the tasks, the split names checked, each toolset's tools) and an in-process `Server`; and `_tools.py`, which turns an env's tools into typed, documented Python methods from their JSON Schema, as transformers' `get_json_schema` reads them. `trl.py` (`trl` extra): `AgentEnv(url)` over a server `openenv/` makes, with `dataset` (a row per task) and `factory` (TRL's `environment_factory`: a generated class per toolset whose methods are the tools, a background `reset` and an owned `get_reward`); `AgentEnv.local(bundle)` serves the bundle in process. `nemo_gym/` (`nemo-gym` extra): `python -m agent_env.training.nemo_gym SOURCE ROOT` writes a bundle's NeMo Gym resources server into a Gym root, `NEMO_GYM_EXTRA_ROOTS`'s or a Gym checkout's (`app.py`; a `pyproject.toml`, since uv applies the nearest one's `[tool.uv]` and a checkout's `exclude-dependencies` would drop `cryptography`; a config with `simple_agent` and the `legacy_agent` environment server, and `data/.jsonl` rows with each task's tools); `server.AgentEnvResourcesServer` is that server, run in the venv Gym builds: a row's seed starts its episode, a catch-all route calls its tools, and `verify` ends the turns left open with the model's last reply and returns the reward, masking a rollout the environment failed; with `reward_components` it also returns each score by step id, their sum the reward, for NeMo RL's GDPO. `verl/` (`verl` extra, installed into verl's own environment): `python -m agent_env.training.verl URL OUT` writes a server's tasks as verl's `.parquet` rows (the task's prompt, `agent_name: agentenv`, the row's tools in `extra_info`) and `agent_loop.yaml`; `loop.AgentEnvAgentLoop` is verl's `ToolAgentLoop` with the row's tools, each call a call on the row's episode, the last reply ending the open turn and the episode's reward its `reward_score`, playing a rollout the env failed again on a new episode up to `attempts` times within `reset_timeout` of its start and raising `RolloutFailed` for one that fails every time, which verl's V1 trainer (the default) leaves out of the loss, the group's baseline and the eval reward; it hands verl's parsers typed tool schemas (`qwen3_coder` converts arguments by them), raises on a first prompt over `rollout.prompt_length`, which verl would cut from the left, and on an episode that ends unscored without failing, and warns once a tool of an answer verl cuts to `max_tool_response_length`. The writer reads a server that asks for a token with the one `--token-var` names, which the loop reads in each worker, and with `--tokenizer` counts each prompt with its tools as verl renders it (`--chat-template-kwargs`), refusing rows over `--max-prompt-length` before writing. `skyrl/` (`skyrl` extra, installed into SkyRL's own environment): `python -m agent_env.training.skyrl URL OUT --tokenizer MODEL` writes a server's tasks as SkyRL's `.parquet` rows (`env_class: agentenv`, the episode's split, index, task id and tools, the server's settings), each with its tools rendered into its system prompt by the model's chat template, checked to render as `tools=` would, since SkyRL's generator passes the template no tools, and refuses rows over `--max-prompt-length` counted as SkyRL's filter and generator count them; `calls.parse` reads a reply's tool calls in the `hermes` or `qwen3_coder` format, typed by the row's schema; `env.AgentEnvTextEnv` (registered as `agentenv`) is skyrl-gym's `BaseTextEnv` on the row's episode, run on the event loop the sessions share (`_openenv.on_session_loop`), its steps running the reply's calls and a reply without one ending the open turn, raising `EpisodeFailed` when the env fails it, with scores, tool calls and unparsed calls as its metrics; `generator.AgentEnvGenerator` is SkyRL's `SkyRLGymGenerator` playing a failed rollout again on a new episode up to `attempts` times within `reset_timeout`, and masking the group of one that fails every attempt in training (an eval scores it 0); `python -m agent_env.training.skyrl.train` (or `trainer.override_entrypoint=agent_env.training.skyrl.train`) runs SkyRL's trainer with both, its fully async one with `trainer.fully_async.enabled`, passing the rows' token variables to the trainer's Ray task. `tinker.py` (`tinker` extra, installed into tinker-cookbook's environment): `AgentEnvDatasetBuilder` is `rl.train`'s dataset builder, each rollout an episode played in process on the Episode API, not through a server: the turn's tools rendered by the cookbook's renderer, each sampled message recorded before its calls run, arguments typed by the tool's schema (a renderer that writes them as untyped text reads digits as numbers), a reply without a call ending the turn, and the reward made by `openenv/reward.py`'s rules; a rollout the cookbook cuts is scored once its group is done, a failed one raises for the cookbook's retry, and each group is filed in the run index as `/train|eval/step-` (`_episodes.RunGroup`). | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | | `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image, or, given `--image-ref ID=REF`, as that registry image pinned by digest with its context beside it, `image = { ref, dockerfile }`, which a run pulls and builds only when it can't; `task_images` lists the images an import writes, downloading nothing, so an export can push them first) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | diff --git a/Makefile b/Makefile index 1b557f909..3109d4a99 100644 --- a/Makefile +++ b/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 tinker-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 Tinker tier plays episodes through agent_env.training.tinker on tinker-cookbook itself, with a stand-in renderer, +# so nothing samples. tinker-cookbook and torch aren't in the dev extra: install them first, as +# .github/workflows/tinker.yml does. +tinker-test: ## Play episodes through the Tinker helper on tinker-cookbook itself (needs tinker-cookbook and CPU torch) + $(PYTHON) -m pytest tst/tinker --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) diff --git a/pyproject.toml b/pyproject.toml index af9a34c95..6d331a481 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,6 +92,11 @@ skyrl = [ "pyarrow>=19", "transformers>=5", ] +# agent_env.training.tinker: tinker-cookbook's rl.train on a bundle's tasks, each rollout an episode played in process; +# openenv for the reward rules build_app applies. tinker-cookbook, Tinker and torch are the trainer's to pin. +tinker = [ + "agentenv-framework[openenv]", +] # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml. aws = [ "boto3>=1.35.0", @@ -140,7 +145,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,harness,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,tinker,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 diff --git a/src/agent_env/training/_episodes.py b/src/agent_env/training/_episodes.py new file mode 100644 index 000000000..1d3e07ef4 --- /dev/null +++ b/src/agent_env/training/_episodes.py @@ -0,0 +1,76 @@ +"""What a trainer's adapter needs without the trainer installed: a tool call's arguments typed by its tool's schema, +and the run group a trainer files its episodes in.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Literal + + +_JSON_TYPES = {bool: "boolean", int: "integer", float: "number", type(None): "null"} + + +def typed_arguments(arguments: Any, schema: Mapping[str, Any]) -> Any: + """``arguments`` with each scalar whose type the tool's schema doesn't allow, but a string, as the string it was + written as. Renderers that write parameters as untyped text (Qwen3.5's ``19122``) parse + each with ``json.loads``, reading digits as a number; Qwen's own parser types values by the schema, as this does.""" + if not isinstance(arguments, dict): + return arguments + properties = schema.get("properties") or {} + typed = dict(arguments) + for name, value in arguments.items(): + allowed = _allowed_types(properties.get(name) or {}) + kind = _JSON_TYPES.get(type(value)) + widened = kind == "integer" and "number" in allowed + if kind is not None and "string" in allowed and kind not in allowed and not widened: + typed[name] = json.dumps(value) + return typed + + +def _allowed_types(schema: Mapping[str, Any]) -> set[str]: + declared = schema.get("type") + allowed = set(declared) if isinstance(declared, list) else {declared} if declared else set() + for option in schema.get("anyOf") or schema.get("oneOf") or []: + allowed |= _allowed_types(option) + return allowed + + +def run_group_id(run: str, phase: Literal["train", "eval"], step: int) -> str: + """The run group a trainer files a group of episodes in: its training run, whether they were played to train or + to evaluate, and the training step whose policy played them. Every trainer names its groups so, and the explorer + lists a task's runs by them.""" + return f"{run}/{phase}/step-{step:04d}" + + +class TrainingSteps: + """The training step a trainer's batches are played at: a training batch's index, and for an evaluation the step + after the last training batch handed out, for a trainer that, as ``rl.train`` does, evaluates a step's policy + before it hands out that step's batch and once more after the last.""" + + def __init__(self) -> None: + self._last_trained: int | None = None + + def trained(self, index: int) -> int: + self._last_trained = index + return index + + def evaluated(self) -> int: + return 0 if self._last_trained is None else self._last_trained + 1 + + +@dataclass(frozen=True) +class RunGroup: + """Where a trainer files a group of episodes: under the training batch ``trained`` names, fixed when the batch is + handed out, or, for an evaluation, under the step it is played at, read each time it is played, since a trainer + may make its evaluation's groups once and play them at every evaluation, as ``rl.train`` does.""" + + run: str + steps: TrainingSteps + trained: int | None = None + + def id(self) -> str: + if self.trained is not None: + return run_group_id(self.run, "train", self.trained) + return run_group_id(self.run, "eval", self.steps.evaluated()) diff --git a/src/agent_env/training/tinker.py b/src/agent_env/training/tinker.py new file mode 100644 index 000000000..5a80e593a --- /dev/null +++ b/src/agent_env/training/tinker.py @@ -0,0 +1,443 @@ +"""Tinker's RL on a bundle's tasks: tinker-cookbook's ``rl.train``, each rollout one episode of a task +(``agent_env.task.EpisodeSpec``) whose agent the model being trained plays. ``AgentEnvDatasetBuilder`` is the +``dataset_builder`` of ``rl.train.Config``: each batch is groups of episodes, a group playing one task, and every +episode is a ``Task.run`` of its own that the task's verifiers score as ``agent-env run`` would. + +The model reads the turn's system prompt, its tools (the envs' tools, and ``bash`` when its agent works in a sandbox) +as the cookbook's renderer writes them, and the turn's prompt. Each sampled message is recorded on the turn before its +tool calls run; a message without tool calls ends the turn, and the run's next turn, if it has one, comes as the next +user message. Once the run ends, the rollout's reward is the task's reward of its scores by the rules ``build_app`` +applies (``agent_env.openenv.reward``: the score at the step id ``reward`` names, or the task's one score), and each +score is a metric ``score/``. + +An episode lives on the event loop that started it, so every rollout runs on ``rl.train``'s: set no rollout executor. +A rollout the cookbook ends itself (a structural parse failure, the context or token limit) is scored too, once its +group's rollouts are done: its turn ends with no reply, the verifiers score where the episode got to, and that score is +added to the cookbook's own reward for the cut, ``failed_reward``. The metric ``cut`` is 1 for such a rollout.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import random +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +try: + import chz + from tinker_cookbook import tokenizer_utils + from tinker_cookbook.renderers import get_renderer + from tinker_cookbook.renderers.base import Message, Renderer, ToolCall, ToolSpec, get_text_content + from tinker_cookbook.rl.message_env import EnvFromMessageEnv, MessageEnv, MessageStepResult + from tinker_cookbook.rl.types import Env, EnvGroupBuilder, Metrics, RLDataset, RLDatasetBuilder, Trajectory +except ModuleNotFoundError as exc: + if (exc.name or "").split(".")[0] not in ("chz", "tinker", "tinker_cookbook"): + raise + raise ImportError("agent_env.training.tinker runs in tinker-cookbook's environment: install tinker-cookbook and " + "'agentenv-framework[tinker]' there") from exc + +from agent_env.bundle import BundleKind, parse_bundle +from agent_env.bundle.run import prepare_bundle +from agent_env.store.routing import namespace_routing +from agent_env.task import Episode, EpisodeError, EpisodeSpec, Task, Turn +from agent_env.openenv.reward import reward_of +from agent_env.training._episodes import RunGroup, TrainingSteps, typed_arguments + +logger = logging.getLogger(__name__) + +Reward = Callable[[dict[str, float]], float | None] + +#: Per event loop, the episodes that may be under way at once, by limit: a semaphore can't be pickled with a builder. +_SLOTS: dict[tuple[int, int], asyncio.Semaphore] = {} + + +def _slots(limit: int) -> asyncio.Semaphore: + key = (id(asyncio.get_running_loop()), limit) + if key not in _SLOTS: + _SLOTS[key] = asyncio.Semaphore(limit) + return _SLOTS[key] + + +class EpisodeEnv(MessageEnv): + """One episode of ``spec``'s task as a cookbook ``MessageEnv``: started at the first observation, played turn by + turn, rewarded from its scores once the run ends. ``max_turns`` bounds the model's messages; at the bound the turn + ends with what the model last said, and every later turn with no reply, so the verifiers score where it got to.""" + + def __init__(self, name: str, spec: EpisodeSpec, *, renderer: Renderer, reward: Reward, model_name: str, + sandbox: str | None, max_turns: int, tool_timeout: float, max_concurrent_episodes: int, + run_group_id: str): + self._name = name + self._run_group_id = run_group_id + self._spec = spec + self._renderer = renderer + self._reward = reward + self._model_name = model_name + self._sandbox = sandbox + self._max_turns = max_turns + self._tool_timeout = tool_timeout + self._max_concurrent = max_concurrent_episodes + self.history: list[Message] = [] + self._episode: Episode | None = None + self._turn: Turn | None = None + self._tools: Any = None + self._specs: list[ToolSpec] | None = None + self._holding = False + self._scored = False + self._messages = 0 + self._tool_calls = 0 + self._tool_errors = 0 + + async def initial_observation(self) -> list[Message]: + await _slots(self._max_concurrent).acquire() + self._holding = True + try: + self._episode = await self._spec.start(sandbox=self._sandbox, run_group_id=self._run_group_id) + turn = await self._episode.next_turn() + if turn is None: + raise EpisodeError(f"episode {self._episode.id} ended before its first turn") + await self._open(turn) + except BaseException: + await self.close() # a failed episode gives its slot back now, not when its group is cleaned up + raise + prefix = self._renderer.create_conversation_prefix_with_tools(tools=self._specs or [], + system_prompt=self._system_prompt(turn)) + self.history = [*prefix, {"role": "user", "content": turn.prompt}] + return self.history + + async def step(self, message: Message) -> MessageStepResult: + try: + return await self._step(message) + except BaseException: + await self.close() + raise + + async def _step(self, message: Message) -> MessageStepResult: + self.history.append(message) + self._messages += 1 + text = get_text_content(message) + calls = [(call, call.id or f"call_{self._messages}_{n}", typed_arguments(_arguments(call), self._schema(call))) + for n, call in enumerate(message.get("tool_calls") or [])] + # A call whose arguments aren't an object can't run, so it isn't recorded: the turn runs every recorded call. + await self._turn.record_model_call( + _chat_message(text, [call for call in calls if isinstance(call[2], dict)]), model=self._model_name) + logs: dict[str, str | int | float] = {"task": self._name, "episode": self._episode.id, "assistant": text} + for n, (call, call_id, arguments) in enumerate(calls): + if isinstance(arguments, dict): + result = await self._run(call, call_id, arguments) + else: + self._tool_errors += 1 + result = _tool_message(call, call_id, "error: the arguments aren't a JSON object") + self.history.append(result) + logs[f"tool_call_{n}"] = f"{call.function.name}({json.dumps(arguments)}) -> {str(result['content'])[:600]}" + unparsed = message.get("unparsed_tool_calls") or [] + for each in unparsed: + self._tool_errors += 1 + self.history.append({"role": "user", "content": f"Your tool call couldn't be parsed ({each.error}). " + "Call the tool again, with JSON arguments."}) + if (calls or unparsed) and self._messages < self._max_turns: + return MessageStepResult(reward=0.0, episode_done=False, next_messages=self.history, logs=logs) + result = await self._end_turn("" if calls else text) + result.logs = {**logs, **result.logs} + return result + + async def close(self) -> None: + """Abandon the episode unless its run has ended, and give back its slot. Safe to call more than once.""" + try: + if self._episode is not None and not self._episode.finished: + await self._episode.abandon("its rollout ended before the episode did") + finally: + self._release() + + async def _open(self, turn: Turn) -> None: + self._turn = turn + self._tools = await turn.tools(timeout=self._tool_timeout).__aenter__() + specs: list[ToolSpec] = [{"name": tool.name, "description": tool.description or "", + "parameters": tool.inputSchema or {"type": "object", "properties": {}}} + for tool in self._tools.list()] + if self._specs is not None and specs != self._specs: + raise EpisodeError(f"turn {turn.step_id!r} offers other tools than the episode's first turn, and the " + "renderer writes the tools once, at the start of the conversation") + self._specs = specs + + def _schema(self, call: ToolCall) -> dict: + return next((spec["parameters"] for spec in self._specs or [] if spec["name"] == call.function.name), {}) + + def _system_prompt(self, turn: Turn) -> str: + if turn.skills: + raise EpisodeError(f"turn {turn.step_id!r} gives the agent skills {[s.name for s in turn.skills]}, which " + "this module doesn't show the model yet") + return turn.system_prompt or "" + + async def _run(self, call: ToolCall, call_id: str, arguments: dict) -> Message: + self._tool_calls += 1 + name = call.function.name + offered = [spec["name"] for spec in self._specs or []] + if name not in offered: + self._tool_errors += 1 + return _tool_message(call, call_id, f"error: there is no tool {name!r}; the tools are {offered}") + result = await self._tools.call(name, arguments) + self._tool_errors += bool(result.isError) + text = "".join(_text(part) for part in result.content) + return _tool_message(call, call_id, f"error: {text}" if result.isError else text) + + async def score_cut(self) -> tuple[float, dict[str, float]]: + """The verifiers' reward for an episode whose rollout the cookbook cut, its open turn ended with no reply, and + the metric ``cut``; nothing for one ``step`` already rewarded.""" + if self._scored: + return 0.0, {"cut": 0.0} + try: + if self._turn is not None: + await self._turn.end("") + reward, metrics, _ = await self._finish() + except BaseException: + await self.close() + raise + return reward, {**metrics, "cut": 1.0} + + async def _end_turn(self, reply: str) -> MessageStepResult: + await self._turn.end(reply) + while (turn := await self._episode.next_turn()) is not None: + if self._messages < self._max_turns: + await self._open(turn) + self.history.append({"role": "user", "content": turn.prompt}) + return MessageStepResult(reward=0.0, episode_done=False, next_messages=self.history) + await turn.end("") + reward, metrics, scores = await self._finish() + return MessageStepResult(reward=reward, episode_done=True, next_messages=self.history, metrics=metrics, + logs={"scores": json.dumps(scores)}) + + async def _finish(self) -> tuple[float, dict[str, float], dict[str, float]]: + """End every turn the run has left with no reply, and reward the scores the run records.""" + while (turn := await self._episode.next_turn()) is not None: + await turn.end("") + scores = await self._episode.result() + self._release() + self._scored = True + reward = self._reward(scores) + if reward is None: + raise EpisodeError(f"episode {self._episode.id} recorded no score to make a reward of: {scores}") + metrics = {f"score/{step}": score for step, score in scores.items()} + metrics |= {"messages": self._messages, "tool_calls": self._tool_calls, "tool_errors": self._tool_errors} + return reward, metrics, scores + + def _release(self) -> None: + if self._holding: + self._holding = False + _slots(self._max_concurrent).release() + + +def _arguments(call: ToolCall) -> Any: + try: + return json.loads(call.function.arguments or "{}") + except json.JSONDecodeError: + return None + + +def _chat_message(text: str, calls: list[tuple[ToolCall, str, dict]]) -> dict[str, Any]: + """The model's message as an OpenAI chat message, with the tool calls that will run.""" + message: dict[str, Any] = {"role": "assistant", "content": text} + if calls: + message["tool_calls"] = [{"id": call_id, "type": "function", + "function": {"name": call.function.name, "arguments": json.dumps(arguments)}} + for call, call_id, arguments in calls] + return message + + +def _tool_message(call: ToolCall, call_id: str, content: str) -> Message: + return {"role": "tool", "content": content, "tool_call_id": call_id, "name": call.function.name} + + +def _text(part: Any) -> str: + kind = part.get("type") if isinstance(part, dict) else getattr(part, "type", None) + if kind == "text": + return part.get("text", "") if isinstance(part, dict) else part.text + return f"[{kind} content]" + + +class EpisodeGroupBuilder(EnvGroupBuilder): + """``group_size`` episodes of one task, each its own run, filed in ``run_group``, which the explorer shows + together. ``make_envs`` may be called again for a retry, and ``cleanup`` abandons every episode still under way.""" + + def __init__(self, name: str, spec: EpisodeSpec, reward: Reward, settings: _Settings, group_size: int, + run_group: RunGroup): + self.name = name + self.run_group = run_group + self.spec = spec + self.reward = reward + self.settings = settings + self.group_size = group_size + self._made: list[EpisodeEnv] = [] + + async def make_envs(self) -> Sequence[Env]: + renderer = self.settings.renderer() + run_group_id = self.run_group.id() + envs = [] + for _ in range(self.group_size): + episode_env = EpisodeEnv(self.name, self.spec, renderer=renderer, reward=self.reward, + model_name=self.settings.model_name, sandbox=self.settings.sandbox, + max_turns=self.settings.max_turns, tool_timeout=self.settings.tool_timeout, + max_concurrent_episodes=self.settings.max_concurrent_episodes, + run_group_id=run_group_id) + self._made.append(episode_env) + envs.append(EnvFromMessageEnv(renderer, episode_env, failed_parse_reward=self.settings.failed_reward, + context_overflow_reward=self.settings.failed_reward, + max_trajectory_tokens=self.settings.max_trajectory_tokens)) + return envs + + async def compute_group_rewards(self, trajectory_group: list[Trajectory], + env_group: Sequence[Env]) -> list[tuple[float, Metrics]]: + return list(await asyncio.gather(*(env.message_env.score_cut() for env in env_group))) + + async def cleanup(self) -> None: + made, self._made = self._made, [] + for result in await asyncio.gather(*(env.close() for env in made), return_exceptions=True): + if isinstance(result, BaseException): + logger.warning("task %s: an episode didn't close: %s", self.name, result) + + def logging_tags(self) -> list[str]: + return [self.settings.tag] + + +class _Settings: + """What every group of a dataset shares.""" + + def __init__(self, *, model_name: str, renderer_name: str, sandbox: str | None, max_turns: int, + tool_timeout: float, max_concurrent_episodes: int, max_trajectory_tokens: int | None, + failed_reward: float, tag: str): + self.model_name = model_name + self.renderer_name = renderer_name + self.sandbox = sandbox + self.max_turns = max_turns + self.tool_timeout = tool_timeout + self.max_concurrent_episodes = max_concurrent_episodes + self.max_trajectory_tokens = max_trajectory_tokens + self.failed_reward = failed_reward + self.tag = tag + + def renderer(self) -> Renderer: + return get_renderer(self.renderer_name, tokenizer_utils.get_tokenizer(self.model_name), + model_name=self.model_name) + + +class EpisodeDataset(RLDataset): + """A split's tasks in batches of ``groups_per_batch`` groups, shuffled anew each of ``epochs`` passes; a group's + episodes are filed under the run, whether they train or evaluate, and the training step that plays them.""" + + def __init__(self, tasks: list[tuple[str, EpisodeSpec, Reward]], settings: _Settings, *, group_size: int, + groups_per_batch: int, epochs: int, seed: int, shuffle: bool, run_name: str, + steps: TrainingSteps, trains: bool): + self._run_name = run_name + self._steps = steps + self._trains = trains + self._tasks = tasks + self._settings = settings + self._group_size = group_size + self._groups_per_batch = min(groups_per_batch, len(tasks)) + self._epochs = epochs + self._seed = seed + self._shuffle = shuffle + self._per_epoch = len(tasks) // self._groups_per_batch + + def get_batch(self, index: int) -> Sequence[EnvGroupBuilder]: + epoch, at = divmod(index, self._per_epoch) + order = list(range(len(self._tasks))) + if self._shuffle: + random.Random(self._seed + epoch).shuffle(order) + picked = order[at * self._groups_per_batch:(at + 1) * self._groups_per_batch] + group = RunGroup(self._run_name, self._steps, self._steps.trained(index) if self._trains else None) + return [EpisodeGroupBuilder(name, spec, reward, self._settings, self._group_size, group) + for name, spec, reward in (self._tasks[i] for i in picked)] + + def __len__(self) -> int: + return self._per_epoch * self._epochs + + +@chz.chz +class AgentEnvDatasetBuilder(RLDatasetBuilder): + """The tasks of the bundle folder ``bundle`` for ``rl.train``, prepared as ``agent-env run`` prepares them, their + episodes filed in run groups named ``/train|eval/step-``. ``train_split`` and ``test_split`` are + ``all`` or an eval of the bundle. ``reward`` is the step id whose score is the reward, needed + when a task records more than one. ``sandbox`` is where episodes deploy, as ``agent-env run --sandbox`` sets it, + and ``max_concurrent_episodes`` bounds the runs under way at once. A rollout the cookbook ends itself is scored + once its group is done, and given ``failed_reward`` on top.""" + + bundle: str + run_name: str + model_name: str + renderer_name: str + group_size: int + groups_per_batch: int + train_split: str = "all" + test_split: str | None = None + test_group_size: int = 1 + reward: str | None = None + sandbox: str | None = None + max_turns: int = 30 + tool_timeout: float = 60.0 + turn_timeout: float | None = None + max_concurrent_episodes: int = 8 + max_trajectory_tokens: int | None = None + failed_reward: float = 0.0 + epochs: int = 1 + seed: int = 0 + + async def __call__(self) -> tuple[RLDataset, RLDataset | None]: + wanted = [split for split in (self.train_split, self.test_split) if split is not None] + splits = await asyncio.to_thread(_splits, Path(self.bundle).expanduser(), wanted, self.sandbox) + settings = _Settings(model_name=self.model_name, renderer_name=self.renderer_name, sandbox=self.sandbox, + max_turns=self.max_turns, tool_timeout=self.tool_timeout, + max_concurrent_episodes=self.max_concurrent_episodes, + max_trajectory_tokens=self.max_trajectory_tokens, failed_reward=self.failed_reward, + tag=Path(self.bundle).name) + steps = TrainingSteps() + train = EpisodeDataset(playable(splits[self.train_split], self.reward, self.turn_timeout), settings, + group_size=self.group_size, groups_per_batch=self.groups_per_batch, epochs=self.epochs, + seed=self.seed, shuffle=True, run_name=self.run_name, steps=steps, trains=True) + test = None + if self.test_split is not None: + tasks = playable(splits[self.test_split], self.reward, self.turn_timeout) + test = EpisodeDataset(tasks, settings, group_size=self.test_group_size, groups_per_batch=len(tasks), + epochs=1, seed=self.seed, shuffle=False, run_name=self.run_name, + steps=steps, trains=False) + return train, test + + +def _splits(bundle: Path, splits: list[str], sandbox: str | None) -> dict[str, list[tuple[str, Task]]]: + """Each split's tasks with their names, ``all`` or an eval's, prepared as ``agent-env run`` prepares them before + its first task, without the agent the trainer plays (``played``), as the OpenEnv server prepares a bundle.""" + with namespace_routing(): + every = [entry.name for entry in parse_bundle(bundle).entries if entry.kind is BundleKind.TASK] + prepared = prepare_bundle(bundle, tasks=every, evals=[split for split in splits if split != "all"], + sandbox=sandbox, on_progress=logger.info, played=True) + tasks = [(entry.name, task) for entry, task in prepared.tasks] + by_id = {task.id: (name, task) for name, task in tasks} + found = {"all": tasks, **{entry.entry.name: [by_id[ref.id] for ref in entry.references] + for entry in prepared.plan.evals}} + return {split: found[split] for split in splits} + + +def playable(tasks: list[tuple[str, Task]], reward: str | None, + turn_timeout: float | None) -> list[tuple[str, EpisodeSpec, Reward]]: + """Each task that can run as an episode, with its reward; the rest are logged and left out. Raises ValueError when + a task's scores make no reward, since a rollout must be rewarded, or when no task is left.""" + episodes, refused, unscorable = [], {}, {} + for name, task in tasks: + try: + spec = EpisodeSpec.from_task(task, turn_timeout=turn_timeout) + except EpisodeError as e: + refused[name] = str(e) + continue + task_reward, problem = reward_of(task, reward) + if problem or not any(getattr(step, "verifier_id", None) for step in task.steps): + unscorable[name] = problem or f"task {task.id!r} records no score" + continue + episodes.append((name, spec, task_reward)) + for why in refused.values(): + logger.warning("left out: %s", why) + if unscorable: + raise ValueError("reward= can't make a reward of every task's scores:\n" + "\n".join(unscorable.values())) + if not episodes: + raise ValueError(f"no task can be played as an episode: {refused}") + return episodes diff --git a/tst/tinker/__init__.py b/tst/tinker/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/tinker/conftest.py b/tst/tinker/conftest.py new file mode 100644 index 000000000..5cd349748 --- /dev/null +++ b/tst/tinker/conftest.py @@ -0,0 +1,6 @@ +"""The Episode API's fake steps, stand-in teardown and tool sessions, and the local stores, which these tests play +through the Tinker helper's episodes.""" + +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 diff --git a/tst/tinker/episode_env_test.py b/tst/tinker/episode_env_test.py new file mode 100644 index 000000000..5c591b8ce --- /dev/null +++ b/tst/tinker/episode_env_test.py @@ -0,0 +1,133 @@ +"""The Tinker helper's episodes on the Episode API's fake steps, as tinker-cookbook's rollouts step them: a rollout's +tool calls run on its turn and its reply ends it with the task's reward, filed in its run group; the turn bound ends a +turn; a rollout the cookbook cut is scored once its group is done; a bundle's splits as prepared for a trainer to play; +and the tasks a dataset plays.""" + +import json +import logging + +import pytest + +from agent_env.openenv.reward import _score_at +from agent_env.runner import store as run_store +from agent_env.task import Task +from tst.unit.bundle._support import layout +from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _spec +from tst.util.capabilities import TINKER, missing_capability_reason + +tinker = pytest.importorskip("agent_env.training.tinker", reason=missing_capability_reason(TINKER), + exc_type=ImportError) +base = pytest.importorskip("tinker_cookbook.renderers.base", reason=missing_capability_reason(TINKER)) + +GROUP = "dock-run/train/step-0003" + + +class _Renderer: + """Stands in for a cookbook renderer: the system prompt and the names of the tools it is given, as one message.""" + + def create_conversation_prefix_with_tools(self, tools, system_prompt=""): + return [{"role": "system", "content": json.dumps({"system": system_prompt, + "tools": [tool["name"] for tool in tools]})}] + + +def _env(*steps, max_turns=8): + return tinker.EpisodeEnv("dock", _spec(*steps), renderer=_Renderer(), reward=_score_at("replies"), + model_name="policy", sandbox=None, max_turns=max_turns, tool_timeout=5, + max_concurrent_episodes=4, run_group_id=GROUP) + + +def _call(name, **arguments): + return base.ToolCall(function=base.ToolCall.FunctionBody(name=name, arguments=json.dumps(arguments)), + id=f"call-{name}") + + +@pytest.mark.asyncio +async def test_a_rollouts_calls_run_on_its_turn_and_its_reply_ends_it_with_the_tasks_reward_in_its_group( + teardowns, tools): + env = _env(ITEMS, ALICE, TURN, REPLIES) + + system, user = await env.initial_observation() + assert json.loads(system["content"]) == {"system": "You keep the items.", "tools": ["add_item"]} + assert user == {"role": "user", "content": "Add a crate."} + called = await env.step({"role": "assistant", "content": "", "tool_calls": [_call("add_item", item="crate")]}) + assert not called.episode_done + assert called.next_messages[-1] == {"role": "tool", "content": "add_item {'item': 'crate'}", + "tool_call_id": "call-add_item", "name": "add_item"} + ended = await env.step({"role": "assistant", "content": "Added a crate."}) + + assert (ended.episode_done, ended.reward) == (True, 1.0) + assert {key: ended.metrics[key] for key in ("score/replies", "messages", "tool_calls", "tool_errors")} == { + "score/replies": 1.0, "messages": 2, "tool_calls": 1, "tool_errors": 0} + run = run_store.get_run(f"episode-{ended.logs['episode']}") + assert (run.status, run.overrides["metadata"]["run_group_id"]) == ("COMPLETED", GROUP) + assert await env.score_cut() == (0.0, {"cut": 0.0}) + + +@pytest.mark.asyncio +async def test_a_turn_ends_at_the_bound_on_the_models_messages_with_no_reply(teardowns, tools): + env = _env(ITEMS, ALICE, TURN, REPLIES, max_turns=1) + await env.initial_observation() + + ended = await env.step({"role": "assistant", "content": "", "tool_calls": [_call("add_item", item="crate")]}) + + assert (ended.episode_done, ended.reward, ended.metrics["tool_calls"]) == (True, 0.0, 1) + + +@pytest.mark.asyncio +async def test_a_rollout_the_cookbook_cut_is_scored_once_its_group_is_done(teardowns, tools): + env = _env(ITEMS, ALICE, TURN, REPLIES) + await env.initial_observation() + + reward, metrics = await env.score_cut() + + assert (reward, metrics["score/replies"], metrics["cut"]) == (0.0, 0.0, 1.0) + assert await env.score_cut() == (0.0, {"cut": 0.0}) + + +def test_a_split_is_an_evals_tasks_prepared_without_the_agent_the_trainer_plays(teardowns, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + played = json.dumps([ITEMS, {**ALICE, "a2a_agent_id": "nobody", "env_ids": []}, TURN, REPLIES]) + root = layout(tmp_path / "dock", {"tasks/crate.json": played, "tasks/again.json": played, + "evals/train.toml": 'tasks = ["crate"]\n'}) + + splits = tinker._splits(root, ["all", "train"], None) + + assert {split: sorted(name for name, _ in tasks) for split, tasks in splits.items()} == { + "all": ["again", "crate"], "train": ["crate"]} + + +ENV = {"id": "env", "type": "deploy_env", "env_id": "retail"} +AGENT = {"id": "agent", "type": "deploy_agent", "agent_name": "policy", "env_ids": ["retail"], + "role": "agent@retail.example", "system_prompt": "You run the store."} +PLAY = {"id": "play", "type": "prompt_agent", "agent_name": "policy", "prompt_id": "play", "prompt": "Help."} +BOB = {"id": "bob", "type": "deploy_agent", "agent_name": "bob", "a2a_agent_id": "echo"} +ASK_BOB = {"id": "ask-bob", "type": "prompt_agent", "agent_name": "bob", "prompt": "Check."} + + +def _verifier(name): + return {"id": name, "type": "env_outcome_verifier", "env_id": "retail", "file_artifact_id": "check", + "verifier_id": name} + + +def _task(id, *steps): + return id, Task.from_dict({"id": id, "version": 1, "steps": [dict(step) for step in steps]}) + + +def test_tasks_that_cant_run_as_an_episode_are_left_out_and_the_rest_rewarded(caplog): + played = _task("played", ENV, AGENT, PLAY, _verifier("records"), _verifier("db_match")) + two_agents = _task("two-agents", ENV, AGENT, PLAY, BOB, ASK_BOB, _verifier("records")) + with caplog.at_level(logging.WARNING): + [(name, spec, reward)] = tinker.playable([played, two_agents], "records", None) + assert (name, spec.trainer_agents) == ("played", ("policy",)) + assert reward({"records": 0.25, "db_match": 0.0}) == 0.25 + assert "two-agents" in caplog.text + + +@pytest.mark.parametrize(("tasks", "reward", "refusal"), [ + ([_task("unscored", ENV, AGENT, PLAY)], None, "records no score"), + ([_task("two", ENV, AGENT, PLAY, _verifier("records"), _verifier("db_match"))], None, "name the one"), + ([_task("two-agents", ENV, AGENT, PLAY, BOB, ASK_BOB, _verifier("records"))], "records", "no task can be played"), +]) +def test_a_dataset_whose_tasks_make_no_reward_or_no_episode_is_refused(tasks, reward, refusal): + with pytest.raises(ValueError, match=refusal): + tinker.playable(tasks, reward, None) diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 1b441befc..db01d36ce 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -32,7 +32,7 @@ def find_spec(self, fullname, path=None, target=None): _NOT_CORE = ( "agent_env.providers.sandbox_providers.sail_vm", "agent_env.providers.sandbox_providers.smol_vm", "agent_env.openenv", "agent_env.training.nemo_gym", "agent_env.training.verl", "agent_env.training.skyrl", - "agent_env.training.harness", + "agent_env.training.harness", "agent_env.training.tinker", "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) diff --git a/tst/unit/training/episodes_test.py b/tst/unit/training/episodes_test.py new file mode 100644 index 000000000..4af7c31c1 --- /dev/null +++ b/tst/unit/training/episodes_test.py @@ -0,0 +1,44 @@ +"""The arguments a rollout's tool calls run with, typed by each tool's schema, and the run groups a trainer files its +episodes in by training step.""" + +from agent_env.training._episodes import RunGroup, TrainingSteps, run_group_id, typed_arguments + +LOOKUP = {"type": "object", "properties": { + "zip": {"type": "string"}, + "note": {"type": ["string", "null"]}, + "order": {"anyOf": [{"type": "string"}, {"type": "null"}]}, + "count": {"type": "integer"}, + "price": {"type": "number"}, + "gift": {"type": "boolean"}, + "items": {"type": "array", "items": {"type": "string"}}, +}} + +def test_a_scalar_its_schema_wants_as_a_string_is_given_back_as_the_string_it_was_written_as(): + assert typed_arguments({"zip": 19122, "note": 7, "order": True}, LOOKUP) == { + "zip": "19122", "note": "7", "order": "true"} + + +def test_values_their_schema_allows_are_left_as_parsed(): + arguments = {"zip": "19122", "note": None, "count": 3, "price": 2, "gift": False, "items": ["1151293680"]} + assert typed_arguments(arguments, LOOKUP) == arguments + + +def test_an_argument_the_schema_doesnt_name_and_arguments_that_arent_an_object_are_left_alone(): + assert typed_arguments({"other": 5}, LOOKUP) == {"other": 5} + assert typed_arguments([19122], LOOKUP) == [19122] + + +def test_a_group_is_filed_under_its_run_phase_and_step_and_an_evaluation_at_the_step_it_measures(): + steps = TrainingSteps() + assert steps.evaluated() == 0 + assert [steps.trained(index) for index in range(4)] == [0, 1, 2, 3] + assert steps.evaluated() == 4 + assert run_group_id("retail-run2", "eval", steps.evaluated()) == "retail-run2/eval/step-0004" + + +def test_an_evaluations_group_made_once_is_filed_at_the_step_each_time_it_is_played(): + steps = TrainingSteps() + evaluation = RunGroup("retail-run2", steps) + assert evaluation.id() == "retail-run2/eval/step-0000" + trained = RunGroup("retail-run2", steps, steps.trained(0)) + assert (trained.id(), evaluation.id()) == ("retail-run2/train/step-0000", "retail-run2/eval/step-0001") diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index 2eddbb726..cd36ec350 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -53,6 +53,10 @@ #: and torch aren't in the dev extra, and only that tier's CI job installs them. SKYRL = "skyrl_installed" +#: The Tinker tier (``tst/tinker``) plays episodes through ``agent_env.training.tinker`` on tinker-cookbook itself; +#: tinker-cookbook and torch aren't in the dev extra, and only that tier's CI job installs them. +TINKER = "tinker_installed" + #: ``collect_artifacts`` sizes a file with ``stat -c %s``, which only GNU stat takes (macOS ships BSD stat), #: so collecting off a local sandbox needs GNU stat first on PATH. GNU_STAT = "gnu_stat" diff --git a/uv.lock b/uv.lock index 7b050c38a..3e1d08130 100644 --- a/uv.lock +++ b/uv.lock @@ -167,6 +167,9 @@ tenki = [ tensorlake = [ { name = "tensorlake" }, ] +tinker = [ + { name = "openenv" }, +] trl = [ { name = "datasets" }, { name = "openenv" }, @@ -188,10 +191,11 @@ verl = [ [package.metadata] requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, - { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "harness", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "harness", "openenv", "sail", "skyrl", "tensorlake", "tenki", "tinker", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'harness'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'nemo-gym'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'skyrl'" }, + { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'tinker'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'trl'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'verl'" }, { name = "agentenv-framework-protocol", extras = ["computer"], editable = "packages/agentenv-protocol" }, @@ -257,7 +261,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, ] -provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "tinker", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol"