[Feature] step_and_maybe_reset in env

[vmoens]

Contribution

This PR proposes the step_and_maybe_reset in EnvBase.
This method executes a step followed by a reset, if necessary.

We also make reset more robust by ensuring that partial resets are handled uniformly. This is necessary since step_and_maybe_reset must take care of this functionality, and from our perspective handling partial resets is the responsibility of reset (the user should not have to worry about data not updated properly).

This has repercussions on the logic behind TransformedEnv._reset and BatchedEnv._reset.

I'm now considering having batched envs calling reset and not _reset to make sure that the data is well presented, since now the update of the input tensordict with the tensordict_reset occurs after _reset (hence, the output of _reset in SerialEnv is incomplete).
This could introduce some overhead but that's of limited impact since now step_and_maybe_reset is there to handle things faster.

[matteobettini]

This crashes when tensordict is None

[vmoens]

Yeah sorry about that, I can't test petting zoo locally so I'm always moving in the dark...

[vmoens]

@matteobettini @albertbou92 @BY571 this should be (almost) mergeable.
Tests are passing (VMAS sporadically failing after tests complete and old deps broken because of #1622)

[matteobettini]

Why would vmas sporadically fail?

[vmoens]

Why would vmas sporadically fail?

Not sure you can have a look.
All tests pass but it fails when closing.
Seems like an issue with tensors not released on CUDA, it happens often with VMAS. It isn't even a flaky test, since they all pass, it's just a weird exit status

[matteobettini]

Why would vmas sporadically fail?

Not sure you can have a look. All tests pass but it fails when closing. Seems like an issue with tensors not released on CUDA, it happens often with VMAS. It isn't even a flaky test, since they all pass, it's just a weird exit status

Was this happening before this PR?

[vmoens]

it happens often with vmas

I was referring to VMAS CI beyond this PR

[matteobettini]

Some final questions/comments

[matteobettini]

I am not following this.
We are taking the done keys (with all the terminated and truncated entries) and replacing a "done" ending for all and making it a set which we call reset keys.

This is counter intuitive as reset_keys have a _reset ending and not a done ending.
This change seems to come from the fact that you aim to use this function in 2 contexts:

• normally on the root td with the reset_keys as input
• on the "next" td in collectors with the done_keys as input

I think we should try to write this better, here are some suggestions:

• always call the function on the root td and pass the keys with preappended "next" if you want to use that
• do this key filtering and conversion outside of the function and let the function just operate on a set of keys to be cconsidered as reset keys

[vmoens]

I don't see the problem, can you elaborate why we should try to write this better?
Is it a naming problem? We can rename the function _aggregate_stop or smth similar.

always call the function on the root td and pass the keys with preappended "next" if you want to use that
That introduces some unwanted overhead when we can directly access "next" and read the done_keys. Recall that td.get(("next", "key")) is considerably slwoer than next_td.get("key") as we do here.

do this key filtering and conversion outside of the function and let the function just operate on a set of keys to be cconsidered as reset keys

What's your suggestion for _update_traj_ids in collectors.py for instance? We don't have a "_reset" in the "next" tensordict, but I think this function does its job of aggregating the done signals to read what the trajectory ids are.
What I understand is that the confusion comes from the "reset" in the function name, but what this function really does is just aggregating end-of-trajectory signals (either reset or done) to the root.
Given this, I don't see why it should be changed. It's a private function, properly tested and I think it serves its purpose.

[matteobettini]

connecting to the comment below, we could have

• _aggeregate_keys(keys=) which can be called on anything

alternatively we could have both _aggeregate_dones and _aggregate_resets where one calls the other

[matteobettini]

I assume we have to do this becuse there are no reset keys anymore visible by the collector.
Since this is quite counterintitive, what about a _aggregate_dones()?
I also have other suggestions in the other comment relating to this function.

[vmoens]

yep we can rename the function. I would rather say that it's more direct to aggregate the done rather than the "_reset" which come from the dones

[matteobettini]

Help me understand this a little bit better.

In an axample case (e.g., pettingzoo) where i have

{
"done": [False],
"agents":{"done:[True, False]}
}

Is the any_done triggered?

If so, this is a problem for envs like PettingZoo where _reset() will be called with {"_reset": [False]}

[vmoens]

The best thing it trying it out :)

from tensordict import TensorDict
from torchrl.envs.utils import _terminated_or_truncated

data = TensorDict({"done": [False], ("agent", "done"): [True, False]}, [])
print(_terminated_or_truncated(data))

which returns True
So what you're saying is that it should be False since there's a False at the root?
I can correct that

[matteobettini]

I think it should to follow the dominance rule we imposed right?
Or at least in this context definitely becuase we do not want to call reset.
I don't know in what other contexts this function is used though, but if its primary use is to decide when to call reset then yes

[vmoens]

@matteobettini I addressed all your comments

[matteobettini]

LGTM