Skip to content

[BugFix] Fix kwargs dropping in AsyncEnvPool.reset - #4521

Open
coder-jayp wants to merge 1 commit into
pytorch:mainfrom
coder-jayp:fix/envpool-kwargs
Open

coder-jayp wants to merge 1 commit into
pytorch:mainfrom
coder-jayp:fix/envpool-kwargs

Conversation

@coder-jayp

Copy link
Copy Markdown
Contributor

Description

This PR addresses the follow-up issue where AsyncEnvPool.reset drops **kwargs before they reach the sub-environments, which caused features like set_state=True to fail silently.

Fixes #4520

Changes:

  • Safely forwarded **kwargs through AsyncEnvPool.reset to async_reset_send.
  • Appended kwargs into the worker queue tuples inside _send_worker_batches.
  • Gracefully unpacked varying length tuples inside the multiprocessing _worker_exec and _env_exec (providing a {} fallback for legacy/internal messages like step or init_shm to guarantee zero regressions).
  • Passed kwargs directly into env.reset(**kwargs).
  • Replicated the identical logic for the ThreadPoolExecutor backend in ThreadingAsyncEnvPool.
  • Added a formal regression test test_async_pool_kwargs_forwarding to ensure kwargs are always forwarded.

Motivation and Context

The maintainer requested a proper forwarding mechanism in a follow-up PR to #4501. This implementation relies strictly on native tuple forwarding/unpacking to achieve cross-process boundary communication with absolutely zero API breakage or heavy abstraction changes.

@pytorch-bot

pytorch-bot Bot commented Oct 3, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⚠️ 16 Awaiting Approval

As of commit 2c8543e with merge base 06d57a0 (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

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 Oct 3, 2026
@github-actions

github-actions Bot commented Oct 3, 2026

Copy link
Copy Markdown
Contributor

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Fix kwargs dropping in AsyncEnvPool.reset

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[Algorithm] new algo [Algorithm] Add new RL objective
[BE] BE [BE] Improve error messages
[Benchmark] or [Benchmarks] Benchmarks [Benchmark] Add collector benchmark
[BugFix] BugFix [BugFix] Fix memory leak in collector
[Example] or [Examples] Examples [Example] Add training script
[Feature] Feature [Feature] Add new optimizer
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactoring [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Tests [Tests] Add unit tests for buffer
[Trainer] or [Trainers] Trainers [Trainer] Add trainer config
[Environment] or [Environments] Environments [Environments] Add Gymnasium support
[Data] Data [Data] Fix replay buffer sampling
[LLM] llm/ [LLM] Add reward model integration
[Minor] small change [Minor] Fix typo in error message
[Performance] or [Perf] Performance [Performance] Optimize tensor ops
[BC-Breaking] bc breaking [BC-Breaking] Remove deprecated API
[Deprecation] Deprecation [Deprecation] Mark old function
[Algorithm] or [Algorithms] new algo [Algorithm] Add new objective
[Quality] Quality [Quality] Fix typos and add codespell
[Versioning] versioning [Versioning] Bump release version
[WIP] WIP [WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

- Passed **kwargs through AsyncEnvPool.reset to async_reset_send.
- Updated _send_worker_batches to append kwargs to the worker queue requests.
- Updated _worker_exec and _env_exec to gracefully unpack tuple lengths (falling back to {} for old messages).
- Updated ThreadingAsyncEnvPool reset functions to pass kwargs.
- Added regression test test_async_pool_kwargs_forwarding to guarantee kwargs are not dropped.
@github-actions

github-actions Bot commented Oct 3, 2026

Copy link
Copy Markdown
Contributor

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Fix kwargs dropping in AsyncEnvPool.reset

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[Algorithm] new algo [Algorithm] Add new RL objective
[BE] BE [BE] Improve error messages
[Benchmark] or [Benchmarks] Benchmarks [Benchmark] Add collector benchmark
[BugFix] BugFix [BugFix] Fix memory leak in collector
[Example] or [Examples] Examples [Example] Add training script
[Feature] Feature [Feature] Add new optimizer
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactoring [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Tests [Tests] Add unit tests for buffer
[Trainer] or [Trainers] Trainers [Trainer] Add trainer config
[Environment] or [Environments] Environments [Environments] Add Gymnasium support
[Data] Data [Data] Fix replay buffer sampling
[LLM] llm/ [LLM] Add reward model integration
[Minor] small change [Minor] Fix typo in error message
[Performance] or [Perf] Performance [Performance] Optimize tensor ops
[BC-Breaking] bc breaking [BC-Breaking] Remove deprecated API
[Deprecation] Deprecation [Deprecation] Mark old function
[Algorithm] or [Algorithms] new algo [Algorithm] Add new objective
[Quality] Quality [Quality] Fix typos and add codespell
[Versioning] versioning [Versioning] Bump release version
[WIP] WIP [WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

@coder-jayp coder-jayp changed the title Fix kwargs dropping in AsyncEnvPool.reset [BugFix] Fix kwargs dropping in AsyncEnvPool.reset Oct 3, 2026
@github-actions github-actions Bot added the BugFix label Oct 3, 2026
@coder-jayp

Copy link
Copy Markdown
Contributor Author

@torchrlbot reviewer @theap06

@github-actions
github-actions Bot requested a review from theap06 October 3, 2026 16:09
@github-actions

github-actions Bot commented Oct 3, 2026

Copy link
Copy Markdown
Contributor

Requested review from @theap06 (requested by @coder-jayp).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

BugFix CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] AsyncEnvPool / ParallelEnv silently drops kwargs during reset()

1 participant