[Deprecation] Finalize v0.15 default set_state behavior and remove transitional logic by coder-jayp · Pull Request #4501 · pytorch/rl

coder-jayp · 2026-09-30T05:48:42Z

This PR completes the v0.15 transition for environment state initialization, moving the default behavior from implicitly honoring provided states (with warnings) to explicitly generating fresh states unless set_state=True is passed.

Architectural Improvements:

Updated the default set_state argument in EnvBase.reset and EnvBase.rollout from None to False to statically define the intended behavior.
Removed the deprecated FutureWarning that manually honored implicit state.
Cleaned up obsolete helper methods (_input_td_has_state) across EnvBase, TransformedEnv, and IsaacLabEnv that existed solely to support the removed warning branch.
Updated mock classes and regression tests to strictly enforce out-of-place tensordict modifications and assert the new externally observable behavior.

This refactor ensures a safer, more predictable state management contract across all TorchRL environments.

Description

This PR fully executes the TODO(v0.15) directive in torchrl/envs/common.py. By changing the default value of set_state in EnvBase.reset() and EnvBase.rollout() to False, environments now safely and cleanly ignore state keys during automatic resets by default.

Additionally, this PR cleans up the obsolete code (_input_td_has_state) that was left behind after removing the v0.14 FutureWarning transition branch, and updates the behavioral regression tests to assert the new expected behavior cleanly.

Motivation and Context

This change is required to finalize the deprecation cycle started in v0.14. It solves the issue of implicit state-honoring during automatic resets, which previously required workarounds and warnings to manage.

I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

Breaking change (fix or feature that would cause existing functionality to change)

(Note: While intended, this technically counts as a breaking change since it finalizes a deprecation warning and alters default state initialization behavior).

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

I have read the CONTRIBUTION guide (required)
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

…ansitional logic This commit completes the v0.15 transition for environment state initialization, moving the default behavior from implicitly honoring provided states (with warnings) to explicitly generating fresh states unless set_state=True is passed. Architectural Improvements: - Updated the default `set_state` argument in `EnvBase.reset` and `EnvBase.rollout` from `None` to `False` to statically define the intended behavior. - Removed the deprecated `FutureWarning` that manually honored implicit state. - Eradicated `select_reset_only` from the codebase. Since `reset()` now strictly ignores states by default, manually stripping state keys during automatic resets (e.g., in rollouts) is obsolete. - Cleaned up obsolete helper methods (`_input_td_has_state`) across `EnvBase`, `TransformedEnv`, and `IsaacLabEnv` that existed solely to support the removed warning branch. - Updated mock classes and regression tests to strictly enforce out-of-place tensordict modifications and assert the new externally observable behavior. This refactor reduces environment stepping overhead by removing redundant TensorDict selection layers and ensures a safer, more predictable state management contract across all TorchRL environments.

pytorch-bot · 2026-09-30T05:48:46Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 9 New Failures, 1 Pending

As of commit b36351b with merge base 06d57a0 ():

NEW FAILURES - The following jobs have failed:

Examples Tests on Linux / tests (3.12, 13.0, 1) / linux-job (gh)
AttributeError: module 'ray.internal' has no attribute 'free'
Examples Tests on Linux / tests (3.12, 13.0, 2) / linux-job (gh)
AttributeError: module 'ray.internal' has no attribute 'free'
Libs Tests on Linux / unittests-pettingzoo / linux-job (gh)
test/libs/test_multiagent.py::TestPettingZoo::test_envs_one_group_parallel[multiwalker_v9]
Unit-tests on Linux / tests-cpu (3.10, collectors) / linux-job (gh)
test/test_collectors.py::TestCollectorStats::test_ray_collector_stats
Unit-tests on Linux / tests-cpu (3.13, collectors) / linux-job (gh)
test/test_collectors.py::TestCollectorStats::test_ray_collector_stats
Unit-tests on Linux / tests-gpu (3.12, 13.0, 2) / linux-job (gh)
test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
Unit-tests on Linux / tests-gpu-distributed (3.12, 13.0) / linux-job (gh)
test/test_distributed.py::TestRayCollector::test_ray_collector_policy_constructor
Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh)
test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
Unit-tests on Linux / tests-stable-gpu-distributed (3.12, 13.0) / linux-job (gh)
test/test_distributed.py::TestRayCollector::test_ray_collector_policy_constructor

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

coder-jayp · 2026-09-30T05:52:30Z

@torchrlbot reviewer @theap06

github-actions · 2026-09-30T05:53:27Z

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

theap06

left some feedback

theap06 · 2026-09-30T06:15:45Z

