[Feature Request] Allow subclassing ReplayBufferSamples and RolloutBufferSamples

[franco-ruggeri]

🚀 Feature

ReplayBufferSamples, RolloutBufferSamples, DictReplayBufferSamples, and DictRolloutBufferSamples are currently NamedTuples. As such, subclassing them is not supported. Converting them to dataclasses would support subclassing.

Motivation

Some RL algorithms require additional fields from the replay/rollout buffer.

An example is action masking. SB3 contrib implements MaskableRolloutBufferSamples but has to type ignore in several methods as a workaround for the unsupported subclassing of NamedTuples (see here). If RolloutBufferSamples were a dataclass, such a workaround would not be necessary, and the solution would be cleaner.

There are many other examples in novel research methods. This refactor would improve modularity and make it easier to implement new algorithms (e.g., in SB3 contrib).

Pitch

ReplayBufferSamples, RolloutBufferSamples, DictReplayBufferSamples, and DictRolloutBufferSamples should be refactored to be data classes.

Alternatives

Nothing comes to my mind other than dataclasses.

Additional context

No response

Checklist

• I have checked that there is no similar issue in the repo
• If I'm requesting a new feature, I have proposed alternatives

[araffin]

Hello,
thanks for the proposal.
I think I chose named tuples over dataclasses for two reasons (fb4e662):

1. Python 3.6 was still supported
2. it would allow to unpack things easily

So, now 1. is not needed and we need to check if 2. is actually used. If not, we can move to dataclasses indeed =)

[franco-ruggeri]

Hello, thanks for the response! Regarding the second point, data classes have several utilities: astuple() and asdict(). So I think it is possible to unpack them very conveniently.

[araffin]

In that case, feel free to submit a PR =) (note that you will need to do another one for SB3 contrib).

I also need to check once the PR is ready if it affects saved replay buffer (it should not but we never know).