From 6b598323ae070bb0a998d25230f6e11eca4cbe61 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 8 Dec 2020 15:49:07 +0100 Subject: [PATCH] Add eval success rate logging (#255) * Add eval success rate logging * Fix name clash * Log data * Bump version --- docs/misc/changelog.rst | 3 +- stable_baselines3/common/callbacks.py | 40 +++++++++++++++++++++++++++ stable_baselines3/version.txt | 2 +- tests/test_callbacks.py | 23 ++++++++++++++- 4 files changed, 65 insertions(+), 3 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index a690651..fc616ee 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Pre-Release 0.11.0a1 (WIP) +Pre-Release 0.11.0a2 (WIP) ------------------------------- Breaking Changes: @@ -25,6 +25,7 @@ New Features: with given Gym wrappers. - Added ``monitor_kwargs`` parameter to ``make_vec_env`` and ``make_atari_env`` - Wrap the environments automatically with a ``Monitor`` wrapper when possible. +- ``EvalCallback`` now logs the success rate when available (``is_success`` must be present in the info dict) Bug Fixes: ^^^^^^^^^^ diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 34aac9b..8e89d91 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -317,6 +317,9 @@ class EvalCallback(EventCallback): self.evaluations_results = [] self.evaluations_timesteps = [] self.evaluations_length = [] + # For computing success rate + self._is_success_buffer = [] + self.evaluations_successes = [] def _init_callback(self) -> None: # Does not work in some corner cases, where the wrapper is not the same @@ -329,12 +332,34 @@ class EvalCallback(EventCallback): if self.log_path is not None: os.makedirs(os.path.dirname(self.log_path), exist_ok=True) + def _log_success_callback(self, locals_: Dict[str, Any], globals_: Dict[str, Any]) -> None: + """ + Callback passed to the ``evaluate_policy`` function + in order to log the success rate (when applicable), + for instance when using HER. + + :param locals_: + :param globals_: + """ + info = locals_["info"] + # VecEnv: unpack + if not isinstance(info, dict): + info = info[0] + + if locals_["done"]: + maybe_is_success = info.get("is_success") + if maybe_is_success is not None: + self._is_success_buffer.append(maybe_is_success) + def _on_step(self) -> bool: if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0: # Sync training and eval env if there is VecNormalize sync_envs_normalization(self.training_env, self.eval_env) + # Reset success rate buffer + self._is_success_buffer = [] + episode_rewards, episode_lengths = evaluate_policy( self.model, self.eval_env, @@ -343,17 +368,26 @@ class EvalCallback(EventCallback): deterministic=self.deterministic, return_episode_rewards=True, warn=self.warn, + callback=self._log_success_callback, ) if self.log_path is not None: self.evaluations_timesteps.append(self.num_timesteps) self.evaluations_results.append(episode_rewards) self.evaluations_length.append(episode_lengths) + + kwargs = {} + # Save success log if present + if len(self._is_success_buffer) > 0: + self.evaluations_successes.append(self._is_success_buffer) + kwargs = dict(successes=self.evaluations_successes) + np.savez( self.log_path, timesteps=self.evaluations_timesteps, results=self.evaluations_results, ep_lengths=self.evaluations_length, + **kwargs, ) mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards) @@ -367,6 +401,12 @@ class EvalCallback(EventCallback): self.logger.record("eval/mean_reward", float(mean_reward)) self.logger.record("eval/mean_ep_length", mean_ep_length) + if len(self._is_success_buffer) > 0: + success_rate = np.mean(self._is_success_buffer) + if self.verbose > 0: + print(f"Success rate: {100 * success_rate:.2f}%") + self.logger.record("eval/success_rate", success_rate) + if mean_reward > self.best_mean_reward: if self.verbose > 0: print("New best mean reward!") diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index e0cbcd5..a09c7eb 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.11.0a1 +0.11.0a2 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 144494a..d86a4d6 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -2,9 +2,11 @@ import os import shutil import gym +import numpy as np import pytest -from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3 +from stable_baselines3 import A2C, DDPG, DQN, HER, PPO, SAC, TD3 +from stable_baselines3.common.bit_flipping_env import BitFlippingEnv from stable_baselines3.common.callbacks import ( CallbackList, CheckpointCallback, @@ -14,6 +16,8 @@ from stable_baselines3.common.callbacks import ( StopTrainingOnRewardThreshold, ) from stable_baselines3.common.env_util import make_vec_env +from stable_baselines3.common.vec_env import DummyVecEnv +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper @pytest.mark.parametrize("model_class", [A2C, PPO, SAC, TD3, DQN, DDPG]) @@ -97,3 +101,20 @@ def select_env(model_class) -> str: return "CartPole-v0" else: return "Pendulum-v0" + + +def test_eval_success_logging(tmp_path): + n_bits = 2 + env = BitFlippingEnv(n_bits=n_bits) + eval_env = DummyVecEnv([lambda: BitFlippingEnv(n_bits=n_bits)]) + eval_callback = EvalCallback( + ObsDictWrapper(eval_env), + eval_freq=250, + log_path=tmp_path, + warn=False, + ) + model = HER("MlpPolicy", env, DQN, learning_starts=100, seed=0, max_episode_length=n_bits) + model.learn(500, callback=eval_callback) + assert len(eval_callback._is_success_buffer) > 0 + # More than 50% success rate + assert np.mean(eval_callback._is_success_buffer) > 0.5