[Performance] Avoid parent env construction in ParallelEnv startup

[vmoens]

Summary

• use one homogeneous VLA environment factory with per-worker create_env_kwargs, preserving task, instruction, group, worker-offset, seed, and init-state assignment
• add opt-in ParallelEnv(metadata_from_workers=True) startup so real workers provide metadata before initialization and the parent never creates shadow environments
• validate worker schemas, report startup failures, clean up partially started workers, and reject the unsupported shared-buffer combination explicitly
• always close temporary environments created by ordinary metadata callables while preserving EnvCreator and instantiated-environment ownership
• enable worker-originated metadata for the LIBERO recipe and document the generic API and production benchmark workflow

Why

A sequence of distinct per-worker callables makes ParallelEnv treat a homogeneous LIBERO workload as heterogeneous. Metadata discovery then constructs every environment serially in the subcollector parent before constructing the same environments again in the real workers. At 320 environments this creates 320 unnecessary MuJoCo/EGL environments and contaminates the parent with EGL state.

The worker metadata handshake removes that shadow wave while keeping the default ParallelEnv behavior unchanged. The initial implementation deliberately uses direct pipe communication; use_buffers=True fails clearly rather than selecting an unsafe path.

Benchmark

Production LIBERO was measured on the 5 subcollector x 64 environment topology. The verified worker-metadata/spawn run produced:

• parent metadata constructions: 0
• real worker constructions: 320
• all subcollectors ready: 367.93 s
• first reset and random step complete: 458.73 s
• peak descendants: 329
• outer exit codes: [0, 0, 0, 0, 0]
• remaining environment processes: none
• lingering worker processes: none
• instructions: 64 per subcollector with 8 distinct task instructions and the expected group-id blocks

A controlled one-subcollector/64-environment comparison on the same node was:

Mode	Parent metadata envs	Real envs	All ready	First batch
legacy independent callables, spawn	64	64	168.53 s	183.52 s
homogeneous factory, spawn	1	64	43.74 s	59.05 s
worker metadata, spawn	0	64	39.77 s	54.19 s
worker metadata, forkserver	0	64	39.42 s	54.29 s

Worker metadata reduced readiness latency by 4.24x and first-batch latency by 3.39x versus the legacy path. Forkserver was not materially faster than spawn in the paired run, so the recipe keeps spawn as the safer default. Outer subcollectors also remain on spawn.

Complete commands, output, construction markers, and JSON summaries are stored under:

/root/artifacts/libero_parallel_env_cold_start/5ab5a17f60/

The aggregate is comparison.json; each run directory contains command.txt, run.log, and summary.json.

Validation

• uv run pytest -q test/envs/test_parallel.py (83 passed, 66 skipped)
• uv run pytest -q sota-implementations/vla_grpo/test_openvla.py (14 passed, 22 skipped for unavailable optional dependencies)
• uv run pytest -q sota-implementations/vla_grpo/test_openvla.py -k 'factory or instruction or group' (10 passed)
• focused worker-metadata startup, failure, compatibility, and cleanup tests (7 passed)
• pre-commit hooks on all changed files
• production 320-environment LIBERO startup, first step, and clean shutdown

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 5ab5a17 with merge base 9f397fd ():
💚 Looks good so far! There are no failures yet. 💚

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

[theap06]

The ParallelEnv in batched_env throws a TypeError at the start. The metaclass converts single-worker ParallelEnvs to SerialEnv but forwards the new flag, which SerialEnv rejects.

[vmoens]

It's weird that our tests did not catch this, the SerialEnv fallback is used in the tests for sure no?

[theap06]

@vmoens The fallback is covered, but never together with the new flag: test_serial_for_single predates this PR. The flag only reaches SerialEnv on the num_workers==1 && serial_for_single=True early-return in _PEnvMeta.call, which pops serial_for_single but not metadata_from_workers before return SerialEnv(*args, **kwargs). so the failing line is never executed in CI. minor nitpick tbh