Refactor Tests + Add Helpers (hill-a#508)

* Add helpers

* Refactor some tests

* Continue refactoring

* Fix for codacy

* Fixes for travis

* Clean up imports

* Fix syntax error

* Fix VecEnv constructor

* Fix perf check in tests

* Seed identity env + minor updates

* Allow more diff after training again

* Try to fix travis non-determinism

* Add tests for the new helpers

* Codacy fixes

* Fix callback logic

* Address comments

* Address review comments

* Make codacy happy

* Fix docstring indentation

* Update README example

* Remove use_subprocess and update doc

Author: araffin

.gitignore

@@ -12,6 +12,7 @@
 __pycache__/
 _build/
 *.npz
+*.zip
 
 # Setuptools distribution and build folders.
 /dist/

README.md

@@ -113,7 +113,9 @@ from stable_baselines.common.vec_env import DummyVecEnv
 from stable_baselines import PPO2
 
 env = gym.make('CartPole-v1')
-env = DummyVecEnv([lambda: env])  # The algorithms require a vectorized environment to run
+# Optional: PPO2 requires a vectorized environment to run
+# the env is now wrapped automatically when passing it to the constructor
+# env = DummyVecEnv([lambda: env])
 
 model = PPO2(MlpPolicy, env, verbose=1)
 model.learn(total_timesteps=10000)

docs/common/evaluation.rst

@@ -0,0 +1,7 @@
+.. _eval:
+
+Evaluation Helper
+=================
+
+.. automodule:: stable_baselines.common.evaluation
+  :members:

docs/guide/examples.rst

@@ -56,6 +56,8 @@ In the following example, we will train, save and load a DQN model on the Lunar
   import gym
 
   from stable_baselines import DQN
+  from stable_baselines.common.evaluation import evaluate_policy
+
 
   # Create environment
   env = gym.make('LunarLander-v2')
@@ -71,6 +73,9 @@ In the following example, we will train, save and load a DQN model on the Lunar
   # Load the trained agent
   model = DQN.load("dqn_lunar")
 
+  # Evaluate the agent
+  mean_reward, n_steps = evaluate_policy(model, model.get_env(), n_eval_episodes=10)
+
   # Enjoy trained agent
   obs = env.reset()
   for i in range(1000):
@@ -98,7 +103,7 @@ Multiprocessing: Unleashing the Power of Vectorized Environments
 
   from stable_baselines.common.policies import MlpPolicy
   from stable_baselines.common.vec_env import SubprocVecEnv
-  from stable_baselines.common import set_global_seeds
+  from stable_baselines.common import set_global_seeds, make_vec_env
   from stable_baselines import ACKTR
 
   def make_env(env_id, rank, seed=0):
@@ -123,6 +128,10 @@ Multiprocessing: Unleashing the Power of Vectorized Environments
       # Create the vectorized environment
       env = SubprocVecEnv([make_env(env_id, i) for i in range(num_cpu)])
 
+      # Stable Baselines provides you with make_vec_env() helper
+      # which does exactly the previous steps for you:
+      # env = make_vec_env(env_id, n_envs=num_cpu, seed=0)
+
       model = ACKTR(MlpPolicy, env, verbose=1)
       model.learn(total_timesteps=25000)
 
@@ -340,8 +349,6 @@ A2C policy gradient updates on the model.
   import gym
   import numpy as np
 
-  from stable_baselines.common.policies import MlpPolicy
-  from stable_baselines.common.vec_env import DummyVecEnv
   from stable_baselines import A2C
 
   def mutate(params):
@@ -365,9 +372,8 @@ A2C policy gradient updates on the model.
 
   # Create env
   env = gym.make('CartPole-v1')
-  env = DummyVecEnv([lambda: env])
   # Create policy with a small network
-  model = A2C(MlpPolicy, env, ent_coef=0.0, learning_rate=0.1,
+  model = A2C('MlpPolicy', env, ent_coef=0.0, learning_rate=0.1,
               policy_kwargs={'net_arch': [8, ]})
 
   # Use traditional actor-critic policy gradient updates to
@@ -546,6 +552,9 @@ You can also move from learning on one environment to another for `continual lea
       obs, rewards, dones, info = env.step(action)
       env.render()
 
+  # Close the processes
+  env.close()
+
   # The number of environments must be identical when changing environments
   env = make_atari_env('SpaceInvadersNoFrameskip-v4', num_env=8, seed=0)
 
@@ -558,6 +567,7 @@ You can also move from learning on one environment to another for `continual lea
       action, _states = model.predict(obs)
       obs, rewards, dones, info = env.step(action)
       env.render()
+  env.close()
 
 
 Record a Video
@@ -591,6 +601,7 @@ Record a mp4 video (here using a random agent).
   for _ in range(video_length + 1):
     action = [env.action_space.sample()]
     obs, _, _, _ = env.step(action)
+  # Save the video
   env.close()
 
 
@@ -606,10 +617,9 @@ Bonus: Make a GIF of a Trained Agent
   import imageio
   import numpy as np
 
-  from stable_baselines.common.policies import MlpPolicy
   from stable_baselines import A2C
 
-  model = A2C(MlpPolicy, "LunarLander-v2").learn(100000)
+  model = A2C("MlpPolicy", "LunarLander-v2").learn(100000)
 
   images = []
   obs = model.env.reset()

docs/guide/quickstart.rst

@@ -17,7 +17,9 @@ Here is a quick example of how to train and run PPO2 on a cartpole environment:
   from stable_baselines import PPO2
 
   env = gym.make('CartPole-v1')
-  env = DummyVecEnv([lambda: env])  # The algorithms require a vectorized environment to run
+  # Optional: PPO2 requires a vectorized environment to run
+  # the env is now wrapped automatically when passing it to the constructor
+  # env = DummyVecEnv([lambda: env])
 
   model = PPO2(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=10000)

docs/index.rst

@@ -80,6 +80,7 @@ This toolset is a fork of OpenAI Baselines, with a major structural refactoring,
   common/tf_utils
   common/cmd_utils
   common/schedules
+  common/evaluation
 
 .. toctree::
   :maxdepth: 1

docs/misc/changelog.rst

@@ -13,10 +13,16 @@ Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - The `seed` argument has been moved from `learn()` method to model constructor
   in order to have reproducible results
+- `allow_early_resets` of the `Monitor` wrapper now default to `True`
+- `make_atari_env` now returns a `DummyVecEnv` by default (instead of a `SubprocVecEnv`)
+  this usually improves performance.
 
 New Features:
 ^^^^^^^^^^^^^
 - Add `n_cpu_tf_sess` to model constructor to choose the number of threads used by Tensorflow
+- Environments are automatically wrapped in a `DummyVecEnv` if needed when passing them to the model constructor
+- Added `stable_baselines.common.make_vec_env` helper to simplify VecEnv creation
+- Added `stable_baselines.common.evaluation.evaluate_policy` helper to simplify model evaluation
 - `VecNormalize` now supports being pickled and unpickled.
 - Add parameter `exploration_initial_eps` to DQN. (@jdossgollin)
 - Add type checking and PEP 561 compliance.
@@ -38,6 +44,7 @@ Deprecations:
 Others:
 ^^^^^^^
 - Add upper bound for Tensorflow version (<2.0.0).
+- Refactored test to remove duplicated code
 - Add pull request template
 
 Documentation:
@@ -46,8 +53,11 @@ Documentation:
 - Add Snake Game AI project (@pedrohbtp)
 - Add note on the support Tensorflow versions.
 - Remove unnecessary steps required for Windows installation.
+- Remove `DummyVecEnv` creation when not needed
+- Added `make_vec_env` to the examples to simplify VecEnv creation
 - Add QuaRL project (@srivatsankrishnan)
 - Add Pwnagotchi project (@evilsocket)
+- Fix multiprocessing example (@rusu24edward)
 - Fix `result_plotter` example
 - Fix typo in algos.rst, "containes" to "contains" (@SyllogismRXS)
 
@@ -530,4 +540,4 @@ Thanks to @bjmuld @iambenzo @iandanforth @r7vme @brendenpetersen @huvar @abhiskk
 @EliasHasle @mrakgr @Bleyddyn @antoine-galataud @junhyeokahn @AdamGleave @keshaviyengar @tperol
 @XMaster96 @kantneel @Pastafarianist @GerardMaggiolino @PatrickWalter214 @yutingsz @sc420 @Aaahh @billtubbs
 @Miffyli @dwiel @miguelrass @qxcv @jaberkow @eavelardev @ruifeng96150 @pedrohbtp @srivatsankrishnan @evilsocket
-@MarvineGothic @jdossgollin @SyllogismRXS
+@MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward

docs/modules/a2c.rst

@@ -49,12 +49,11 @@ Train a A2C agent on `CartPole-v1` using 4 processes.
   import gym
 
   from stable_baselines.common.policies import MlpPolicy
-  from stable_baselines.common.vec_env import SubprocVecEnv
+  from stable_baselines.common import make_vec_env
   from stable_baselines import A2C
 
-  # multiprocess environment
-  n_cpu = 4
-  env = SubprocVecEnv([lambda: gym.make('CartPole-v1') for i in range(n_cpu)])
+  # Parallel environments
+  env = make_vec_env('CartPole-v1', n_envs=4)
 
   model = A2C(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=25000)

docs/modules/acer.rst

@@ -43,12 +43,11 @@ Example
   import gym
 
   from stable_baselines.common.policies import MlpPolicy, MlpLstmPolicy, MlpLnLstmPolicy
-  from stable_baselines.common.vec_env import SubprocVecEnv
+  from stable_baselines.common import make_vec_env
   from stable_baselines import ACER
 
   # multiprocess environment
-  n_cpu = 4
-  env = SubprocVecEnv([lambda: gym.make('CartPole-v1') for i in range(n_cpu)])
+  env = make_vec_env('CartPole-v1', n_envs=4)
 
   model = ACER(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=25000)

docs/modules/acktr.rst

@@ -44,12 +44,11 @@ Example
   import gym
 
   from stable_baselines.common.policies import MlpPolicy, MlpLstmPolicy, MlpLnLstmPolicy
-  from stable_baselines.common.vec_env import SubprocVecEnv
+  from stable_baselines.common import make_vec_env
   from stable_baselines import ACKTR
 
   # multiprocess environment
-  n_cpu = 4
-  env = SubprocVecEnv([lambda: gym.make('CartPole-v1') for i in range(n_cpu)])
+  env = make_vec_env('CartPole-v1', n_envs=4)
 
   model = ACKTR(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=25000)

docs/modules/ddpg.rst

@@ -63,12 +63,10 @@ Example
   import numpy as np
 
   from stable_baselines.ddpg.policies import MlpPolicy
-  from stable_baselines.common.vec_env import DummyVecEnv
-  from stable_baselines.ddpg.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise, AdaptiveParamNoiseSpec
+  from stable_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise, AdaptiveParamNoiseSpec
   from stable_baselines import DDPG
 
   env = gym.make('MountainCarContinuous-v0')
-  env = DummyVecEnv([lambda: env])
 
   # the noise objects for DDPG
   n_actions = env.action_space.shape[-1]
@@ -148,7 +146,6 @@ You can easily define a custom architecture for the policy network:
   import gym
 
   from stable_baselines.ddpg.policies import FeedForwardPolicy
-  from stable_baselines.common.vec_env import DummyVecEnv
   from stable_baselines import DDPG
 
   # Custom MLP policy of two layers of size 16 each
@@ -159,10 +156,7 @@ You can easily define a custom architecture for the policy network:
                                              layer_norm=False,
                                              feature_extraction="mlp")
 
-  # Create and wrap the environment
-  env = gym.make('Pendulum-v0')
-  env = DummyVecEnv([lambda: env])
 
-  model = DDPG(CustomDDPGPolicy, env, verbose=1)
+  model = DDPG(CustomDDPGPolicy, 'Pendulum-v0', verbose=1)
   # Train the agent
   model.learn(total_timesteps=100000)

docs/modules/gail.rst

@@ -111,7 +111,7 @@ Example
   # Load the expert dataset
   dataset = ExpertDataset(expert_path='expert_pendulum.npz', traj_limitation=10, verbose=1)
 
-  model = GAIL("MlpPolicy", 'Pendulum-v0', dataset, verbose=1)
+  model = GAIL('MlpPolicy', 'Pendulum-v0', dataset, verbose=1)
   # Note: in practice, you need to train for 1M steps to have a working policy
   model.learn(total_timesteps=1000)
   model.save("gail_pendulum")

docs/modules/ppo1.rst

@@ -59,11 +59,9 @@ Example
   import gym
 
   from stable_baselines.common.policies import MlpPolicy
-  from stable_baselines.common.vec_env import DummyVecEnv
   from stable_baselines import PPO1
 
   env = gym.make('CartPole-v1')
-  env = DummyVecEnv([lambda: env])
 
   model = PPO1(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=25000)

docs/modules/ppo2.rst

@@ -61,12 +61,11 @@ Train a PPO agent on `CartPole-v1` using 4 processes.
    import gym
 
    from stable_baselines.common.policies import MlpPolicy
-   from stable_baselines.common.vec_env import SubprocVecEnv
+   from stable_baselines.common import make_vec_env
    from stable_baselines import PPO2
 
    # multiprocess environment
-   n_cpu = 4
-   env = SubprocVecEnv([lambda: gym.make('CartPole-v1') for i in range(n_cpu)])
+   env = make_vec_env('CartPole-v1', n_envs=4)
 
    model = PPO2(MlpPolicy, env, verbose=1)
    model.learn(total_timesteps=25000)

docs/modules/sac.rst

@@ -75,11 +75,9 @@ Example
   import numpy as np
 
   from stable_baselines.sac.policies import MlpPolicy
-  from stable_baselines.common.vec_env import DummyVecEnv
   from stable_baselines import SAC
 
   env = gym.make('Pendulum-v0')
-  env = DummyVecEnv([lambda: env])
 
   model = SAC(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=50000, log_interval=10)

docs/modules/td3.rst

@@ -73,7 +73,6 @@ Example
   from stable_baselines.ddpg.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
 
   env = gym.make('Pendulum-v0')
-  env = DummyVecEnv([lambda: env])
 
   # The noise objects for TD3
   n_actions = env.action_space.shape[-1]

docs/modules/trpo.rst

@@ -49,11 +49,9 @@ Example
   import gym
 
   from stable_baselines.common.policies import MlpPolicy
-  from stable_baselines.common.vec_env import DummyVecEnv
   from stable_baselines import TRPO
 
   env = gym.make('CartPole-v1')
-  env = DummyVecEnv([lambda: env])
 
   model = TRPO(MlpPolicy, env, verbose=1)
   model.learn(total_timesteps=25000)

setup.py

@@ -80,7 +80,9 @@
 from stable_baselines import PPO2
 
 env = gym.make('CartPole-v1')
-env = DummyVecEnv([lambda: env])  # The algorithms require a vectorized environment to run
+# Optional: PPO2 requires a vectorized environment to run
+# the env is now wrapped automatically when passing it to the constructor
+# env = DummyVecEnv([lambda: env])
 
 model = PPO2(MlpPolicy, env, verbose=1)
 model.learn(total_timesteps=10000)

stable_baselines/bench/monitor.py

@@ -14,7 +14,7 @@ class Monitor(Wrapper):
     EXT = "monitor.csv"
     file_handler = None
 
-    def __init__(self, env, filename, allow_early_resets=False, reset_keywords=(), info_keywords=()):
+    def __init__(self, env, filename, allow_early_resets=True, reset_keywords=(), info_keywords=()):
         """
         A monitor wrapper for Gym environments, it is used to know the episode reward, length, time and other data.
 

stable_baselines/common/__init__.py

@@ -6,3 +6,4 @@
 from stable_baselines.common.misc_util import zipsame, set_global_seeds, boolean_flag
 from stable_baselines.common.base_class import BaseRLModel, ActorCriticRLModel, OffPolicyRLModel, SetVerbosity, \
     TensorboardWriter
+from stable_baselines.common.cmd_util import make_vec_env

stable_baselines/common/base_class.py

@@ -70,7 +70,12 @@ def __init__(self, policy, env, verbose=0, *, requires_vec_env, policy_base,
                 if isinstance(env, VecEnv):
                     self.n_envs = env.num_envs
                 else:
-                    raise ValueError("Error: the model requires a vectorized environment, please use a VecEnv wrapper.")
+                    # The model requires a VecEnv
+                    # wrap it in a DummyVecEnv to avoid error
+                    self.env = DummyVecEnv([lambda: env])
+                    if self.verbose >= 1:
+                        print("Wrapping the env in a DummyVecEnv.")
+                    self.n_envs = 1
             else:
                 if isinstance(env, VecEnv):
                     if env.num_envs == 1:

stable_baselines/common/cmd_util.py

@@ -3,6 +3,7 @@
 """
 
 import os
+import warnings
 
 import gym
 
@@ -14,20 +15,80 @@
 from stable_baselines.common.vec_env import DummyVecEnv, SubprocVecEnv
 
 
+def make_vec_env(env_id, n_envs=1, seed=None, start_index=0,
+                 monitor_dir=None, wrapper_class=None,
+                 env_kwargs=None, vec_env_cls=None, vec_env_kwargs=None):
+    """
+    Create a wrapped, monitored `VecEnv`.
+    By default it uses a `DummyVecEnv` which is usually faster
+    than a `SubprocVecEnv`.
+
+    :param env_id: (str or Type[gym.Env]) the environment ID or the environment class
+    :param n_envs: (int) the number of environments you wish to have in parallel
+    :param seed: (int) the inital seed for the random number generator
+    :param start_index: (int) start rank index
+    :param monitor_dir: (str) Path to a folder where the monitor files will be saved.
+        If None, no file will be written, however, the env will still be wrapped
+        in a Monitor wrapper to provide additional information about training.
+    :param wrapper_class: (gym.Wrapper or callable) Additional wrapper to use on the environment.
+        This can also be a function with single argument that wraps the environment in many things.
+    :param env_kwargs: (dict) Optional keyword argument to pass to the env constructor
+    :param vec_env_cls: (Type[VecEnv]) A custom `VecEnv` class constructor. Default: None.
+    :param vec_env_kwargs: (dict) Keyword arguments to pass to the `VecEnv` class constructor.
+    :return: (VecEnv) The wrapped environment
+    """
+    env_kwargs = {} if env_kwargs is None else env_kwargs
+    vec_env_kwargs = {} if vec_env_kwargs is None else vec_env_kwargs
+
+    def make_env(rank):
+        def _init():
+            if isinstance(env_id, str):
+                env = gym.make(env_id)
+                if len(env_kwargs) > 0:
+                    warnings.warn("No environment class was passed (only an env ID) so `env_kwargs` will be ignored")
+            else:
+                env = env_id(**env_kwargs)
+            if seed is not None:
+                env.seed(seed + rank)
+                env.action_space.seed(seed + rank)
+            # Wrap the env in a Monitor wrapper
+            # to have additional training information
+            monitor_path = os.path.join(monitor_dir, str(rank)) if monitor_dir is not None else None
+            # Create the monitor folder if needed
+            if monitor_path is not None:
+                os.makedirs(monitor_dir, exist_ok=True)
+            env = Monitor(env, filename=monitor_path)
+            # Optionally, wrap the environment with the provided wrapper
+            if wrapper_class is not None:
+                env = wrapper_class(env)
+            return env
+        return _init
+
+    # No custom VecEnv is passed
+    if vec_env_cls is None:
+        # Default: use a DummyVecEnv
+        vec_env_cls = DummyVecEnv
+
+    return vec_env_cls([make_env(i + start_index) for i in range(n_envs)], **vec_env_kwargs)
+
+
 def make_atari_env(env_id, num_env, seed, wrapper_kwargs=None,
-                   start_index=0, allow_early_resets=True, start_method=None):
+                   start_index=0, allow_early_resets=True,
+                   start_method=None, use_subprocess=False):
     """
-    Create a wrapped, monitored SubprocVecEnv for Atari.
+    Create a wrapped, monitored VecEnv for Atari.
 
     :param env_id: (str) the environment ID
     :param num_env: (int) the number of environment you wish to have in subprocesses
     :param seed: (int) the inital seed for RNG
     :param wrapper_kwargs: (dict) the parameters for wrap_deepmind function
     :param start_index: (int) start rank index
     :param allow_early_resets: (bool) allows early reset of the environment
-    :return: (Gym Environment) The atari environment
     :param start_method: (str) method used to start the subprocesses.
         See SubprocVecEnv doc for more information
+    :param use_subprocess: (bool) Whether to use `SubprocVecEnv` or `DummyVecEnv` when
+        `num_env` > 1, `DummyVecEnv` is usually faster. Default: False
+    :return: (VecEnv) The atari environment
     """
     if wrapper_kwargs is None:
         wrapper_kwargs = {}
@@ -43,8 +104,8 @@ def _thunk():
     set_global_seeds(seed)
 
     # When using one environment, no need to start subprocesses
-    if num_env == 1:
-        return DummyVecEnv([make_env(0)])
+    if num_env == 1 or not use_subprocess:
+        return DummyVecEnv([make_env(i + start_index) for i in range(num_env)])
 
     return SubprocVecEnv([make_env(i + start_index) for i in range(num_env)],
                          start_method=start_method)

stable_baselines/common/evaluation.py

@@ -0,0 +1,52 @@
+import numpy as np
+
+from stable_baselines.common.vec_env import VecEnv
+
+
+def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
+                    render=False, callback=None, reward_threshold=None,
+                    return_episode_rewards=False):
+    """
+    Runs policy for `n_eval_episodes` episodes and returns average reward.
+    This is made to work only with one env.
+
+    :param model: (BaseRLModel) The RL agent you want to evaluate.
+    :param env: (gym.Env or VecEnv) The gym environment. In the case of a `VecEnv`
+        this must contain only one environment.
+    :param n_eval_episodes: (int) Number of episode to evaluate the agent
+    :param deterministic: (bool) Whether to use deterministic or stochastic actions
+    :param render: (bool) Whether to render the environement or not
+    :param callback: (callable) callback function to do additional checks,
+        called after each step.
+    :param reward_threshold: (float) Minimum expected reward per episode,
+        this will raise an error if the performance is not met
+    :param return_episode_rewards: (bool) If True, a list of reward per episode
+        will be returned instead of the mean.
+    :return: (float, int) Mean reward per episode, total number of steps
+        returns ([float], int) when `return_episode_rewards` is True
+    """
+    if isinstance(env, VecEnv):
+        assert env.num_envs == 1, "You must pass only one environment when using this function"
+
+    episode_rewards, n_steps = [], 0
+    for _ in range(n_eval_episodes):
+        obs = env.reset()
+        done, state = False, None
+        episode_reward = 0.0
+        while not done:
+            action, state = model.predict(obs, state=state, deterministic=deterministic)
+            obs, reward, done, _info = env.step(action)
+            episode_reward += reward
+            if callback is not None:
+                callback(locals(), globals())
+            n_steps += 1
+            if render:
+                env.render()
+        episode_rewards.append(episode_reward)
+    mean_reward = np.mean(episode_rewards)
+    if reward_threshold is not None:
+        assert mean_reward > reward_threshold, 'Mean reward below threshold: '\
+                                         '{:.2f} < {:.2f}'.format(mean_reward, reward_threshold)
+    if return_episode_rewards:
+        return episode_rewards, n_steps
+    return mean_reward, n_steps

