[Performance] Eliminate parent LIBERO environment construction

[vmoens]

Summary

• add an opt-in ParallelEnv(metadata_from_workers=True) startup handshake so real workers report and validate metadata without parent-side shadow environments
• use one homogeneous VLA-GRPO environment factory plus create_env_kwargs for worker-specific assignment, while keeping spawn as the configured worker start method
• close callable-created metadata environments on every path and bound/reap worker shutdown, including native renderer teardown
• add regression coverage, documentation, and a reusable startup benchmark

Root cause

VLA-GRPO built each LIBERO worker from a distinct partial, so ParallelEnv classified the batch as heterogeneous and constructed every environment serially in the parent for metadata. At 320 environments, that produced 320 temporary MuJoCo/EGL environments before constructing the 320 long-lived workers.

Benchmark

The committed synthetic benchmark (benchmarks/benchmark_parallel_env_startup.py) with 20 workers and a 10 ms constructor delay measured:

Mode	Parent constructions	Worker constructions	Closed environments	Ready time (construct + reset)
distinct factories, spawn	20	20	40	2.98 s
homogeneous factory, spawn	1	20	21	2.11 s
worker metadata, spawn	0	20	20	2.02 s
worker metadata, forkserver	0	20	20	1.82 s

The production LIBERO configuration (4 subcollectors x 80 environments on H200s) measured:

Mode	Parent metadata envs	Real workers ready	MultiCollector ready / capped observation	First step	Shutdown
distinct factories, spawn	4 started, 0 completed	0	capped at 300 s in the first serial metadata construction per parent	not reached	not reached
worker metadata, spawn	0	320 / 320	1235.5 s	1236.1 s	236.1 s max
worker metadata, forkserver	0	0 / 320	capped at 995 s; unsafe EGL device behavior	not reached	terminated

The successful spawn run retained all 10 task instructions per subcollector, the expected task/group offsets, worker indices, and seeds. Peak child-process counts were 81-82 per subcollector, the first environment step completed, and shutdown reported zero orphan worker processes. The forkserver run accumulated EGL allocations on GPU 0 rather than the requested per-subcollector render GPUs, so it is not enabled. Complete benchmark commands, logs, per-collector JSON, and process snapshots were retained with the run artifacts.

spawn remains the VLA default. forkserver is reported as benchmark evidence only rather than being selected from a single startup stress run.

Validation

• pytest -q test/envs/test_parallel.py
• pytest -q sota-implementations/vla_grpo/test_openvla.py -k 'factory or instruction or group'
• pytest -q test/envs/test_special.py -k callable_metadata_env
• pre-commit hooks for all changed files

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4028

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Reviewed the rebased head efac773, including the latest follow-up commit. The SerialEnv fallback now removes ParallelEnv-only arguments, worker construction failures preserve the full traceback, temporary metadata envs close on failure, and the configurable shutdown path reaps slow workers; these changes address the failure modes described in the updated commit and PR body. I ran the six metadata_from_workers tests plus test_callable_metadata_env_closes_when_extraction_fails locally: 7 passed. I did not find a blocking code issue in the current diff. Residual risk is the real LIBERO/OpenVLA and high-worker-count path, which I could not reproduce locally; the freshly triggered full CI is also still in progress.