DLR-RM / DLR-RM/stable-baselines3
[Feature Request] Allow subclassing ReplayBufferSamples and RolloutBufferSamples
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### 🚀 Feature
`ReplayBufferSamples`, `RolloutBufferSamples`, `DictReplayBufferSamples`, and `DictRolloutBufferSamples` are currently `NamedTuple`s. As such, subclassing them is not supported. Converting them to [`dataclass`es](https://docs.python.org/3/library/dataclasses.html) 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 `NamedTuple`s (see [here](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/blob/master/sb3_contrib/common/maskable/buffers.py#L88)). 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
- [x] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo
- [x] If I'm requesting a new feature, I have proposed alternatives
Contributor guide
Assessment
This issue has not been assessed yet.