Add eval success rate logging (#255)

* Add eval success rate logging

* Fix name clash

* Log data

* Bump version

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:
 ^^^^^^^^^^

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!")

stable_baselines3/version.txt

@@ -1 +1 @@
-0.11.0a1
+0.11.0a2

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