I’ve opened a pull request to the original stable-baselines3-contrib repository, which is currently pending review by the maintainers.
In the meantime, the feature is already available via our fork under the NeuroGym organization: https://github.com/neurogym/stable-baselines3-contrib (branch: rnn_policy_addition), and can be installed directly for immediate use.
Remeber to update the note in the docs about such feature, when and if gets merged in the original repository.