[earakely-scale]

What changes

TRL turns an environment_factory's methods into tool schemas with transformers' get_json_schema, which reads only a function's type hints and Google docstring. A type hint can't carry an object's properties, so a tool argument such as an array of objects reached the model as {"type": "array", "items": {"type": "object"}}, with its fields, required lists, enums and bounds only in prose. On the GridSim desk, untrained Qwen3-4B's main failure mode is malformed nested JSON for the offer tools, which take arguments like this.

Neither TRL (1.16, 1.17, main) nor transformers (5.21, 5.22, main) has a supported way to attach a schema to a callable. TRL lists an environment's tools as functions on the class, fetches each from the instance with getattr, and passes what it fetched straight to apply_chat_template(tools=...) and parse_response(tools=...). Transformers renders a dict tool as given and converts a function or method with get_json_schema.

• SchemaTool (agent_env/training/_tools.py) is both things at once. It is a dict holding {"type": "function", "function": {name, description, parameters: inputSchema}}, which transformers renders unchanged. It also has __name__ and __call__, which is all TRL needs to call it.
• Setup per instance. Each environment instance sets a SchemaTool over every tool method except end_turn. The class methods stay as they are, because TRL finds the tools through them.
• end_turn keeps its current schema.
• Failure behaviour, both made explicit:
  • If TRL ever takes the tools from the class instead of the instance, the method schemas come back. Those are today's behaviour, and a test pins them.
  • A SchemaTool has no docstring, so get_json_schema on one raises DocstringParsingException instead of building a wrong schema. On Python 3.14 it would otherwise silently produce a schema with no properties.
  • The factory checks at build time that transformers renders each SchemaTool as its own schema. Otherwise it raises a RuntimeError naming the tools.
• Dependency: transformers>=5.2 moves from dev into the trl extra. trl.py now imports transformers for that build-time check, and test_a_module_using_an_extra_imports_only_what_one_extra_and_core_declare requires a module's imports to be declared by its extra. TRL already requires transformers, so this adds nothing to a TRL install, and the extra still doesn't install TRL. 5.2 is the version TRL's environment_factory needs. The lockfile only moves the extra marker; resolution is unchanged.

submit_dam_offers on the GridSim desk, before and after

These come from AgentEnv.local on the desk-v4 bench bundle, rendered by Qwen3-4B's chat template.

Before, curves (as_offers the same way):

{"type": "array", "items": {"type": "object"}, "nullable": true,
 "description": "curves Each entry: hours (string or array, required): operating hours: \"HE07-HE10\", \"HE07\", \"all\", or a list like [7, 8]; points (array): [[mw, price], ...] from max charge (negative MW) to max discharge; MW and price non-decreasing; 2-10 points."}

as_offers[].product showed as a bare string, and its enum was lost.

After, exactly what the env publishes:

{"type": "array", "items": {"type": "object", "properties": {
  "hours": {"description": "operating hours: \"HE07-HE10\", \"HE07\", \"all\", or a list like [7, 8]", "anyOf": [{"type": "string"}, {"type": "array", "items": {"type": "integer"}}]},
  "points": {"type": "array", "items": {"type": "array", "items": {"type": "number"}, "minItems": 2, "maxItems": 2}, "description": "[[mw, price], ...] from max charge (negative MW) to max discharge; MW and price non-decreasing; 2-10 points"}},
  "required": ["hours"]}}

as_offers[].product is now {"type": "string", "enum": ["regup", "regdn", "rrspfr", "rrsffr", "ecrs", "nspin"]}.

Through the new factory, all 11 desk tools render as published and end_turn keeps its method schema.

Token cost: the system prompt with the desk's 12 tools goes from 3,241 to 3,573 Qwen3-4B tokens (+332). submit_dam_offers alone goes from 558 to 730.

Versions checked

• TRL 1.16.0 and 1.17.0: GRPOTrainer and AsyncGRPO's rollout worker find, render, parse and call tools the same way.
• transformers 5.21.0 (the lock), 5.22.0 and main: chat_template_utils is identical. The build-time check also passes on the 5.2.0 floor.
• tst/trl now pins that a real GRPOTrainer step shows the model the published nested schema and end_turn's current one. It passes locally with TRL 1.16.0 and 1.17.0, transformers 5.21.0, CPU torch 2.15.0 and Python 3.12. The tier's workflow runs the same matrix.

Tests

• tst/unit/training/tools_test.py:
  • a SchemaTool renders through apply_chat_template with Qwen3's chat template to exactly its schema;
  • parse_response, with TRL's Qwen3 response template, types a call's arguments from that schema;
  • the call goes through;
  • get_json_schema refuses the object.
  • Qwen3-4B's template is saved as tst/data/training/qwen3_chat_template.jinja, since CI has no tokenizer cache. A word-level tokenizer built in memory carries it.
• tst/unit/training/trl_test.py:
  • a tool with nested arrays, objects, an enum and bounds is served by build_app and fetched as TRL fetches it;
  • it renders and parses to exactly the published schema;
  • end_turn renders to its current schema;
  • the class's methods still render the method schemas;
  • calls reach the env and the rollout scores;
  • a transformers that renders a tool differently stops the factory.
• tst/unit/openenv/environment_test.py: the fake tool session _Tools can now offer whole EnvironmentTools, not only names. That is how the trl tests serve a nested inputSchema through the real server. Existing callers that pass names behave as before.
• make unit-test: 11,781 passed. Two tests fail on the Mac, and neither is affected by this change:
  • hermes/test_hermes_bridge.py::test_the_event_pipe_is_not_left_open_in_the_server reads /proc and fails on unmodified main too.
  • gateway/process_scope_socket_test.py::test_an_armed_time_trigger_keeps_firing_after_a_client_leaves passes when run alone.
• ruff check is clean, and uv lock --check passes.

Trainer scripts need no change to benefit: GRPOTrainer(environment_factory=env.factory, ...) and AsyncGRPO pick up the schemas as they are.

🤖 Generated with Claude Code

Confidence Score: 5/5

The PR appears safe to merge; no actionable issue was found.

Summary

SchemaTool lets TRL call each tool while giving transformers the full published input schema.

• TRL training prompts now show each tool's full nested input schema.

earakely-scale explicitly documents the class-level method-schema fallback as intentional. Tests were inspected, not run during this review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Published tool schema] --> B[Class tool method]
  A --> C[SchemaTool dictionary]
  B --> D[Bound instance method]
  D --> C
  C --> E[Transformers renders full schema]
  C --> F[TRL calls tool]
  F --> D
  D --> G[Environment tool call]

Loading

_{Reviews (1) · Last reviewed commit: "fix(training): TRL tools carry their MCP..." · Reviewed by Greptile}