Sync identity envs

stable_baselines3/common/identity_env.py

@@ -1,7 +1,7 @@
-from typing import List, Union
+from typing import List, Union, Optional
 
 import numpy as np
-from gym import Env
+from gym import Env, Space
 from gym.spaces import Discrete, MultiDiscrete, MultiBinary, Box
 
 
@@ -9,22 +9,36 @@ from stable_baselines3.common.type_aliases import GymStepReturn, GymObs
 
 
 class IdentityEnv(Env):
-    def __init__(self, dim, ep_length=100):
+    def __init__(self,
+                 dim: Optional[int] = None,
+                 space: Optional[Space] = None,
+                 ep_length: int = 100):
         """
         Identity environment for testing purposes
 
-        :param dim: (int) the size of the dimensions you want to learn
-        :param ep_length: (int) the length of each episodes in timesteps
+        :param dim: the size of the action and observation dimension you want
+            to learn. Provide at most one of ``dim`` and ``space``. If both are
+            None, then initialization proceeds with ``dim=1`` and ``space=None``.
+        :param space: the action and observation space. Provide at most one of
+            ``dim`` and ``space``.
+        :param ep_length: the length of each episode in timesteps
         """
-        self.action_space = Discrete(dim)
-        self.observation_space = self.action_space
+        if space is None:
+            if dim is None:
+                dim = 1
+            space = Discrete(dim)
+        else:
+            assert dim is None, "arguments for both 'dim' and 'space' provided: at most one allowed"
+
+        self.action_space = self.observation_space = space
         self.ep_length = ep_length
         self.current_step = 0
-        self.dim = dim
+        self.num_resets = -1  # Becomes 0 after __init__ exits.
         self.reset()
 
     def reset(self) -> GymObs:
         self.current_step = 0
+        self.num_resets += 1
         self._choose_next_state()
         return self.state
 
@@ -55,18 +69,11 @@ class IdentityEnvBox(IdentityEnv):
         :param low: (float) the lower bound of the box dim
         :param high: (float) the upper bound of the box dim
         :param eps: (float) the epsilon bound for correct value
-        :param ep_length: (int) the length of each episodes in timesteps
+        :param ep_length: (int) the length of each episode in timesteps
         """
-        super(IdentityEnvBox, self).__init__(1, ep_length)
-        self.action_space = Box(low=low, high=high, shape=(1,), dtype=np.float32)
-        self.observation_space = self.action_space
+        space = Box(low=low, high=high, shape=(1,), dtype=np.float32)
+        super().__init__(ep_length=ep_length, space=space)
         self.eps = eps
-        self.reset()
-
-    def reset(self) -> np.ndarray:
-        self.current_step = 0
-        self._choose_next_state()
-        return self.state
 
     def step(self, action: np.ndarray) -> GymStepReturn:
         reward = self._get_reward(action)
@@ -75,39 +82,32 @@ class IdentityEnvBox(IdentityEnv):
         done = self.current_step >= self.ep_length
         return self.state, reward, done, {}
 
-    def _choose_next_state(self) -> None:
-        self.state = self.observation_space.sample()
-
     def _get_reward(self, action: np.ndarray) -> float:
         return 1.0 if (self.state - self.eps) <= action <= (self.state + self.eps) else 0.0
 
 
 class IdentityEnvMultiDiscrete(IdentityEnv):
-    def __init__(self, dim: int, ep_length: int = 100):
+    def __init__(self, dim: int = 1, ep_length: int = 100):
         """
         Identity environment for testing purposes
 
         :param dim: (int) the size of the dimensions you want to learn
-        :param ep_length: (int) the length of each episodes in timesteps
+        :param ep_length: (int) the length of each episode in timesteps
         """
-        super(IdentityEnvMultiDiscrete, self).__init__(dim, ep_length)
-        self.action_space = MultiDiscrete([dim, dim])
-        self.observation_space = self.action_space
-        self.reset()
+        space = MultiDiscrete([dim, dim])
+        super().__init__(ep_length=ep_length, space=space)
 
 
 class IdentityEnvMultiBinary(IdentityEnv):
-    def __init__(self, dim: int, ep_length: int = 100):
+    def __init__(self, dim: int = 1, ep_length: int = 100):
         """
         Identity environment for testing purposes
 
         :param dim: (int) the size of the dimensions you want to learn
