Recursivity is still not implemented here right?

@matteobettini I'm not even sure we need this function anymore. We wanted to avoid that people would have "done_agent1" and "done_agent2" and rather ("agent1", "done") IIRC.
Now we will allow "done_spec" to contain "done" and "truncated" at any level. Per se the done_spec could even be richer. So I don't think we need this check anymore: it's good practice to have only one done per level but I don't think that will break anywhere if you have more.

The only exception I see is computing the "_reset" signal which usually reflects done | truncated. In other words, _reset reflects somethings like

def _make_reset(data, done_spec):
    out = {}
    for key, item in done_spec.items():
        if isinstance(done_spec, CompositeSpec):
            out[key] = _make_reset(data.get(key), item)
        else:
            out["_reset"] = out.get("_reset", False) | data.get(key)
     return data.update(out)

This means tht if you put done_agent1 and done_agent2 at the same level you'll get a single reset for both. Not ideal but not really our business if the MARL API is well documented IMO.

In collectors we use "done" to compute the trajectory ids but this is an extra feature that could apply only in settings where done is at the root (or in other restricted settings). The trajectory id doesn't influence the collector behaviour, which only cares about the number of frames (number of elements of the root tensordict).

For the rest, loss modules have parametric dones so I think we could rename "done" to "merylstreep" and they could still work. They will need refactoring to account for "truncated" anyway.

So the situation is a little more complex, let me try to explain it here:

Since we decided that our naming convention is to have one "_reset" key per tensordict (due to the fact that if we had more we would have to change the name), then this key has to reflect the shape of the "done" in that tensordict.

If we allowed more then one leaf key in each td of the done_spec then the problem is what shape should the _reset reflect. We cannot do like vincent said cause the shapes might be arbitrary.

Therefore the convention established is that at least a key with the name "done" (and exactly done) should exist in each td of the done_spec and the reset will mimic its shape.

If now we want to change the convention and say that the constraint now is that all the values in the same tensordict of the done_spec should have the same shape, that is fine by me, but then we have to check and enforce that.

if we do like in this PR and assume to always have truncated in the done_spec, then we have to carefully adjust all the components of the codebase that create the _reset key.

Like here

and here

For example.

These lines will currently create a _reset for each done_spec.keys(True,True).
But if the done_spec now contains more than one leaf per td we have to change all this.

If now we want to change the convention and say that the constraint now is that all the values in the same tensordict of the done_spec should have the same shape, that is fine by me, but then we have to check and enforce that.

Per se, if someone writes an env that does not require _reset (which we use internally but is not exposed to users) it's fine. What I mean is that if a data structure has "done_agent1" of shape X and "done_agent2" of shape Y, computing a "_reset" key will fail but that only means that the env cannot be used within a collector.
Since we want to provide independent, modular components, I think it's a design choice that is ok not to enforce that all envs written by any user would fit in a collector. We need to make sure that the collector tells you specifically what is wrong with your env though.

What about letting envs have the done_spec they want, and let the collectors / parallel envs spot the potential problems during construction?

Yeah but also the rollout function uses the reset logic. So how would we treat that?

Imagine someone wants to use TorchRL to write a multiagent env that does some job with the "done_agentX" structure.
I think it's cool if they do that and put that in their app. I don't think they should get an exception or a warning that tells them that they're going to be in a dead end when using a collector. Maybe they don't want to use the collectors.
Let's not prevent legitimate use of one class to prevent bug in another if the usage of the other is optional, wdyt?

Yes but vectorized envs will need to be fed a _reset flag upon reset in the env.rollout function. so you are saying we will provide that with our best effort but if we cannot we just give a warning and say that in collector this will be an error?

Yeah but also the rollout function uses the reset logic. So how would we treat that?

Oh yeah good point.
This is a bit tricky: per se, we set a _reset in rollout and hope for the best but we don't enforce that it's being read anywhere.
The thing is that we don't really have any agency about what reset(tensordict) does under the hood. For ParallelEnv and similar we know what to do but in other cases I'm not 100% sure of what could be done generically.

But yeah, since we use _reset within the env class it makes sense that it's computed in a consistent way within the library.
So what we would opt for is:

• during construction, walk through the specs and check that at each level, all done specs have a matching shape
• during execution, compute _reset as the combination of all _reset keys at each level

during execution, compute _reset as the combination of all _reset keys at each level

-> as the combination of all done_spec_keys at each level

but yes that could work!

The method would benefit from a quick docstring or better name? I'm not sure to understand what it does without more context.

This doesn't seem to match its current output (tuple of 3 elements). Missing the truncated output?

yes I still need to go through the docstrings but thanks for reminding me!

If done is exactly terminated then what does "we do not support the terminated signal" mean?

I agree it's confusing. Maybe

In torchrl, a `done` signal means that a trajectory is terminated (what is referred to as `termination` in gymnasium).

Yep! That's clear!

This means eventual bootstrapping will be done correctly for new versions of gym and incorrectly for old versions (as when truncated they set done to true and specify their truncation bit in the info variable). A note should be added about this somewhere.

Let me rephrase what you're saying here:
if older version of gym (<0.26) a truncation signal could be retrieved from the info.
Currently, we don't use it so assuming that done=truncated will not work properly with value functions, is that right?

If so, I think that we should advise users of adopting a dedicated info reader to write the truncated signal, which will solve the problem.

