Repository navigation
[BugFix] Fix kwargs dropping in AsyncEnvPool.reset - #4521
coder-jayp wants to merge 1 commit into
Conversation
🔗 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.
|
|
| 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.
92e5cad to
2c8543e
Compare
|
| 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]).
|
@torchrlbot reviewer @theap06 |
|
Requested review from @theap06 (requested by @coder-jayp). |
Description
This PR addresses the follow-up issue where
AsyncEnvPool.resetdrops**kwargsbefore they reach the sub-environments, which caused features likeset_state=Trueto fail silently.Fixes #4520
Changes:
**kwargsthroughAsyncEnvPool.resettoasync_reset_send.kwargsinto the worker queue tuples inside_send_worker_batches._worker_execand_env_exec(providing a{}fallback for legacy/internal messages likesteporinit_shmto guarantee zero regressions).env.reset(**kwargs).ThreadPoolExecutorbackend inThreadingAsyncEnvPool.test_async_pool_kwargs_forwardingto 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.