Skip to content

[Performance] Eliminate parent LIBERO environment construction - #4028

Open
vmoens wants to merge 2 commits into
pytorch:mainfrom
vmoens:codex/libero-parallel-env-cold-start-single-pr
Open

[Performance] Eliminate parent LIBERO environment construction#4028
vmoens wants to merge 2 commits into
pytorch:mainfrom
vmoens:codex/libero-parallel-env-cold-start-single-pr

Conversation

@vmoens

@vmoens vmoens commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

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

pytorch-bot Bot commented Jul 16, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

✅ No Failures

As of commit 35911c6 with merge base f2a2337 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 16, 2026
@github-actions github-actions Bot added Documentation Improvements or additions to documentation Benchmarks rl/benchmark changes sota-implementations/ Performance Performance issue or suggestion for improvement labels Jul 16, 2026
@vmoens
vmoens marked this pull request as ready for review July 17, 2026 08:26
…tracebacks

- Drop ParallelEnv-only kwargs (mp_start_method, metadata_from_workers) when
  serial_for_single falls back to SerialEnv instead of raising a TypeError,
  and retain mp_start_method when the env stays parallel
- Expose the worker shutdown grace period as a shutdown_timeout constructor
  argument (also settable via configure_parallel) instead of a hardcoded 10s
- Send the full worker traceback in metadata_error messages so construction
  failures are debuggable from the parent
- Close transient metadata envs with raise_if_closed=False so a close failure
  cannot mask the metadata result or the original error
- Document the identical-schema requirement and sequential shutdown of
  metadata_from_workers in the docstring and the vectorized-envs docs

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Benchmarks rl/benchmark changes CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Documentation Improvements or additions to documentation Performance Performance issue or suggestion for improvement sota-implementations/

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant