From 8831eff163bc061727e16b7f728162776dc5ba27 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 7 Jan 2020 14:00:03 +0100 Subject: [PATCH] Unify evaluation --- torchy_baselines/cem_rl/cem_rl.py | 5 ++-- torchy_baselines/common/base_class.py | 34 ++++++++++++++++++---- torchy_baselines/common/evaluation.py | 42 +++++++++++++++++++++++---- torchy_baselines/ppo/ppo.py | 21 ++++---------- torchy_baselines/sac/sac.py | 13 ++------- torchy_baselines/td3/td3.py | 15 ++-------- 6 files changed, 80 insertions(+), 50 deletions(-) diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 41efdf3..e3a322b 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/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 diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 6200bc1..041ce83 100644 --- a/torchy_baselines/common/base_class.py +++ b/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 diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py index bc8fc70..04ae296 100644 --- a/torchy_baselines/common/evaluation.py +++ b/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 diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 66f459d..8641552 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/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 diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index e4f24d2..f3f45fb 100644 --- a/torchy_baselines/sac/sac.py +++ b/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 diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 9b5cafd..a24e0db 100644 --- a/torchy_baselines/td3/td3.py +++ b/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