tests/test_action_space.py

@@ -4,6 +4,7 @@
 from stable_baselines import A2C, PPO1, PPO2, TRPO
 from stable_baselines.common.identity_env import IdentityEnvMultiBinary, IdentityEnvMultiDiscrete
 from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines.common.evaluation import evaluate_policy
 
 MODEL_LIST = [
     A2C,
@@ -27,14 +28,8 @@ def test_identity_multidiscrete(model_class):
 
     model = model_class("MlpPolicy", env)
     model.learn(total_timesteps=1000)
-
-    n_trials = 1000
-    reward_sum = 0
+    evaluate_policy(model, env, n_eval_episodes=5)
     obs = env.reset()
-    for _ in range(n_trials):
-        action, _ = model.predict(obs)
-        obs, reward, _, _ = env.step(action)
-        reward_sum += reward
 
     assert np.array(model.action_probability(obs)).shape == (2, 1, 10), \
         "Error: action_probability not returning correct shape"
@@ -56,14 +51,8 @@ def test_identity_multibinary(model_class):
 
     model = model_class("MlpPolicy", env)
     model.learn(total_timesteps=1000)
-
-    n_trials = 1000
-    reward_sum = 0
+    evaluate_policy(model, env, n_eval_episodes=5)
     obs = env.reset()
-    for _ in range(n_trials):
-        action, _ = model.predict(obs)
-        obs, reward, _, _ = env.step(action)
-        reward_sum += reward
 
     assert model.action_probability(obs).shape == (1, 10), \
         "Error: action_probability not returning correct shape"

tests/test_auto_vec_detection.py

@@ -1,10 +1,23 @@
 import pytest
 import numpy as np
 
-from stable_baselines import A2C, ACER, ACKTR, DDPG, DQN, PPO1, PPO2, SAC, TRPO
+from stable_baselines import A2C, ACER, ACKTR, DDPG, DQN, PPO1, PPO2, SAC, TRPO, TD3
 from stable_baselines.common.vec_env import DummyVecEnv
 from stable_baselines.common.identity_env import IdentityEnv, IdentityEnvBox, IdentityEnvMultiBinary, \
     IdentityEnvMultiDiscrete
+from stable_baselines.common.evaluation import evaluate_policy
+
+
+def check_shape(make_env, model_class, shape_1, shape_2):
+    model = model_class(policy="MlpPolicy", env=DummyVecEnv([make_env]))
+
+    env0 = make_env()
+    env1 = DummyVecEnv([make_env])
+
+    for env, expected_shape in [(env0, shape_1), (env1, shape_2)]:
+        def callback(locals_, _globals):
+            assert np.array(locals_['action']).shape == expected_shape
+        evaluate_policy(model, env, n_eval_episodes=5, callback=callback)
 
 
 @pytest.mark.slow
@@ -15,46 +28,18 @@ def test_identity(model_class):
 
     :param model_class: (BaseRLModel) the RL model
     """
-    model = model_class(policy="MlpPolicy", env=DummyVecEnv([lambda: IdentityEnv(dim=10)]))
-
-    env0 = IdentityEnv(dim=10)
-    env1 = DummyVecEnv([lambda: IdentityEnv(dim=10)])
-
-    n_trials = 100
-    for env, expected_shape in [(env0, ()), (env1, (1,))]:
-        obs = env.reset()
-        for _ in range(n_trials):
-            action, _ = model.predict(obs)
-            assert np.array(action).shape == expected_shape
-            obs, _, _, _ = env.step(action)
-
-    # Free memory
-    del model, env0, env1
+    check_shape(lambda: IdentityEnv(dim=10), model_class, (), (1,))
 
 
 @pytest.mark.slow
-@pytest.mark.parametrize("model_class", [A2C, DDPG, PPO1, PPO2, SAC, TRPO])
+@pytest.mark.parametrize("model_class", [A2C, DDPG, PPO1, PPO2, SAC, TRPO, TD3])
 def test_identity_box(model_class):
     """
     test the Box environment vectorisation detection
 
     :param model_class: (BaseRLModel) the RL model
     """
-    model = model_class(policy="MlpPolicy", env=DummyVecEnv([lambda: IdentityEnvBox(eps=0.5)]))
-
-    env0 = IdentityEnvBox()
-    env1 = DummyVecEnv([lambda: IdentityEnvBox(eps=0.5)])
-
-    n_trials = 100
-    for env, expected_shape in [(env0, (1,)), (env1, (1, 1))]:
-        obs = env.reset()
-        for _ in range(n_trials):
-            action, _ = model.predict(obs)
-            assert np.array(action).shape == expected_shape
-            obs, _, _, _ = env.step(action)
-
-    # Free memory
-    del model, env0, env1
+    check_shape(lambda: IdentityEnvBox(eps=0.5), model_class, (1,), (1, 1))
 
 
 @pytest.mark.slow
@@ -65,21 +50,7 @@ def test_identity_multi_binary(model_class):
 
     :param model_class: (BaseRLModel) the RL model
     """
-    model = model_class(policy="MlpPolicy", env=DummyVecEnv([lambda: IdentityEnvMultiBinary(dim=10)]))
-
-    env0 = IdentityEnvMultiBinary(dim=10)
-    env1 = DummyVecEnv([lambda: IdentityEnvMultiBinary(dim=10)])
-
-    n_trials = 100
-    for env, expected_shape in [(env0, (10,)), (env1, (1, 10))]:
-        obs = env.reset()
-        for _ in range(n_trials):
-            action, _ = model.predict(obs)
-            assert np.array(action).shape == expected_shape
-            obs, _, _, _ = env.step(action)
-
-    # Free memory
-    del model, env0, env1
+    check_shape(lambda: IdentityEnvMultiBinary(dim=10), model_class, (10,), (1, 10))
 
 
 @pytest.mark.slow
@@ -90,18 +61,4 @@ def test_identity_multi_discrete(model_class):
 
     :param model_class: (BaseRLModel) the RL model
     """
-    model = model_class(policy="MlpPolicy", env=DummyVecEnv([lambda: IdentityEnvMultiDiscrete(dim=10)]))
-
-    env0 = IdentityEnvMultiDiscrete(dim=10)
-    env1 = DummyVecEnv([lambda: IdentityEnvMultiDiscrete(dim=10)])
-
-    n_trials = 100
-    for env, expected_shape in [(env0, (2,)), (env1, (1, 2))]:
-        obs = env.reset()
-        for _ in range(n_trials):
-            action, _ = model.predict(obs)
-            assert np.array(action).shape == expected_shape
-            obs, _, _, _ = env.step(action)
-
-    # Free memory
-    del model, env0, env1
+    check_shape(lambda: IdentityEnvMultiDiscrete(dim=10), model_class, (2,), (1, 2))

tests/test_continuous.py

@@ -8,14 +8,14 @@
 from stable_baselines import A2C, ACKTR, SAC, DDPG, PPO1, PPO2, TRPO, TD3
 # TODO: add support for continuous actions
 # from stable_baselines.acer import ACER
-from stable_baselines.common import set_global_seeds
 from stable_baselines.common.vec_env import DummyVecEnv
 from stable_baselines.common.identity_env import IdentityEnvBox
 from stable_baselines.ddpg import AdaptiveParamNoiseSpec, NormalActionNoise
+from stable_baselines.common.evaluation import evaluate_policy
 from tests.test_common import _assert_eq
 
 
-N_TRIALS = 1000
+N_EVAL_EPISODES = 20
 NUM_TIMESTEPS = 15000
 
 MODEL_LIST = [
@@ -44,18 +44,10 @@ def test_model_manipulation(request, model_class):
         env = DummyVecEnv([lambda: IdentityEnvBox(eps=0.5)])
 
         # create and train
-        model = model_class(policy="MlpPolicy", env=env)
+        model = model_class(policy="MlpPolicy", env=env, seed=0)
         model.learn(total_timesteps=NUM_TIMESTEPS)
 
-        # predict and measure the acc reward
-        acc_reward = 0
-        set_global_seeds(0)
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, reward, _, _ = env.step(action)
-            acc_reward += reward
-        acc_reward = sum(acc_reward) / N_TRIALS
+        acc_reward, _ = evaluate_policy(model, env, n_eval_episodes=N_EVAL_EPISODES)
 
         # saving
         model_fname = './test_model_{}.zip'.format(request.node.name)
@@ -70,16 +62,9 @@ def test_model_manipulation(request, model_class):
         env = DummyVecEnv([lambda: IdentityEnvBox(eps=0.5)])
         model.set_env(env)
 
-        # predict the same output before saving
-        loaded_acc_reward = 0
-        set_global_seeds(0)
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, reward, _, _ = env.step(action)
-            loaded_acc_reward += reward
-        loaded_acc_reward = sum(loaded_acc_reward) / N_TRIALS
+        loaded_acc_reward, _ = evaluate_policy(model, env, n_eval_episodes=N_EVAL_EPISODES)
 
+        obs = env.reset()
         with pytest.warns(None) as record:
             act_prob = model.action_probability(obs)
 
@@ -124,27 +109,25 @@ def test_model_manipulation(request, model_class):
         # loaded_acc_reward = 0
         # set_global_seeds(0)
         # obs = env.reset()
-        # for _ in range(N_TRIALS):
+        # for _ in range(N_EVAL_EPISODES):
         #     action, _ = model.predict(obs)
         #     obs, reward, _, _ = env.step(action)
         #     loaded_acc_reward += reward
-        # loaded_acc_reward = sum(loaded_acc_reward) / N_TRIALS
+        # loaded_acc_reward = sum(loaded_acc_reward) / N_EVAL_EPISODES
         # # assert <10% diff
         # assert abs(acc_reward - loaded_acc_reward) / max(acc_reward, loaded_acc_reward) < 0.1, \
         #     "Error: the prediction seems to have changed between pre learning and post learning"
 
         # predict new values
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, _, _, _ = env.step(action)
+
+        evaluate_policy(model, env, n_eval_episodes=N_EVAL_EPISODES)
 
         # Free memory
         del model, env
 
     finally:
-        if os.path.exists("./test_model.zip"):
-            os.remove("./test_model.zip")
+        if os.path.exists(model_fname):
+            os.remove(model_fname)
 
 
 def test_ddpg():

tests/test_gail.py

@@ -8,8 +8,10 @@
  TD3, TRPO, SAC
 from stable_baselines.common.cmd_util import make_atari_env
 from stable_baselines.common.vec_env import VecFrameStack
+from stable_baselines.common.evaluation import evaluate_policy
 from stable_baselines.gail import ExpertDataset, generate_expert_traj
 
+
 EXPERT_PATH_PENDULUM = "stable_baselines/gail/dataset/expert_pendulum.npz"
 EXPERT_PATH_DISCRETE = "stable_baselines/gail/dataset/expert_cartpole.npz"
 
@@ -36,13 +38,7 @@ def test_gail(expert_env):
     model = model.load("GAIL-{}".format(env_id), env=env)
     model.learn(1000)
 
-    obs = env.reset()
-
-    for _ in range(1000):
-        action, _ = model.predict(obs)
-        obs, _, done, _ = env.step(action)
-        if done:
-            obs = env.reset()
+    evaluate_policy(model, env, n_eval_episodes=5)
     del dataset, model
 
 @pytest.mark.parametrize("generate_env", [

tests/test_identity.py

@@ -5,7 +5,7 @@
 from stable_baselines.ddpg import NormalActionNoise
 from stable_baselines.common.identity_env import IdentityEnv, IdentityEnvBox
 from stable_baselines.common.vec_env import DummyVecEnv
-from stable_baselines.common import set_global_seeds
+from stable_baselines.common.evaluation import evaluate_policy
 
 
 # Hyperparameters for learning identity for each RL model
@@ -39,24 +39,16 @@ def test_identity(model_name):
     env = DummyVecEnv([lambda: IdentityEnv(10)])
 
     model = LEARN_FUNC_DICT[model_name](env)
+    evaluate_policy(model, env, n_eval_episodes=20, reward_threshold=0.9)
 
-    n_trials = 1000
-    reward_sum = 0
-    set_global_seeds(0)
     obs = env.reset()
-    for _ in range(n_trials):
-        action, _ = model.predict(obs)
-        obs, reward, _, _ = env.step(action)
-        reward_sum += reward
-
     assert model.action_probability(obs).shape == (1, 10), "Error: action_probability not returning correct shape"
     action = env.action_space.sample()
     action_prob = model.action_probability(obs, actions=action)
     assert np.prod(action_prob.shape) == 1, "Error: not scalar probability"
     action_logprob = model.action_probability(obs, actions=action, logp=True)
     assert np.allclose(action_prob, np.exp(action_logprob)), (action_prob, action_logprob)
 
-    assert reward_sum > 0.9 * n_trials
     # Free memory
     del model, env
 
@@ -80,14 +72,6 @@ def test_identity_continuous(model_class):
                          action_noise=action_noise, buffer_size=int(1e6))
     model.learn(total_timesteps=20000)
 
-    n_trials = 1000
-    reward_sum = 0
-    set_global_seeds(0)
-    obs = env.reset()
-    for _ in range(n_trials):
-        action, _ = model.predict(obs)
-        obs, reward, _, _ = env.step(action)
-        reward_sum += reward
-    assert reward_sum > 0.9 * n_trials
+    evaluate_policy(model, env, n_eval_episodes=20, reward_threshold=0.9)
     # Free memory
     del model, env

tests/test_lstm_policy.py

@@ -11,6 +11,7 @@
 from stable_baselines.common.vec_env import SubprocVecEnv
 from stable_baselines.common.vec_env.vec_normalize import VecNormalize
 from stable_baselines.ppo2.ppo2 import safe_mean
+from stable_baselines.common.evaluation import evaluate_policy
 
 
 class CustomLSTMPolicy1(LstmPolicy):
@@ -85,11 +86,7 @@ def test_lstm_policy(request, model_class, policy):
         model.learn(total_timesteps=100)
 
         env = model.get_env()
-        # predict and measure the acc reward
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, _, _, _ = env.step(action)
+        evaluate_policy(model, env, n_eval_episodes=10)
         # saving
         model.save(model_fname)
         del model, env

tests/test_save.py

@@ -7,12 +7,12 @@
 import numpy as np
 
 from stable_baselines import A2C, ACER, ACKTR, DQN, PPO1, PPO2, TRPO
-from stable_baselines.common import set_global_seeds
 from stable_baselines.common.identity_env import IdentityEnv
 from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines.common.evaluation import evaluate_policy
 from stable_baselines.common.policies import MlpPolicy, FeedForwardPolicy
 
-N_TRIALS = 2000
+N_EVAL_EPISODES = 100
 
 MODEL_LIST = [
     A2C,
@@ -44,11 +44,11 @@ def test_model_manipulation(request, model_class, storage_method, store_format):
     works and that the action prediction works
 
     :param model_class: (BaseRLModel) A RL model
-    :param storage_method: (str) Should file be saved to a file ("path") or to a buffer 
+    :param storage_method: (str) Should file be saved to a file ("path") or to a buffer
         ("file-like")
     :param store_format: (str) Save format, either "zip" or "cloudpickle".
     """
-    
+
     # Use postfix ".model" so we can remove the file later
     model_fname = './test_model_{}.model'.format(request.node.name)
     store_as_cloudpickle = store_format == "cloudpickle"
@@ -57,20 +57,12 @@ def test_model_manipulation(request, model_class, storage_method, store_format):
         env = DummyVecEnv([lambda: IdentityEnv(10)])
 
         # create and train
-        model = model_class(policy="MlpPolicy", env=env)
-        model.learn(total_timesteps=50000)
+        model = model_class(policy="MlpPolicy", env=env, seed=0)
+        model.learn(total_timesteps=10000)
 
-        # predict and measure the acc reward
-        acc_reward = 0
-        set_global_seeds(0)
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            # Test action probability method
-            model.action_probability(obs)
-            obs, reward, _, _ = env.step(action)
-            acc_reward += reward
-        acc_reward = sum(acc_reward) / N_TRIALS
+        env.envs[0].action_space.seed(0)
+        mean_reward, _ = evaluate_policy(model, env, deterministic=True,
+                                         n_eval_episodes=N_EVAL_EPISODES)
 
         # test action probability for given (obs, action) pair
         env = model.get_env()
@@ -107,37 +99,24 @@ def test_model_manipulation(request, model_class, storage_method, store_format):
         model.set_env(env)
 
         # predict the same output before saving
-        loaded_acc_reward = 0
-        set_global_seeds(0)
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, reward, _, _ = env.step(action)
-            loaded_acc_reward += reward
-        loaded_acc_reward = sum(loaded_acc_reward) / N_TRIALS
-        assert abs(acc_reward - loaded_acc_reward) < 0.1, "Error: the prediction seems to have changed between " \
-                                                          "loading and saving"
+        env.envs[0].action_space.seed(0)
+        loaded_mean_reward, _ = evaluate_policy(model, env, deterministic=True, n_eval_episodes=N_EVAL_EPISODES)
+        # Allow 10% diff
+        assert abs((mean_reward - loaded_mean_reward) / mean_reward) < 0.1, "Error: the prediction seems to have changed between " \
+                                                                            "loading and saving"
 
         # learn post loading
         model.learn(total_timesteps=100)
 
         # validate no reset post learning
-        loaded_acc_reward = 0
-        set_global_seeds(0)
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, reward, _, _ = env.step(action)
-            loaded_acc_reward += reward
-        loaded_acc_reward = sum(loaded_acc_reward) / N_TRIALS
-        assert abs(acc_reward - loaded_acc_reward) < 0.1, "Error: the prediction seems to have changed between " \
-                                                          "pre learning and post learning"
+        env.envs[0].action_space.seed(0)
+        loaded_mean_reward, _ = evaluate_policy(model, env, deterministic=True, n_eval_episodes=N_EVAL_EPISODES)
+
+        assert abs((mean_reward - loaded_mean_reward) / mean_reward) < 0.15, "Error: the prediction seems to have changed between " \
+                                                                            "pre learning and post learning"
 
         # predict new values
-        obs = env.reset()
-        for _ in range(N_TRIALS):
-            action, _ = model.predict(obs)
-            obs, _, _, _ = env.step(action)
+        evaluate_policy(model, env, n_eval_episodes=N_EVAL_EPISODES)
 
         del model, env
 

tests/test_utils.py

@@ -0,0 +1,78 @@
+import os
+import shutil
+
+import pytest
+import gym
+
+from stable_baselines import A2C
+from stable_baselines.bench.monitor import Monitor
+from stable_baselines.common.evaluation import evaluate_policy
+from stable_baselines.common.cmd_util import make_vec_env
+from stable_baselines.common.vec_env import DummyVecEnv, SubprocVecEnv
+
+
+@pytest.mark.parametrize("env_id", ['CartPole-v1', lambda: gym.make('CartPole-v1')])
+@pytest.mark.parametrize("n_envs", [1, 2])
+@pytest.mark.parametrize("vec_env_cls", [None, SubprocVecEnv])
+@pytest.mark.parametrize("wrapper_class", [None, gym.wrappers.TimeLimit])
+def test_make_vec_env(env_id, n_envs, vec_env_cls, wrapper_class):
+    env = make_vec_env(env_id, n_envs, vec_env_cls=vec_env_cls,
+                       wrapper_class=wrapper_class, monitor_dir=None, seed=0)
+
+    assert env.num_envs == n_envs
+
+    if vec_env_cls is None:
+        assert isinstance(env, DummyVecEnv)
+        if wrapper_class is not None:
+            assert isinstance(env.envs[0], wrapper_class)
+        else:
+            assert isinstance(env.envs[0], Monitor)
+    else:
+        assert isinstance(env, SubprocVecEnv)
+    # Kill subprocesses
+    env.close()
+
+
+def test_custom_vec_env():
+    """
+    Stand alone test for a special case (passing a custom VecEnv class) to avoid doubling the number of tests.
+    """
+    monitor_dir = 'logs/test_make_vec_env/'
+    env = make_vec_env('CartPole-v1', n_envs=1,
+                       monitor_dir=monitor_dir, seed=0,
+                       vec_env_cls=SubprocVecEnv, vec_env_kwargs={'start_method': None})
+
+
+    assert env.num_envs == 1
+    assert isinstance(env, SubprocVecEnv)
+    assert os.path.isdir('logs/test_make_vec_env/')
+    # Kill subprocess
+    env.close()
+    # Cleanup folder
+    shutil.rmtree(monitor_dir)
+
+    # This should fail because DummyVecEnv does not have any keyword argument
+    with pytest.raises(TypeError):
+        make_vec_env('CartPole-v1', n_envs=1, vec_env_kwargs={'dummy': False})
+
+
+def test_evaluate_policy():
+    model = A2C('MlpPolicy', 'Pendulum-v0', seed=0)
+    n_steps_per_episode, n_eval_episodes = 200, 2
+    model.n_callback_calls = 0
+
+    def dummy_callback(locals_, _globals):
+        locals_['model'].n_callback_calls += 1
+
+    _, n_steps = evaluate_policy(model, model.get_env(), n_eval_episodes, deterministic=True,
+                                 render=False, callback=dummy_callback, reward_threshold=None,
+                                 return_episode_rewards=False)
+    assert n_steps == n_steps_per_episode * n_eval_episodes
+    assert n_steps == model.n_callback_calls
+
+    # Reaching a mean reward of zero is impossible with the Pendulum env
+    with pytest.raises(AssertionError):
+        evaluate_policy(model, model.get_env(), n_eval_episodes, reward_threshold=0.0)
+
+    episode_rewards, _ = evaluate_policy(model, model.get_env(), n_eval_episodes, return_episode_rewards=True)
+    assert len(episode_rewards) == n_eval_episodes