if older version of gym (<0.26) a truncation signal could be retrieved from the info.

Yes in the TimeLimit.truncated key if it happened. See https://gymnasium.farama.org/tutorials/gymnasium_basics/handling_time_limits/#solution.

assuming that done=truncated will not work properly with value functions

Yes, and to be exact done here with the old gym versions would be gym_done=terminated | truncated. (It was an end_of_traj)

I think that we should advise users of adopting a dedicated info reader to write the truncated signal, which will solve the problem.

Yes, that would be the correct way to record the truncated key and propagate it to the rest of TorchRL's components. However, the gym_done key would still need to be corrected with torchrl_done = gym_done and ~truncated I don't how this could be done by the users.

Realizing once more that variable order can so easily be messed up here. Shout out to Tensordict for solving this!

Maybe the variables in this file referring to the gym done can be renamed to gym_done to avoid confusion with the torchrl done?

good point, it's getting a bit confusing!

@matteobettini

Instead of adding teminated and keeping done, why don't we just use terminated instead of done?

because we want done to be able to quickly know if we should reset. It avoids always going through all the possible keys and carries the summary of the end-of-traj in one tensor.
So now all envs have at least one done and one terminated, possibly a truncated.
It's what we discussed with @skandermoalla IIRC.

but the done would be added automatically by EnvBase no?
i am talking about this mocking class, it could just output terminated and done would be built

Not sure if my input is needed here, but here's how I understand this:

So now all envs have at least one done and one terminated, possibly a truncated.

Yes, this is what is expected.

but the done would be added automatically by EnvBase no?
i am talking about this mocking class, it could just output terminated and done would be built

Yes, the done key can be inferred from the terminated (possibly truncated keys). I haven't looked at the implementation yet to know where this is done, but if the EnvBase takes care of this, then this mock can avoid writing the done key indeed to test the automatic done inference feature.

I think it's best if we show both whenever we can. We can copy the done if you prefer but not showing that it's populated isn't great from a UX perspective IMO
The auto patch is supposedly just there to ensure that everything runs fine even if someone just writes a terminated/done key

If this is our rule then we should apply it to all envs.

let’s just do it uniformily

just did, what did I miss?

the reset key should not be removed here

why? Can you explain? It used to be removed (see diff) and i'm pretty sure tests fail if you don't.

there are many reasons we discussed

• they should alwayes be removed by env.reset() for uniformity
• this removes them only when self._single_task and not otherwise

ok let me try

right so the problem is this:
if there is no reset in the input we assume it's all true.
But the buffer is not necessarily all true (it's more likely that it is all false). So if there's a _reset in the buffer, we don't want to pass it in the result of reset because its value isn't synced.

Now you could argue that we must then update the _reset in the buffer with the one that we got from the input or that we inferred but then the risk is that these _resets do not have the same shape.

Transforms such as reward sum look at the tensordict that resuts from _reset before it is pruned from its "_reset" and decide what to do from there.

So here are the cases we need to account for:

1. the input tensordict does not have a "_reset" (or there is no input tensordict) => we infer that "_reset" is full of True
2. the input tensordict has a reset but its shape mismatch from the parallel env. This can happen if the "_reset" in the buffer has the shape of the "done" which, like you mentioned on some occasions, can have more dims than 1, but the user just passed a "_reset" that matched the number of envs (1-dim). In that case we could update the buffer by extending on the right I guess, but then the problem is if the env on the other end (in the subprocess) reads that "_reset". Then there's a chance it will be corrupt.
3. All is good and the "_reset" match.

I think case 2 is hard to handle. The test pass if we exclude the "_reset". And we update the tensordict that results from ParallelEnv._reset in TransformedEnv._reset so the transforms can see the partial "_reset" if there is.
All in all I don't see why not exluding the "_reset" in ParallelEnv._reset: the only method that will look for a "_reset" are TransformedEnv and EnvBase in their _reset / reset methods, respectively. I don't see anywhere where that will break.
If the point is that it should also be excluded inf the case where there is more than 1 task I totally agree!

``> If the point is that it should also be excluded inf the case where there is more than 1 task I totally agree!

definitely let's uniform this.

my other point is just that i thought we decided that the convention was that all the _resets are removed only in env.reset() and in this particular instance they are removed bofore. this has 2 problems:

• env.reset() cannot perform the done check appropriately (it will see a missing _reset and assume it is all true but that is not the case)(this is a problem if the reset was partial and some dones are true after the reset as they should but they cause the check to fail)
• and that this is not uniform as the batched env _reset() will be the only one to not provide _resets (and this will make env.reset() not perform its checks properly)

this was not a problem before as in env.reset() i was chaching the _resets in the input (before any other call) to check the dones

env.reset() cannot perform the done check appropriately (it will see a missing _reset and assume it is all true but that is not the case)(this is a problem if the reset was partial and some dones are true after the reset as they should but they cause the check to fail)

no because env.reset gets its "_reset" from the input tensordict, not the output

and that this is not uniform as the batched env _reset() will be the only one to not provide _resets

Same I don't think it is true: per se all _reset must return a new tensordict and it's written nowhere in the contract of that function that you must copy any part of the content of the input in the output. In other words: we never read "_reset" from the output of _reset. Check the reset method for more context if need reassurance.

ok got it, my bad.

let's just test that BatchedEnv._reset() does not return _resets in all branches