DLR-RM / stable-baselines3

PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms.
https://stable-baselines3.readthedocs.io
MIT License
9.17k stars 1.71k forks source link

[Feature Request] Allow render in venv to return 4D array instead of tiling #1155

Open Rocamonde opened 2 years ago

Rocamonde commented 2 years ago

🚀 Feature

Support for not tiling images in vectorized environments and instead returning a 4-dimensional array.

Motivation

Untiling the image is fairly non-trivial (the tiling code is 10-15 lines long) and an additional overhead if people just want to keep the images not tiled. This is useful if one wants to record separate videos or process the images of the vectorized environments as different independent samples for experiments.

Pitch

Allow a kwarg in the render method or in the constructor that checks whether the user would like to receive a tiled image. By default, tile the image, so the feature is backwards compatible. If a flag is passed, return a 4-D array instead.

I can propose a PR for this (I made a solution already on my fork) but first we should decide on the API: is the flag passed to render or to init?

araffin commented 2 years ago

is the flag passed to render or to init?

this won't break with gym 0.26, right? I would go for a flag passed to the render method as it allows more flexibility (and allows to have both modes).

Rocamonde commented 2 years ago

I guess it's about whether we want to make that API design choice in coherence with the new changes coming in with #780. In principle nothing should break since the flag is handled internally by the wrapper, not passed down.

So what option do people prefer?