DLR-RM / DLR-RM/stable-baselines3

[Bug]: `is_image_space` works poorly with Gymnasium's `FrameStackObservation`

Open
#2,090 6 comments 0 reactions 0 assignees View on GitHub
bug documentation help wanted
Dominant language
Python
Stars
13.8k
Forks
2.2k
Avg merge
1h 35m
Merged PRs (30d)
2

Description

### 🐛 Bug

MRC 👇
```python
from stable_baselines3.common.preprocessing import is_image_space
import gymnasium as gym
import numpy as np

from gymnasium.wrappers import FrameStackObservation

image_space = gym.spaces.Box(0, 255, (3, 64, 64), np.uint8) # a basic RGB 64x64 image

class DummyEnv(gym.Env):
def __init__(self):
super().__init__()
self.observation_space = image_space

env = DummyEnv()
print(is_image_space(env.observation_space)) # True

env = FrameStackObservation(env, stack_size=2)

print(env.observation_space.shape) # (2, 3, 64, 64), stacked obs on dim=0
print(is_image_space(env.observation_space)) # False
```
From [here](https://github.com/DLR-RM/stable-baselines3/blob/fa21bce04ee625c67f6ea2a7678bf46c39cd226c/stable_baselines3/common/preprocessing.py#L49) it seems to be due to the fact that sb3 expects images to be tensors of dimension strictly equal to 3.

I am wondering why not setting the check to be `>=3` instead of strictly equal. In this way, one would still pass the image check (though I reckon might have problems with `NatureCNN` as I am not sure how it would handle 4-dimensional tensors).

Worth adding that:
- A workaround this bug is to transform the observation using the `TransformObservation` env wrapper from `gymnasium`, merging the two first dimension into one
- This is precisely what sb3 does in their VecFrameStack!

Worth saying this can have very unintended consequence: if one passes an image through Gymnasium's frame stacking and then uses sb3, the image won't be recognized as such, as the feature extractor for that image will be set to be `Flatten` (instead of `NatureCNN`).

Happy to open a PR to change this check, but I wanted to double check it made sense first!

### To Reproduce

see above

### Relevant log output / Error message

```shell
see above
```

### System Info

- OS: macOS-15.0.1-arm64-arm-64bit Darwin Kernel Version 24.0.0: Tue Sep 24 23:36:26 PDT 2024; root:xnu-11215.1.12~1/RELEASE_ARM64_T8103
- Python: 3.12.9
- Stable-Baselines3: 2.5.0
- PyTorch: 2.6.0
- GPU Enabled: False
- Numpy: 2.2.3
- Cloudpickle: 3.1.1
- Gymnasium: 1.0.0

### Checklist

- [x] My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
- [x] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo
- [x] I have read the [documentation](https://stable-baselines3.readthedocs.io/en/master/)
- [x] I have provided a [minimal and working](https://github.com/DLR-RM/stable-baselines3/issues/982#issuecomment-1197044014) example to reproduce the bug
- [x] I've used the [markdown code blocks](https://help.github.com/en/articles/creating-and-highlighting-code-blocks) for both code and stack traces.

Contributor guide

Open the contributing guide

Research direction

Start with the minimal reproduction in the issue and inspect stable_baselines3/common/preprocessing.py, especially is_image_space. Trace how the observation-space result selects NatureCNN versus Flatten, then determine the expected behavior for Gymnasium's FrameStackObservation and verify it against the reported stacked shape. Done means the intended stacked-image behavior is explicitly supported or rejected without silently choosing the wrong extractor.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
computer-vision, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.