Repository: DLR-RM/stable-baselines3 (Python)

Issue:
[Bug]: check_env incorrectly assumes that an environment has >1 step() available before termination
### 🐛 Bug

When calling `stable_baselines3.common.env_checker.check_env`, the code inside calls `step()` on the checked environment even if the previous `step()` returned `terminated == True`, thus leading to exceptions with some environments under some circumstances (depending on the environment's initial state, the action sampled, ...).

The problem occurs at least with `stable-baselines3` `2.0.0a5` through `2.0.0a10` (checked only the two).

I hope I didn't miss something obvious, or didn't do something stupid about the number of steps that an environment should have available.

### To Reproduce

I've created a minimal example to reproduce the bug here: https://github.com/lutogniew/sb3-envchecker-bug

It's a simple dummy environment that returns `terminated = True` after the configured number of calls to `step()`, and throws an exception afterwards until the environment is `reset()`. It seems that with the default configuration, `check_env` triggers if the environment has only 1 `step()` available. If it has at least 2, then the problem isn't triggered.

### Relevant log output / Error message

_No response_

### System Info

- OS: Linux-6.2.0-20-generic-x86_64-with-glibc2.37 # 20-Ubuntu SMP PREEMPT_DYNAMIC Thu Apr  6 07:48:48 UTC 2023
- Python: 3.11.2
- Stable-Baselines3: 2.0.0a10
- PyTorch: 2.0.1+cu117
- GPU Enabled: True
- Numpy: 1.24.3
- Cloudpickle: 2.2.1
- Gymnasium: 0.28.1

### Checklist

- [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 working 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.


Discussion:
Hello,
you are right, this is a corner case when the episode is only one step (normally unlikely).

A simple fix would to call `vec_env.reset()` after that line: https://github.com/DLR-RM/stable-baselines3/blob/1bfb55df0edc6882ecadf499419cd3ecb8fe6a79/stable_baselines3/common/env_checker.py#L112

I would welcome a PR that solves this issue =)
Thanks for taking a look into this!

I'll try to submit a PR within a day or so, as soon as I get my bearings around the project.



Patch:
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
index 16a7737e52..d2a3d60602 100644
--- a/docs/misc/changelog.rst
+++ b/docs/misc/changelog.rst
@@ -3,7 +3,7 @@
 Changelog
 ==========
 
-Release 2.0.0a10 (WIP)
+Release 2.0.0a11 (WIP)
 --------------------------
 
 **Gymnasium support**
@@ -39,6 +39,8 @@ Bug Fixes:
 - Fixed ``VecExtractDictObs`` does not handle terminal observation (@WeberSamuel)
 - Set NumPy version to ``>=1.20`` due to use of ``numpy.typing`` (@troiganto)
 - Fixed loading DQN changes ``target_update_interval`` (@tobirohrer)
+- Fixed env checker to properly reset the env before calling ``step()`` when checking
+  for ``Inf`` and ``NaN`` (@lutogniew)
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -1346,3 +1348,4 @@ And all the contributors:
 @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875
 @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong
 @DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError @harveybellini @DavyMorgan @FieteO @jonasreiher @npit @WeberSamuel @troiganto
+@lutogniew
diff --git a/stable_baselines3/common/env_checker.py b/stable_baselines3/common/env_checker.py
index 058710df9d..b6ce490dfe 100644
--- a/stable_baselines3/common/env_checker.py
+++ b/stable_baselines3/common/env_checker.py
@@ -110,6 +110,7 @@ def _check_unsupported_spaces(env: gym.Env, observation_space: spaces.Space, act
 def _check_nan(env: gym.Env) -> None:
     """Check for Inf and NaN using the VecWrapper."""
     vec_env = VecCheckNan(DummyVecEnv([lambda: env]))
+    vec_env.reset()
     for _ in range(10):
         action = np.array([env.action_space.sample()])
         _, _, _, _ = vec_env.step(action)
diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt
index 7385c4c8b2..d70b1bb71b 100644
--- a/stable_baselines3/version.txt
+++ b/stable_baselines3/version.txt
@@ -1,1 +1,1 @@
-2.0.0a10
+2.0.0a11


Test patch:
diff --git a/tests/test_env_checker.py b/tests/test_env_checker.py
index e855e2137d..c0a5e06102 100644
--- a/tests/test_env_checker.py
+++ b/tests/test_env_checker.py
@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 import gymnasium as gym
 import numpy as np
@@ -112,3 +112,47 @@ def step(self, action):
     test_env = TestEnv()
     with pytest.raises(AssertionError, match=error_message):
         check_env(env=test_env)
+
+
+class LimitedStepsTestEnv(gym.Env):
+    action_space = spaces.Discrete(n=2)
+    observation_space = spaces.Discrete(n=2)
+
+    def __init__(self, steps_before_termination: int = 1):
+        super().__init__()
+
+        assert steps_before_termination >= 1
+        self._steps_before_termination = steps_before_termination
+
+        self._steps_called = 0
+        self._terminated = False
+
+    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict] = None) -> Tuple[int, Dict]:
+        super().reset(seed=seed)
+
+        self._steps_called = 0
+        self._terminated = False
+
+        return 0, {}
+
+    def step(self, action: np.ndarray) -> Tuple[int, float, bool, bool, Dict[str, Any]]:
+        self._steps_called += 1
+
+        assert not self._terminated
+
+        observation = 0
+        reward = 0.0
+        self._terminated = self._steps_called >= self._steps_before_termination
+        truncated = False
+
+        return observation, reward, self._terminated, truncated, {}
+
+    def render(self) -> None:
+        pass
+
+
+def test_check_env_single_step_env():
+    test_env = LimitedStepsTestEnv(steps_before_termination=1)
+
+    # This should not throw
+    check_env(env=test_env, warn=True)