[earakely-scale]

What

AgentEnv.factory's classes now pickle, so TRL's AsyncGRPO can use them.

AsyncGRPO pickles environment_factory into the rollout process it spawns. The factory's classes are built at runtime, so pickle couldn't find them by name and trainer.train() failed:

_pickle.PicklingError: Can't pickle <class 'agent_env.training.trl.dock'>: it's not found as agent_env.training.trl.dock

• How: the classes have their own metaclass, _EnvironmentClass, registered with copyreg. A class pickles as the toolset (the listed tools) and the server it was built from. The spawned process rebuilds it from those, without probing the server again.
• Subclasses: a subclass of a factory class pickles by name, as any class does. A runtime subclass that can't be found by name fails to pickle, rather than silently losing what it overrides.
• Closing: a class rebuilt in the process that holds its AgentEnv is closed by AgentEnv.close() like the original. Each toolset's shared settings carry a key that survives pickling, so a rebuilt class matches the AgentEnv it came from, and only that one.
• Unchanged: GRPOTrainer, which calls the factory in process.

Tests

• New: tst/unit/training/trl_test.py::test_the_factory_pickles_and_plays_in_a_spawned_process_as_asyncgrpos_rollout_process_does. It does what AsyncGRPO's rollout process does: it builds an environment from the factory in a spawned process and plays and scores a rollout. It also checks that a runtime subclass fails to pickle by name. Run against main, it fails with the PicklingError above.
• Extended: the multi-environment test plays its pantry rollout from an unpickled factory dict.

Verified on GPUs

Both trainers ran on the PortSim training pack (portsim-dock-v1) served by build_app on one 4× L40S box, with Qwen3-4B (LoRA r=16) and the framework at this branch's head.

AsyncGRPO (vLLM server on GPU 0, trainer on GPU 1):

• 5 steps of 16 rollouts in 391 s, with 32 rollouts in flight;
• every step's reward came from the env (rewards/portsim-dock-v1);
• about 8 tool calls per rollout, with no tool failures and nothing unscored;
• LoRA adapter syncs took about 0.4 s;
• 2 of the 5 steps had a non-zero gradient.

GRPOTrainer (vLLM colocated on GPU 2):

• 2 steps of 16 rollouts plus an 18-rollout eval: 50 rollouts, none unscored.

Teardown: no episode containers were left after either run ended normally. When the trainer was killed mid-run, the server tore down its 28 live episodes within 30 s.

Install: installing agentenv-framework[trl] on top of trl[vllm,peft,kernels]==1.16.0, vllm==0.29.* and transformers==5.17.* leaves trl, vllm, transformers, torch, peft, accelerate and datasets where they were. It does move four other packages down:

Package	Before	After	Capped by
mcp	2.3.0	1.30.0	agent-env and OpenEnv's FastMCP 3
openai	3.40.0	2.54.0	litellm (every release caps it below 3)
pillow	13.0.0	12.3.0	gradio, through OpenEnv
tomlkit	0.15.2	0.14.0	gradio, through OpenEnv

vLLM 0.29 runs with all four.

What the docs should say about AsyncGRPO

Found while running it:

• The main guard: the training script must run under if __name__ == "__main__":, since the rollout process imports it again.
• GPUs: AsyncGRPO's packed attention loads the kernels-community/flash-attn3 Hub kernel. It's built for compute capability 8.0 and 9.0a, so it runs on A100, L40S, A10G, L4 and H100, but not on Blackwell (RTX PRO 6000, B200).
• Context length: AsyncGRPO 1.16 asks vLLM for max_completion_length new tokens on every turn, however long the conversation has grown. A multi-turn rollout that passes max-model-len minus that gets a 400, which it retries 30 times before failing the run. So the server's --max-model-len has to cover the longest conversation plus max_completion_length. For PortSim's dock tasks, --max-model-len 32768 with max_completion_length=2048 and max_tool_calling_iterations=8 worked.

🤖 Generated with Claude Code

Confidence Score: 5/5

The PR appears safe to merge; the previous cleanup issue is fixed and no new blocking issues were found.

Summary

Makes AgentEnv.factory classes pickleable for AsyncGRPO’s spawned rollout process.

• AsyncGRPO rollouts can now rebuild the environment classes made by AgentEnv.factory.

The previous cleanup finding is fixed. No new actionable issues were found. Tests were inspected, not run.

_{Reviews (2) · Last reviewed commit: "fix(training): AgentEnv.close also close..." · Reviewed by Greptile}