Skip to content

[Performance] Add _skip_maybe_reset flag to bypass auto-reset in step_and_maybe_reset #3560

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
vmoens wants to merge 2 commits into from
Closed

[Performance] Add _skip_maybe_reset flag to bypass auto-reset in step_and_maybe_reset #3560

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aba1d20
Update
vmoens Mar 23, 2026
db0dc99
Update
vmoens Apr 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion torchrl/envs/common.py
View file
Open in desktop
Copy link
Copy Markdown
Collaborator Author

@vmoens vmoens Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't want that.
The proper way of doing auto-reset with torch compile should be to ALWAYS compute the reset and then mask reset and non-rest in tensordict_ using torch.where. We should have a toy auto-reset env that we can compile as an example and we should discuss how to make an extension point out of this but we should NOT skip maybe_reset entirely. I would rather make it a no-op if auto-reset is used with the masking I just talked about, or implement the masking in that maybe_reset which I would find more natural.

Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,9 @@ class EnvBase(nn.Module, metaclass=_EnvPostInit):
are to be expected. Can be ``None``.
is_spec_locked (bool): returns ``True`` if the specs are locked. See the :attr:`spec_locked`
argument above.
_skip_maybe_reset (bool): if ``True``, :meth:`step_and_maybe_reset` will skip calling
:meth:`maybe_reset` after each step. This is useful for auto-resetting environments
that already handle resets inside their :meth:`_step` method. Defaults to ``False``.

Methods:
step (TensorDictBase -> TensorDictBase): step in the environment
Expand Down Expand Up @@ -477,6 +480,7 @@ class EnvBase(nn.Module, metaclass=_EnvPostInit):
_batch_size: torch.Size | None
_device: torch.device | None
_is_spec_locked: bool = False
_skip_maybe_reset: bool = False

def __init__(
self,
Expand Down Expand Up @@ -3829,7 +3833,8 @@ def step_and_maybe_reset(
tensordict_ = self._step_mdp(tensordict)
# if self._post_step_mdp_hooks is not None:
# tensordict_ = self._post_step_mdp_hooks(tensordict_)
tensordict_ = self.maybe_reset(tensordict_)
if not self._skip_maybe_reset:
tensordict_ = self.maybe_reset(tensordict_)
return tensordict, tensordict_

# _post_step_mdp_hooks: Callable[[TensorDictBase], TensorDictBase] | None = None
Expand Down
Loading