[Feature Request] SB3 compatible version of Gym's FrameStack wrapper

🚀 Feature

A version of Gym’s FrameStack wrapper that returns the stacked observations as (channel*n_stack, height, width).

Motivation

The VecFrameStack in SB3 returns stacked observations as (channel*n_stack, height, width), while the non-vectorized version in Gym returns the observations as (n_stack, height, width, channel). Using the non-vectorized version with e.g. the NatureCnn feature extractor is not possible (fails assert in https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/torch_layers.py#L67).

Pitch

Either

• add a non-vectorized FrameStack wrapper (as has been done with Monitor), or
• add a compatibility layer to Gym’s FrameStack that assembles the the observations correctly.

### Checklist

• I have checked that there is no similar issue in the repo (required)

[araffin]

I just want to normalize the rewards.

You can skip obs normalization in that case by passing norm_obs=False (you may need SB3 master version), or as suggested in #693 invert the wrappers (that would be equivalent to doing frame stack on each single env and then normalizing)

[Miffyli]

Reasonable request, but in what situation would you want to apply FrameStack as an individual wrapper per env? On a quick thought this is useful if you want to add some wrappers after FrameStack to do some mods, but these could also be done in a vectorized fashion.

SB3 is built solely around “vectorized” envs as you can always turn single env into a vectorized env of size one, which simplifies the code. If the wrapper is not used in/for SB3 specifically, there is little reason to add it here tbh.