[BugFix] Make TransformedEnv mirror allow_done_after_reset property of base env

[matteobettini]

Fixes #1806

The context is:

Some enviornments can be done immediately after a reset (e.g., some agents can be still dead after a reset in multiagent settings).
TorchRL by default does not allow this and will throw an error.
If you, however, would like to allow your environment to be done after reset, you can set allow_done_after_reset during construction of your environment.
PettingZoo and VMAS need this.

However, in #1806, it is shown that when a transform is added to the environment, this attribute changes.

In this PR, we fix this, by making the attribute of transformend enviornments mirror the one of the base environment.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[vmoens]

This is untested and I'm not sure where it would fail in practice, can you add a test (and some more context?)
As of now I think we should rather have a property that falls back onto the env_base attribute, such that changing the env_base changes the transformed env. That will avoid silent bugs down the line if things are modified dynamically (ie, you won't be able to modify the transformed env and if we end up in that situation we will deal with it accordingly)

[matteobettini]

Yes I can do that.

[matteobettini]

What I did is I now treat it exactly like env._run_type_checks

I made a public setter and getter for it. The getter of the transform calls the one of the base env and the setter throws an error.

[vmoens]

why do we want this to be public? We should be super conservative with public attributes because once it's there we can't take it away!

[matteobettini]

So basically if I did not make it public, implementing the getter and setter solution would not be possible.

It is quite convoluted, but let me explain.

imagine i implement private getter and setters for TransformedEnv with the setter throwing the error. Now, when in the constructor we call the super().__init__()

rl/torchrl/envs/transforms/transforms.py

Line 575 in baea10b

super().__init__(device=None, **kwargs)

that will try to set env._allow_done_on_reset which, however, now calls the setter of the TransofrmedEnv, throwing an error.

[matteobettini]

This is the reason behind my first proposed solution.

[matteobettini]

A possible solution is that the setter of the transformed env does not throw an error but sets the variable in the base env. It seems to me that you wanted to avoid this though

[vmoens]

Why not: set it to None, get it, if the value is not None throw an error, otherwise set the private attribute.
Setting the private attribute will be ok as long as it's None (the setter of this property will be a no-op that passes silently if the value is None and throws an exception otherwise).
The getter on the other hand gets it from the env.

Like this we don't change anything, have the error where appropriate and keep the attribute private.

[matteobettini]

What i understand from the second sentence might work.

i am not super able to understand what you mean in the first sentence.

but i can make the setter a no-op for None values if that is what you are suggesting, that might work. (Maybe just a bit confusing imo)

[matteobettini]

Implemented, let me know if this is what you meant

[vmoens]

LGTM