fix(training): TRL tools carry their MCP input schema, nested types included

pyproject.toml

@@ -59,6 +59,7 @@ openenv = [
 trl = [
     "agentenv-framework[openenv]",
     "datasets>=3",
+    "transformers>=5.2",  # renders the tools as TRL shows them to a model; 5.2 is what TRL's environment_factory needs
 ]
 # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml.
 aws = [
@@ -106,7 +107,6 @@ boat = [
 dev = [
     "agentenv-framework[aws,boat,boxd,daytona,gcp,openenv,sail,tensorlake,tenki,trl,vercel]",  # 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
-    "transformers>=5.2",  # get_json_schema, which reads agent_env.training's tool methods as TRL shows them to a model
     "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
     "griffelib==2.3.0",  # the plugin API check; exact, since what it reports as a break changes between releases

src/agent_env/training/_tools.py

@@ -1,6 +1,6 @@
 """Python methods for an env's tools, typed and documented from each tool's JSON Schema the way transformers'
-``get_json_schema`` reads a function, which is how TRL shows a model an environment's tools. Built without generating
-source; a method calls its instance's ``_call(name, arguments)``."""
+``get_json_schema`` reads a function, which is how TRL finds an environment's tools, and the ``SchemaTool`` each one
+reaches the model as. Built without generating source; a method calls its instance's ``_call(name, arguments)``."""
 
 from __future__ import annotations
 
@@ -47,13 +47,32 @@ def method(self: Any, **arguments: Any) -> str:
         return self._call(name, arguments)
 
     method.__name__ = method.__qualname__ = name
-    method.__doc__ = (tool.get("description") or f"Call the {name} tool.") + (
+    method.__doc__ = _summary(tool) + (
         "\n\nArgs:\n" + "\n".join(described) if described else "")
     method.__annotations__ = annotations
     method.__signature__ = inspect.Signature(parameters, return_annotation=str)
     return method
 
 
+class SchemaTool(dict):
+    """A tool as both its env's JSON Schema, which transformers renders as given, and the callable TRL calls: a method's
+    type hints have no nested types. TRL lists the methods on the class and takes the tools from the instance, which
+    sets one over each; a tool has no docstring, so ``get_json_schema`` refuses it rather than make up a schema."""
+
+    def __init__(self, tool: Mapping[str, Any], call: Callable[..., str]):
+        super().__init__(type="function", function={
+            "name": tool["name"], "description": _summary(tool), "parameters": tool["input_schema"]})
+        self.__name__, self.__doc__ = tool["name"], None
+        self._function = call
+
+    def __call__(self, **arguments: Any) -> str:
+        return self._function(**arguments)
+
+
+def _summary(tool: Mapping[str, Any]) -> str:
+    return tool.get("description") or f"Call the {tool['name']} tool."
+
+
 def _callable_name(name: str) -> bool:
     return name.isidentifier() and not keyword.iskeyword(name) and not name.startswith("_")
 

src/agent_env/training/trl.py

@@ -1,7 +1,7 @@
 """TRL's trainers on a bundle served as an OpenEnv env (``agent_env.openenv.build_app``, or the package ``openenv import``
 writes): ``AgentEnv`` reads the server's tasks into a dataset, and makes the ``environment_factory`` whose environments
 each play one episode of a task over OpenEnv's WebSocket, offer its turn's tools as methods and own its reward.
-``GRPOTrainer`` takes both as they are, verified with TRL 1.16, and AsyncGRPO takes the same ``environment_factory``;
+``GRPOTrainer`` takes both as they are, verified with TRL 1.16 and 1.17, and AsyncGRPO the same ``environment_factory``;
 either needs TRL 1.15 or later, for an environment that owns its reward. Every session runs on one event-loop thread of
 this module's own, since TRL calls an environment synchronously, from its own threads."""
 
@@ -29,16 +29,17 @@
 try:
     from datasets import Dataset
     from openenv.core.generic_client import GenericEnvClient
+    from transformers.utils.chat_template_utils import render_jinja_template
 except ModuleNotFoundError as exc:
-    if exc.name not in ("datasets", "openenv"):
+    if exc.name not in ("datasets", "openenv", "transformers"):
         raise
     raise ImportError("agent_env.training.trl requires pip install 'agentenv-framework[trl]'") from exc
 
 from agent_env.openenv import build_app
 from agent_env.openenv.environment import END_TURN as END_TURN_TOOL
 from agent_env.openenv.environment import TIMEOUT_METADATA, Toolset
 from agent_env.task.teardown import TERMINATE_TIMEOUT_SECONDS
-from agent_env.training._tools import tool_method
+from agent_env.training._tools import SchemaTool, tool_method
 
 logger = logging.getLogger(__name__)
 
@@ -247,8 +248,9 @@ class _Rollout:
 
 class _Environment:
     """One session on the server, playing one episode at a time. Its public methods other than ``reset`` and
-    ``get_reward`` are the tools, which ``AgentEnv.factory`` binds per toolset; every call on the session waits its
-    turn, as AsyncGRPO can call two tools of a completion at once."""
+    ``get_reward`` are the tools, which ``AgentEnv.factory`` binds per toolset, and each but ``end_turn`` is set over
+    with its ``SchemaTool``; every call on the session waits its turn, as AsyncGRPO can call two tools of a completion
+    at once."""
 
     _shared: ClassVar[_Shared]
     _tools: ClassVar[tuple[dict[str, Any], ...]]
@@ -258,6 +260,9 @@ def __init__(self) -> None:
         self._opened = False
         self._rollout: _Rollout | None = None
         self._lock = asyncio.Lock()
+        for tool in self._tools:
+            if tool["name"] != END_TURN:
+                setattr(self, tool["name"], SchemaTool(tool, getattr(self, tool["name"])))
         _environments.add(self)
 
     def reset(self, *, split: str, index: int, task_id: str | None = None, **_: Any) -> None:
@@ -413,9 +418,19 @@ def _environment_class(name: str, tools: tuple[dict[str, Any], ...], shared: _Sh
     methods = {tool["name"]: tool_method(tool) for tool in tools}
     for method in methods.values():
         method.__qualname__ = f"{name}.{method.__name__}"
+    _check_rendered([SchemaTool(tool, methods[tool["name"]]) for tool in tools if tool["name"] != END_TURN])
     return _EnvironmentClass(name, (_Environment,), {**methods, "_shared": shared, "_tools": tools})
 
 
+def _check_rendered(tools: list[SchemaTool]) -> None:
+    """Fails as the trainer sets up, not as the model meets the tools, when transformers doesn't render a tool given
+    as its JSON Schema as it is."""
+    [rendered], _ = render_jinja_template([[]], tools=tools, chat_template="{{ tools | tojson }}")
+    if changed := [tool.__name__ for tool, shown in zip(tools, json.loads(rendered), strict=True) if shown != tool]:
+        raise RuntimeError(f"this transformers renders the tools {changed} other than as their env's JSON Schema, so "
+                           "TRL would show the model other tools than the env's")
+
+
 def _reduce_environment_class(environment: _EnvironmentClass) -> tuple[Any, ...] | str:
     """A class ``_environment_class`` built pickles as what it's built from, and a subclass of one by its name, as any
     class does, so it keeps what it overrides."""

tst/data/training/qwen3_chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

tst/trl/grpo_test.py

@@ -1,14 +1,16 @@
 """TRL's GRPOTrainer trains a step through ``agent_env.training.trl`` on CPU, against the OpenEnv server ``build_app``
 makes over the Episode API's fake steps. TRL marks ``environment_factory`` experimental, so this pins what the helper
-relies on: TRL resets an environment with its row, calls its tools, reads ``get_reward``, and trains on the rewards.
+relies on: TRL shows the model each tool as its env publishes it, resets an environment with its row, calls its tools,
+reads ``get_reward``, and trains on the rewards.
 As in TRL's own environment tests, a tiny model's ``generate`` is stood in for. CI runs it against TRL 1.16, the oldest
 that trains on CPU, and the latest; AsyncGRPO needs a vLLM server, so it isn't covered here."""
 
 from unittest.mock import patch
 
 import pytest
 
-from tst.unit.training.trl_test import PLAYED
+from tst.unit.training.tools_test import tools_in
+from tst.unit.training.trl_test import PLAYED, SHOWN, STACK
 from tst.util.capabilities import TRL, missing_capability_reason
 
 torch = pytest.importorskip("torch", reason=missing_capability_reason(TRL))
@@ -19,9 +21,10 @@
 END = '<tool_call>\n{"name": "end_turn", "arguments": {"response": "Added a crate."}}\n</tool_call>'
 
 
-def test_grpotrainer_trains_a_step_on_a_served_bundle_and_scores_every_rollout(served, teardowns, tmp_path,
+def test_grpotrainer_trains_a_step_on_a_served_bundle_and_scores_every_rollout(served, teardowns, tools, tmp_path,
                                                                                  monkeypatch):
     monkeypatch.setenv("TRL_EXPERIMENTAL_SILENCE", "1")
+    tools("add_item", STACK)
     env = served({"a": PLAYED, "b": PLAYED})
     rewards = []
 
@@ -36,10 +39,12 @@ def get_reward(self):
         use_cpu=True))
     before = {name: param.clone() for name, param in trainer.model.named_parameters()}
 
-    with patch.object(trainer.model, "generate", side_effect=_Policy(trainer.processing_class)):
+    policy = _Policy(trainer.processing_class)
+    with patch.object(trainer.model, "generate", side_effect=policy):
         trainer.train()
 
     [logged] = [row for row in trainer.state.log_history if "tools/call_frequency" in row]
+    assert tools_in(policy.prompt) == SHOWN
     assert sorted(rewards) == [0.0, 0.0, 1.0, 1.0]
     assert (logged["tools/call_frequency"], logged["tools/failure_frequency"], logged["rewards/Recorded/mean"]) == (
         pytest.approx(1.0), 0.0, pytest.approx(0.5))
@@ -50,14 +55,16 @@ def get_reward(self):
 
 class _Policy:
     """Stands in for ``generate``: of each group, the first rollout adds a crate, then ends its turn, then stops, and the
-    second answers without a tool, so its turn is ended empty and scores 0."""
+    second answers without a tool, so its turn is ended empty and scores 0. ``prompt`` is the first it was given."""
 
     def __init__(self, tokenizer):
         self.tokenizer = tokenizer
         self.calls = 0
+        self.prompt = None
 
     def __call__(self, input_ids, **_):
         self.calls += 1
+        self.prompt = self.prompt or self.tokenizer.decode(input_ids[0])
         if self.calls == 1:
             texts = [ADD if row % 2 == 0 else "I won't." for row in range(len(input_ids))]
         else:

tst/unit/openenv/environment_test.py

@@ -21,9 +21,9 @@
 
 
 class _Tools:
-    """Stands in for the proxy's ``ToolSession``: every env offers ``offered``, and a call answers with what ``answers`` maps
-    its tool to, raising it when it's an exception. A session opens once ``release``, a ``threading.Event`` since the
-    in-process proxy runs on a loop of its own, is set, when there is one."""
+    """Stands in for the proxy's ``ToolSession``: every env offers ``offered``, tools or their names, and a call answers
+    with what ``answers`` maps its tool to, raising it when it's an exception. A session opens once ``release``, a
+    ``threading.Event`` since the in-process proxy runs on a loop of its own, is set, when there is one."""
 
     def __init__(self, offered, answers):
         self.offered, self.answers, self.closed, self.release = offered, answers, 0, None
@@ -38,8 +38,9 @@ async def open(self):
                 return self
 
             async def list_tools(self):
-                return [EnvironmentTool(name=name, description=f"{name}.", inputSchema={"type": "object"})
-                        for name in tools.offered]
+                return [tool if isinstance(tool, EnvironmentTool)
+                        else EnvironmentTool(name=tool, description=f"{tool}.", inputSchema={"type": "object"})
+                        for tool in tools.offered]
 
             async def call_tool(self, name, arguments=None):
                 answer = tools.answers.get(name, ToolResult(content=[{"type": "text", "text": f"{name} {arguments}"}]))