Unify evaluation

torchy_baselines/cem_rl/cem_rl.py

@@ -155,11 +155,12 @@ class CEMRL(TD3):
                 self.actor.load_from_vector(self.es.mu)
                 sync_envs_normalization(self.env, eval_env)
 
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
+                mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
                 evaluations.append(mean_reward)
 
                 if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
+                    print("Eval num_timesteps={}, "
+                          "episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
                     print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
 
             actor_steps = 0

torchy_baselines/common/base_class.py

@@ -9,11 +9,12 @@ import gym
 import torch as th
 import numpy as np
 
+from torchy_baselines.common import logger
 from torchy_baselines.common.policies import get_policy_from_name
 from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
-from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
+from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization
 from torchy_baselines.common.monitor import Monitor
-from torchy_baselines.common import logger
+from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.save_util import data_to_json, json_to_data
 
 
@@ -287,9 +288,9 @@ class BaseRLModel(object):
         :param log_interval: (int) The number of timesteps before logging.
         :param tb_log_name: (str) the name of the run for tensorboard log
         :param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging)
-        :param eval_env: (gym.Env)
-        :param eval_freq: (int)
-        :param n_eval_episodes: (int)
+        :param eval_env: (gym.Env) Environment that will be used to evaluate the agent
+        :param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
+        :param n_eval_episodes: (int) Number of episode to evaluate the agent
         :return: (BaseRLModel) the trained model
         """
         pass
@@ -712,3 +713,26 @@ class BaseRLModel(object):
         params_to_save = self.get_policy_parameters()
         opt_params_to_save = self.get_opt_parameters()
         self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
+
+    def _eval_policy(self, eval_freq, eval_env, n_eval_episodes,
+                     timesteps_since_eval, deterministic=True):
+        """
+        Evaluate the current policy on a test environment.
+
+        :param eval_env: (gym.Env) Environment that will be used to evaluate the agent
+        :param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
+        :param n_eval_episodes: (int) Number of episode to evaluate the agent
+        :parma timesteps_since_eval: (int) Number of timesteps since last evaluation
+        :param deterministic: (bool) Whether to use deterministic or stochastic actions
+        :return: (int) Number of timesteps since last evaluation
+        """
+        if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
+            timesteps_since_eval %= eval_freq
+            # Synchronise the normalization stats if needed
+            sync_envs_normalization(self.env, eval_env)
+            mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes, deterministic=deterministic)
+            if self.verbose > 0:
+                print("Eval num_timesteps={}, "
+                      "episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
+                print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+        return timesteps_since_eval

torchy_baselines/common/evaluation.py

@@ -1,10 +1,34 @@
+# Copied from stable_baselines
 import numpy as np
 
+from torchy_baselines.common.vec_env import VecEnv
 
-def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=False):
+
+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 episodes and returns average reward
+    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 environment 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, float) Mean reward per episode, std of reward per episode
+        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()
@@ -12,11 +36,19 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=F
         episode_reward = 0.0
         while not done:
             action = model.predict(obs, deterministic=deterministic)
-            obs, reward, done, _ = env.step(action)
+            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)
-
-    return np.mean(episode_rewards), np.std(episode_rewards)
+    mean_reward = np.mean(episode_rewards)
+    std_reward = np.std(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, std_reward

torchy_baselines/ppo/ppo.py

@@ -14,10 +14,8 @@ except ImportError:
 import numpy as np
 
 from torchy_baselines.common.base_class import BaseRLModel
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.buffers import RolloutBuffer
 from torchy_baselines.common.utils import explained_variance, get_schedule_fn
-from torchy_baselines.common.vec_env import sync_envs_normalization
 from torchy_baselines.common import logger
 from torchy_baselines.ppo.policies import PPOPolicy
 
@@ -303,19 +301,12 @@ class PPO(BaseRLModel):
 
             self.train(self.n_epochs, batch_size=self.batch_size)
 
-            # Evaluate agent
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-                sync_envs_normalization(self.env, eval_env)
-
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
-                if self.tb_writer is not None:
-                    self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
-
-                evaluations.append(mean_reward)
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+            # Evaluate the agent
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
+            # For tensorboard integration
+            # if self.tb_writer is not None:
+            #     self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
 
         return self
 

torchy_baselines/sac/sac.py

@@ -6,9 +6,7 @@ import numpy as np
 
 from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.buffers import ReplayBuffer
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.sac.policies import SACPolicy
-from torchy_baselines.common.vec_env import sync_envs_normalization
 from torchy_baselines.common import logger
 
 
@@ -287,15 +285,8 @@ class SAC(BaseRLModel):
 
                 self.train(gradient_steps, batch_size=self.batch_size)
 
-            # Evaluate episode
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-                sync_envs_normalization(self.env, eval_env)
-                mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
-                evaluations.append(mean_reward)
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}, std_reward={:.2f}".format(self.num_timesteps, mean_reward, std_reward))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
 
         return self
 

torchy_baselines/td3/td3.py

@@ -6,10 +6,7 @@ import numpy as np
 
 from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.buffers import ReplayBuffer
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.td3.policies import TD3Policy
-from torchy_baselines.common.vec_env import sync_envs_normalization
-from torchy_baselines.ppo.policies import MlpPolicy
 
 
 class TD3(BaseRLModel):
@@ -290,15 +287,9 @@ class TD3(BaseRLModel):
                 gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
                 self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
 
-            # Evaluate episode
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-                sync_envs_normalization(self.env, eval_env)
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
-                evaluations.append(mean_reward)
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+            # Evaluate the agent
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
 
         return self