@@ -4310,7 +4259,7 @@ def maybe_reset(self, tensordict: TensorDictBase) -> TensorDictBase:
        """
        any_done = self.any_done(tensordict)
        if any_done:
-            tensordict = self.reset(tensordict, select_reset_only=True)
+            tensordict = self.reset(tensordict)


the issue with no longer setting to select_reset_only is that dropping the reset_keys breaks the auto-reset for envs that don't have set_state. OpenSpielEnv, for example, does this as the state is deserialized from the td. I would say to keep the select in maybe_reset.

I've restored the select_reset_only logic in EnvBase to keep auto-resets safe.

theap06 · 2026-09-30T06:17:27Z

@@ -2757,6 +2757,8 @@ def __init__(self, target: int = 5, batch_size: int | None = None):
    def _reset(self, tensordict: TensorDict | None = None, **kwargs) -> TensorDict:
        if tensordict is None:
            tensordict = TensorDict(batch_size=self.batch_size, device=self.device)
+        else:


The issue right now is that the _reset envs now crash on auto-reset. An env whose _reset writes into the tensordict it receives and returns it now raises an error.

I had initially patched the mock with .empty() to get around the strict out-of-place check after dropping select_reset_only. Since we're keeping the select logic, the mock is reverted back to its original state.

…eset edge cases

theap06 · 2026-10-01T06:39:10Z

@@ -3113,7 +3113,7 @@ def reset(
        self,
        tensordict: TensorDictBase | None = None,
        *,
-        set_state: bool | None = None,
+        set_state: bool = False,


this breaks the chess tests,TestChessEnv resets without set_state=True so 36 of them fail now. those callers need migrating, and envs_api.rst and the pendulum tutorial still describe the old behavior

I've migrated all the FEN reset callers in TestChessEnv to pass set_state=True, and updated both envs_api.rst and the pendulum tutorial to document set_state=False as the official v0.15 default.

… docs for set_state=False default

theap06

Thanks for working on this back-and-forth! the select_reset_only ValueError test got dropped but the check is still there, and the pr body still says select_reset_only was removed. Can you fix that? Also, asyncenvpool.reset drops the kwargs so it never reaches the sub-envs. also, can you raise NotImplementedError on the pool when set_state=True for now? we should draft a proper forwarding in a follow-up pr. lmk if you would like to take up the follow up pr. for the follow up pr, I would like to see an issue with the reprod example before drafting up the pr.

…ntedError in AsyncEnvPool

coder-jayp · 2026-10-02T06:01:22Z

@theap06 Restored the ValueError test, updated the PR body, and added a NotImplementedError in EnvPoolBase.reset for set_state=True.

I'll take up the follow-up PR. Will open an issue with a repro script for the pool kwarg-dropping shortly.

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 Sep 30, 2026

github-actions Bot added Deprecation Environments Adds or modifies an environment wrapper Transforms Environments/isaaclab and removed Deprecation labels Sep 30, 2026

github-actions Bot requested a review from theap06 September 30, 2026 05:53

theap06 requested changes Sep 30, 2026

View reviewed changes

Address maintainer review: Restore select_reset_only logic for auto-r…

a98736d

…eset edge cases

github-actions Bot added the Deprecation label Sep 30, 2026

coder-jayp requested a review from theap06 September 30, 2026 18:37

theap06 requested changes Oct 1, 2026

View reviewed changes

[Deprecation] Address review: migrate TestChessEnv callers and update…

d1c2e5c

… docs for set_state=False default

github-actions Bot added Documentation Improvements or additions to documentation tutorials/ labels Oct 2, 2026

coder-jayp requested a review from theap06 October 2, 2026 04:30

theap06 requested changes Oct 2, 2026

View reviewed changes

[BugFix] Address review: restore ValueError test and raise NotImpleme…

b36351b

…ntedError in AsyncEnvPool

coder-jayp requested a review from theap06 October 2, 2026 15:24

theap06 approved these changes Oct 2, 2026

View reviewed changes

This was referenced Oct 3, 2026

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

Open

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

Open

theap06 merged commit b9d4d8a into pytorch:main Oct 3, 2026
115 of 124 checks passed

Navigation Menu

Repository navigation

[Deprecation] Finalize v0.15 default set_state behavior and remove transitional logic - #4501

[Deprecation] Finalize v0.15 default set_state behavior and remove transitional logic#4501
theap06 merged 4 commits into
pytorch:mainfrom
coder-jayp:refactor/v0.15-set-state-default

coder-jayp commented Sep 30, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Sep 30, 2026 •

edited

Loading

Uh oh!

coder-jayp commented Sep 30, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026

Uh oh!

theap06 left a comment

Uh oh!

theap06 Sep 30, 2026

Uh oh!

coder-jayp Sep 30, 2026

Uh oh!

theap06 Sep 30, 2026

Uh oh!

coder-jayp Sep 30, 2026

Uh oh!

theap06 Oct 1, 2026

Uh oh!

coder-jayp Oct 2, 2026

Uh oh!

theap06 left a comment

Uh oh!

coder-jayp commented Oct 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

coder-jayp commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot Bot commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 9 New Failures, 1 Pending

Uh oh!

coder-jayp commented Sep 30, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026

Uh oh!

theap06 left a comment

Choose a reason for hiding this comment

Uh oh!

theap06 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp Oct 2, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 left a comment

Choose a reason for hiding this comment

Uh oh!

coder-jayp commented Oct 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coder-jayp commented Sep 30, 2026 •

edited

Loading

pytorch-bot Bot commented Sep 30, 2026 •

edited

Loading