-        :param ep_length: (int) the length of each episodes in timesteps
+        :param ep_length: (int) the length of each episode in timesteps
         """
-        super(IdentityEnvMultiBinary, self).__init__(dim, ep_length)
-        self.action_space = MultiBinary(dim)
-        self.observation_space = self.action_space
-        self.reset()
+        space = MultiBinary(dim)
+        super().__init__(ep_length=ep_length, space=space)
 
 
 class FakeImageEnv(Env):

stable_baselines3/common/utils.py

@@ -11,14 +11,17 @@ def set_random_seed(seed: int, using_cuda: bool = False) -> None:
     :param seed: (int)
     :param using_cuda: (bool)
     """
+    # Seed python RNG
     random.seed(seed)
+    # Seed numpy RNG
     np.random.seed(seed)
+    # seed the RNG for all devices (both CPU and CUDA)
     th.manual_seed(seed)
 
     if using_cuda:
-        # Make CuDNN Determinist
+        # Deterministic operations for CuDNN, it may impact performances
         th.backends.cudnn.deterministic = True
-        th.cuda.manual_seed(seed)
+        th.backends.cudnn.benchmark = False
 
 
 # From stable baselines

tests/test_deterministic.py

@@ -0,0 +1,36 @@
+import pytest
+
+from stable_baselines3 import A2C, PPO, SAC, TD3
+from stable_baselines3.common.noise import NormalActionNoise
+
+N_STEPS_TRAINING = 3000
+SEED = 0
+
+
+@pytest.mark.parametrize("algo", [A2C, PPO, SAC, TD3])
+def test_deterministic_training_common(algo):
+    results = [[], []]
+    rewards = [[], []]
+    # Smaller network
+    kwargs = {'policy_kwargs': dict(net_arch=[64])}
+    if algo in [TD3, SAC]:
+        env_id = 'Pendulum-v0'
+        kwargs.update({'action_noise': NormalActionNoise(0.0, 0.1),
+                       'learning_starts': 100})
+    else:
+        env_id = 'CartPole-v1'
+        # if algo == DQN:
+        #     kwargs.update({'learning_starts': 100})
+
+    for i in range(2):
+        model = algo('MlpPolicy', env_id, seed=SEED, **kwargs)
+        model.learn(N_STEPS_TRAINING)
+        env = model.get_env()
+        obs = env.reset()
+        for _ in range(100):
+            action, _ = model.predict(obs, deterministic=False)
+            obs, reward, _, _ = env.step(action)
+            results[i].append(action)
+            rewards[i].append(reward)
+    assert sum(results[0]) == sum(results[1]), results
+    assert sum(rewards[0]) == sum(rewards[1]), rewards

tests/test_envs.py

@@ -5,11 +5,11 @@ import numpy as np
 
 from stable_baselines3.common.env_checker import check_env
 from stable_baselines3.common.bit_flipping_env import BitFlippingEnv
-from stable_baselines3.common.identity_env import (IdentityEnv, IdentityEnvBox,
+from stable_baselines3.common.identity_env import (IdentityEnv, IdentityEnvBox, FakeImageEnv,
                                                   IdentityEnvMultiBinary, IdentityEnvMultiDiscrete,)
 
 ENV_CLASSES = [BitFlippingEnv, IdentityEnv, IdentityEnvBox, IdentityEnvMultiBinary,
-               IdentityEnvMultiDiscrete]
+               IdentityEnvMultiDiscrete, FakeImageEnv]
 
 
 @pytest.mark.parametrize("env_id", ['CartPole-v0', 'Pendulum-v0'])
@@ -43,7 +43,7 @@ def test_high_dimension_action_space():
     Test for continuous action space
     with more than one action.
     """
-    env = gym.make('Pendulum-v0')
+    env = FakeImageEnv()
     # Patch the action space
     env.action_space = spaces.Box(low=-1, high=1, shape=(20,), dtype=np.float32)
     # Patch to avoid error
@@ -68,7 +68,7 @@ def test_high_dimension_action_space():
     spaces.Dict({"position": spaces.Discrete(5)}),
 ])
 def test_non_default_spaces(new_obs_space):
-    env = gym.make('BreakoutNoFrameskip-v4')
+    env = FakeImageEnv()
     env.observation_space = new_obs_space
     # Patch methods to avoid errors
     env.reset = new_obs_space.sample