I think adding the [feature extractors](https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/torch_layers.py) from SB3 is a good next step towards feature parity between the libraries