From 91bbc28c0f01b8a50cea61c113186c5ab8ac3068 Mon Sep 17 00:00:00 2001
From: Adam Gleave <adam@gleave.me>
Date: Tue, 7 Jul 2020 18:39:55 -0700
Subject: [PATCH] Address minor issues after clarification by @araffin

---
 stable_baselines3/common/base_class.py    | 10 ++++++----
 stable_baselines3/common/distributions.py | 14 ++++++++++++++
 stable_baselines3/common/policies.py      |  3 +++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py
index f66220e..4e51367 100644
--- a/stable_baselines3/common/base_class.py
+++ b/stable_baselines3/common/base_class.py
@@ -320,8 +320,8 @@ class BaseAlgorithm(ABC):
                              f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}")
 
         # check if observation space and action space are part of the saved parameters
-        if ("observation_space" not in data or "action_space" not in data) and "env" not in data:
-            raise ValueError("The observation_space and action_space was not given, can't verify new environments")
+        if "observation_space" not in data or "action_space" not in data:
+            raise KeyError("The observation_space and action_space were not given, can't verify new environments")
         # check if given env is valid
         if env is not None:
             check_for_correct_spaces(env, data["observation_space"], data["action_space"])
@@ -425,8 +425,10 @@ class BaseAlgorithm(ABC):
         :return: (Tuple[int, BaseCallback])
         """
         self.start_time = time.time()
-        self.ep_info_buffer = deque(maxlen=100)
-        self.ep_success_buffer = deque(maxlen=100)
+        if self.ep_info_buffer is None or reset_num_timesteps:
+            # Initialize buffers if they don't exist, or reinitialize if resetting counters
+            self.ep_info_buffer = deque(maxlen=100)
+            self.ep_success_buffer = deque(maxlen=100)
 
         if self.action_noise is not None:
             self.action_noise.reset()
diff --git a/stable_baselines3/common/distributions.py b/stable_baselines3/common/distributions.py
index 6fa70f8..1a20928 100644
--- a/stable_baselines3/common/distributions.py
+++ b/stable_baselines3/common/distributions.py
@@ -17,6 +17,20 @@ class Distribution(ABC):
     def __init__(self):
         super(Distribution, self).__init__()
 
+    @abstractmethod
+    def proba_distribution_net(self, *args, **kwargs):
+        """Create the layers and parameters that represent the distribution.
+
+        Subclasses must define this, but the arguments and return type vary between
+        concrete classes."""
+
+    @abstractmethod
+    def proba_distribution(self, *args, **kwargs) -> 'Distribution':
+        """Set parameters of the distribution.
+
+        :return: (Distribution) self
+        """
+
     @abstractmethod
     def log_prob(self, x: th.Tensor) -> th.Tensor:
         """
diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py
index 2aeffd9..7342c46 100644
--- a/stable_baselines3/common/policies.py
+++ b/stable_baselines3/common/policies.py
@@ -140,6 +140,7 @@ class BasePolicy(nn.Module, ABC):
         :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state
             (used in recurrent policies)
         """
+        # TODO (GH/1): add support for RNN policies
         # if state is None:
         #     state = self.initial_state
         # if mask is None:
@@ -438,6 +439,8 @@ class ActorCriticPolicy(BasePolicy):
             self.action_net = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi)
         elif isinstance(self.action_dist, BernoulliDistribution):
             self.action_net = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi)
+        else:
+            raise NotImplementedError(f"Unsupported distribution '{self.action_dist}'.")
 
         self.value_net = nn.Linear(self.mlp_extractor.latent_dim_vf, 1)
         # Init weights: use orthogonal initialization