[earakely-scale]

What

Two fixes in the episode path that agent_env.training.trl trains through. Both showed up training Qwen3-4B on the GridSim training bundle (the run behind the docs' TRL page).

1. A scored episode answers at once. A rollout's reward came back about 2.4 s after its last end_turn. TRL 1.16's GRPOTrainer calls get_reward() on each environment in turn, so on that run it was about 45 s of every 190 s step (32 rollouts). Two waits made it up:

Wait	Before	Now
prompt_agent notices the turn has ended	at its next poll of the proxy, every poll_interval_seconds (2 s by default)	poll_a2a_task takes a wake from whoever ends the task in this process (wake_poll), and an episode's Turn.end() / _fail() wakes it. A woken poll asks again every 50 ms for up to 2 s, since the proxy shows the task as ended a moment (~20 ms) after it accepts the reply; a wake that comes before the poll begins is kept, and the poll takes it up as it starts.
The closing end_turn answers	after the episode's teardown (Episode.result() waits for the whole run)	Episode.result(teardown=False) returns once the run's steps have finished, next_turn() returns None then too, and the OpenEnv env answers with the scores and lets the teardown finish behind (_abandon(..., wait=False)), as a call past its deadline already did. The next reset or close of the session, and the server's shutdown, wait for it.

2. A role-less agent's toolset keys the default role's access. ServedTask.toolset kept a modify_env_tool_access only when its role equalled the played agent's. An agent deployed with no role gets the gateway's default, so a change on default fell out of its key: tasks that differ only by it shared an environment class, and the first rollout that saw other tools than the class's would stop training. The GridSim bundle's Task API showed "access": [] for tasks that all disable two tools on default.

Measured

get_reward() on the GridSim bundle served by AgentEnv.local on local Docker (Mac), per rollout:

0.9.1380	2.45 s, 2.62 s
this branch	0.03–0.04 s, over six rollouts

Played the way GRPOTrainer plays them (each environment reset in turn, a scripted tool call each, then get_reward() each in turn), three rounds:

Environments	Collecting a round's rewards, 0.9.1380	This branch	Containers left after close()
16	24.7–29.6 s	2.3–2.6 s	0 and 0
32	11.5–23.3 s	1.3–2.6 s	0 and 0

Between rounds, 1–2 teardowns are still finishing on this branch (0 on 0.9.1380), so they don't pile up: each session's next reset waits for its own. At 32, Docker Desktop's 4 CPUs on the Mac can't start 32 GridSim envs at once, and both versions leave rollouts unscored at apply_server_config (ReadTimeout): 64 of 96 on this branch, 75 of 96 on 0.9.1380. A native-Docker box with 4 vCPUs ran 32 at once without that.

Behaviour changes

• Release note: env.factory can become a dict for some bundles. A bundle whose played agent has no role, and whose tasks differ in a modify_env_tool_access on default, now serves a toolset per variant, so env.factory is a dict of classes (TRL's multi-environment mode) where it was one class. Before, those tasks shared a class whose tools didn't match some of them.

• A teardown that raises after the episode is scored no longer turns the closing end_turn into an error observation; the reward stands and the session is freed. teardown_run never raises (it is best-effort: failures are logged and reported as still_up), so this changes only what a fake that raises produces; the test that pinned it is renamed and asserts the new outcome.

• Episode.result(teardown=...) is new and keyword-only, defaulting to the old behaviour; next_turn() now returns None once the steps have finished rather than once the teardown has.

Tests

• protocol_test: a woken poll asks at once and again until the task shows as ended, and leaves nothing registered; the fake clock stands in for _nap.
• run_test: a turn with a 30 s poll interval carries on within 2 s of ending; during a held teardown result(teardown=False) returns the scores and the episode isn't finished.
• environment_test: with the teardown held, end_turn answers with the reward and the teardown finishes after; a role-less agent's toolset carries the access change on default; a teardown that raises after scoring leaves the reward.
• trl_test: counts the teardowns after close(), which waits for them.
• make unit-test: 11,724 passed; two packages/agentenv-protocol tests failed under -n auto (test_a2a_proxy::…stops_what_it_left, test_computer_queue::…release_timer) and pass alone on this branch and on main. run_test::test_a_turn_not_taken_by_its_deadline_fails_the_run… is flaky on main as well (the in-process proxy takes ~1.7 s to start against a 5 s wait); left as is.

🤖 Generated with Claude Code

Confidence Score: 5/5

The PR appears safe to merge, with the previously reported non-blocking memory leak still open.

Findings

1.  Late replies retain memory ▶

Fix with agent prompt

### Issue 1
src/agent_env/a2a_agent/protocol.py:undefined-178
`wake_poll()` now keeps an event even when nobody will poll that task again. If a caller ends a timed-out turn before the episode notices it has closed, `ProxyClient.reply()` returns `False`, but `Turn.end()` still calls `wake_poll()`. The finished poll has already removed its entry, so this creates an entry that nothing removes.

Repeated late replies grow `_polls` for the life of the server. Bound how long pending wakes live, or remove them when their episode closes, while preserving wakes that arrive before polling starts.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

Episodes return scores before teardown finishes, and role-less agents include default access changes in their toolset keys.

• Ending a turn wakes its episode run instead of waiting for the next poll.
• OpenEnv returns an episode's reward while teardown finishes in the background.
• Role-less tasks include default-role tool access when TRL groups their tools.

Diagram

sequenceDiagram
    participant Client
    participant Episode
    participant Poll
    participant Cleanup
    Client->>Episode: End turn
    Episode->>Poll: Wake poll
    Poll-->>Episode: Task completed
    Episode->>Cleanup: Start teardown
    Episode-->>Client: Return scores
    Client->>Episode: Reset or close
    Episode->>Cleanup: Wait for teardown

Loading

_{Reviews (3) · Last reviewed commit: "test(openenv): wait for the sessions' te..." · Reviewed by Greptile}

[greptile-apps]

Late replies retain memory

wake_poll() now keeps an event even when nobody will poll that task again. If a caller ends a timed-out turn before the episode notices it has closed, ProxyClient.reply() returns False, but Turn.end() still calls wake_poll(). The finished poll has already removed its entry, so this creates an entry that nothing removes.

Repeated late replies grow _polls for the life of the server. Bound how long pending wakes live, or remove them when their episode closes, while preserving wakes that arrive before polling starts.

Prompt To Fix With AI

This is a comment left during a code review.
Path: src/agent_env/a2a_agent/protocol.py
Line: 178

Comment:
**Late replies retain memory**

`wake_poll()` now keeps an event even when nobody will poll that task again. If a caller ends a timed-out turn before the episode notices it has closed, `ProxyClient.reply()` returns `False`, but `Turn.end()` still calls `wake_poll()`. The finished poll has already removed its entry, so this creates an entry that nothing removes.

Repeated late replies grow `_polls` for the life of the server. Bound how long pending wakes live, or remove them when their episode closes, while preserving wakes that arrive before polling starts.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.