DLR-RM / DLR-RM/stable-baselines3

[Feature Request] Allow subclassing ReplayBufferSamples and RolloutBufferSamples

Open
#2,202 3 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.