[feature request] Channels-first environments without extra transpose wrapper

At present, image-based environments are automatically wrapped in a VecTranspose wrapper when they are passed to a policy. Likewise, observations passed to policy.predict() are automatically transposed so that the last channel appears first. This adds a bit of unnecessary overhead when working with channels-first environments: users need to wrap such environments in a special wrapper that puts channels last. The wrapper is then “undone” by VecTranspose, yielding the original (correct) observation. It would be nice if there was some way to support channels-first environments without two unnecessary transposes.

One way of fixing this would be to simply assume that environments are channels-first; environments that are not channels-first can be wrapped in a VecTranspose wrapper before being passed to an RL algorithm. This has a few advantages:

• It simplifies the behaviour of the code slightly, removing surprises like #103 and eliminating the need for transpose-aware checks like the one in common.utils.check_for_correct_spaces()
• It removes unnecessary overhead for channels-first environments
• It makes it easier for other libraries to interoperate with SB3 policies. For example, imitation trains SB3 policies using behavioural cloning, but that requires all of imitation to be aware of the heuristics that SB3 uses to decide whether to transpose images.

I’ve made a draft PR that does this in #143.

[araffin]

https://github.com/DLR-RM/stable-baselines3/pull/35#issuecomment-636331838

[Miffyli]

Generally I agree with this proposal: VecTranspose happens bit too often in bit too many places around the code to be a clean solution. However since channels-first is not a common convention (e.g. PIL and OpenCV treat images as channels-last), I think it will be a common mistake to feed in channels-last observations. The code should still sanity-check, even in BaseClass, if the first dimension is larger than last dimension. Normally this would be a check_env thing, but this could happen so stealthily I think it could be checked always. I will check the code more in depth next week when I am back on main machine.

@araffin I think we discussed this same at one point and you had some further arguments for the current behaviour?