Fix DQN predict shape for single Gym env (#222)

* Fix DQN predict shape for single Gym env

* Remove unused imports

docs/misc/changelog.rst

@@ -27,6 +27,7 @@ New Features:
 Bug Fixes:
 ^^^^^^^^^^
 - Fixed bug where code added VecTranspose on channel-first image environments (thanks @qxcv)
+- Fixed ``DQN`` predict method when using single ``gym.Env`` with ``deterministic=False``
 - Fixed bug that the arguments order of ``explained_variance()`` in ``ppo.py`` and ``a2c.py`` is not correct (@thisray)
 
 Deprecations:

stable_baselines3/common/policies.py

@@ -3,7 +3,7 @@
 import collections
 from abc import ABC, abstractmethod
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import gym
 import numpy as np

stable_baselines3/common/utils.py

@@ -3,7 +3,7 @@ import os
 import random
 from collections import deque
 from itertools import zip_longest
-from typing import Callable, Iterable, Optional, Union
+from typing import Iterable, Optional, Union
 
 import gym
 import numpy as np

stable_baselines3/dqn/dqn.py

@@ -7,7 +7,7 @@ from torch.nn import functional as F
 from stable_baselines3.common import logger
 from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
 from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback
-from stable_baselines3.common.utils import get_linear_fn, polyak_update
+from stable_baselines3.common.utils import get_linear_fn, is_vectorized_observation, polyak_update
 from stable_baselines3.dqn.policies import DQNPolicy
 
 
@@ -203,8 +203,11 @@ class DQN(OffPolicyAlgorithm):
             (used in recurrent policies)
         """
         if not deterministic and np.random.rand() < self.exploration_rate:
-            n_batch = observation.shape[0]
-            action = np.array([self.action_space.sample() for _ in range(n_batch)])
+            if is_vectorized_observation(observation, self.observation_space):
+                n_batch = observation.shape[0]
+                action = np.array([self.action_space.sample() for _ in range(n_batch)])
+            else:
+                action = np.array(self.action_space.sample())
         else:
             action, state = self.policy.predict(observation, state, mask, deterministic)
         return action, state

stable_baselines3/dqn/policies.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
 import gym
 import torch as th

stable_baselines3/sac/policies.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import gym
 import torch as th

stable_baselines3/td3/policies.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 import gym
 import torch as th

tests/test_predict.py

@@ -59,3 +59,13 @@ def test_predict(model_class, env_id, device):
     vec_env_obs = vec_env.reset()
     action, _ = model.predict(vec_env_obs)
     assert action.shape[0] == vec_env_obs.shape[0]
+
+    # Special case for DQN to check the epsilon greedy exploration
+    if model_class == DQN:
+        model.exploration_rate = 1.0
+        action, _ = model.predict(obs, deterministic=False)
+        assert action.shape == env.action_space.shape
+        assert env.action_space.contains(action)
+
+        action, _ = model.predict(vec_env_obs, deterministic=False)
+        assert action.shape[0] == vec_env_obs.shape[0]