[Bug] Bug in StackedObservations for terminal_observation when using scalar observations

🐛 Bug

When using scalar observation spaces (Box with just one value), the update function of StackedObservations and StackedDictObservations determines an incorrect stack_ax_size, because the extra dimension is not present. The function assumes (num_envs, dim), but receives (num_envs,) because dim is 1.

In StackedObservations: https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/vec_env/stacked_observations.py#L120

In StackedDictObservations: https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/vec_env/stacked_observations.py#L232

To Reproduce

Code:

import numpy as np
import gym.spaces
from stable_baselines3.common.vec_env.stacked_observations import StackedObservations


num_envs = 4
n_stack = 3
obs_space = gym.spaces.Box(low=-1, high=1, shape=(1,), dtype=np.float32)

infos = [{"terminal_observation": 0}] * num_envs
dones = np.array([True] * num_envs)


stacked_observations = StackedObservations(
    num_envs=num_envs,
    n_stack=n_stack,
    observation_space=obs_space,
)

obs = np.zeros((num_envs,) + obs_space.shape)

stacked_observations.update(obs, dones=dones, infos=infos)

Error:

Traceback (most recent call last):
  File "******", line 24, in <module>
    stacked_observations.update(obs, dones=dones, infos=infos)
  File "******stable-baselines3/stable
_baselines3/common/vec_env/stacked_observations.py", line 132, in update
    new_terminal = np.concatenate(
  File "<__array_function__ internals>", line 180, in concatenate
ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 1 dimension(s) and the array at index 1 has 0 dimension(s)

Expected behavior

When calling update, stack_ax_size should be set to 1, but is set to num_envs due to the missing 1 in observations.shape.

### System Info

Describe the characteristic of your environment:

• Stable Baselines 3: 1.4.1a3 (cdaa9ab418aec18f41c7e8e12e0ad28f238553eb)
• Python version: 3.9.7
• Gym version: 0.21.0

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

[ScheiklP]

Oh damn, I thought it did, but there were too many moving parts. check_env correctly catches the incorrect shape. Sorry for the trouble…

[araffin]

Hello, thanks for spotting the bug, we would welcome a PR that fixes it 😉