Training and evaluation: call model.train() and model.eval() (#537)

* training and evaluation: call model.train() and model.eval() to enable and disable dropout and batchnorm

* Add comment documentation

* Fix train and eval for the Actor class

* Run black

* Add github handle to changelog

* Add unit tests for PPO and DQN

* Refactor unit test

* Run black

* unit test: add a dropout layer and check that calling predict with deterministic=True is deterministic

* documentation: add bugfix description to changelog

* unit test: use learning_starts=0, decrease the size of the network and use more training steps

* on policy algorithms: call policy.train() and policy.eval() instead of disable_training and enable_training as it is a th.nn.module

* Rename unit test

* unit test: use drop out probability of 0.5

* Call policy.train and policy.eval

* Fixes + update tests

* Remove unneeded eval

Co-authored-by: David Blom <davidsblom@gmail.com>
Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.2.0a1 (WIP)
+Release 1.2.0a2 (WIP)
 ---------------------------
 
 Breaking Changes:
@@ -16,6 +16,7 @@ New Features:
 
 Bug Fixes:
 ^^^^^^^^^^
+- Fixed model predictions when using batch normalization and dropout layers by calling ``train()`` and ``eval()`` (@davidblom603)
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -737,4 +738,4 @@ And all the contributors:
 @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray
 @tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @JadenTravnik @decodyng @ardabbour @lorenz-h @mschweizer @lorepieri8 @vwxyzjn
 @ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro @amy12xx @juancroldan @benblack769 @bstee615
-@c-rizz @skandermoalla @MihaiAnca13
+@c-rizz @skandermoalla @MihaiAnca13 @davidblom603

stable_baselines3/a2c/a2c.py

@@ -119,6 +119,9 @@ class A2C(OnPolicyAlgorithm):
         Update policy using the currently gathered
         rollout buffer (one gradient step over whole data).
         """
+        # Switch to train mode (this affects batch norm / dropout)
+        self.policy.train()
+
         # Update optimizer learning rate
         self._update_learning_rate(self.policy.optimizer)
 

stable_baselines3/common/base_class.py

@@ -399,6 +399,7 @@ class BaseAlgorithm(ABC):
         :return:
         """
         self.start_time = time.time()
+
         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)

stable_baselines3/common/off_policy_algorithm.py

@@ -536,6 +536,9 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         :param log_interval: Log data every ``log_interval`` episodes
         :return:
         """
+        # Switch to eval mode (this affects batch norm / dropout)
+        self.policy.eval()
+
         episode_rewards, total_timesteps = [], []
         num_collected_steps, num_collected_episodes = 0, 0
 

stable_baselines3/common/on_policy_algorithm.py

@@ -147,6 +147,9 @@ class OnPolicyAlgorithm(BaseAlgorithm):
             collected, False if callback terminated rollout prematurely.
         """
         assert self._last_obs is not None, "No previous observation was provided"
+        # Switch to eval mode (this affects batch norm / dropout)
+        self.policy.eval()
+
         n_steps = 0
         rollout_buffer.reset()
         # Sample new weights for the state dependent exploration

stable_baselines3/common/policies.py

@@ -267,6 +267,8 @@ class BasePolicy(BaseModel):
         #     state = self.initial_state
         # if mask is None:
         #     mask = [False for _ in range(self.n_envs)]
+        # Switch to eval mode (this affects batch norm / dropout)
+        self.eval()
 
         vectorized_env = False
         if isinstance(observation, dict):

stable_baselines3/dqn/dqn.py

@@ -152,6 +152,8 @@ class DQN(OffPolicyAlgorithm):
         self.logger.record("rollout/exploration rate", self.exploration_rate)
 
     def train(self, gradient_steps: int, batch_size: int = 100) -> None:
+        # Switch to train mode (this affects batch norm / dropout)
+        self.policy.train()
         # Update learning rate according to schedule
         self._update_learning_rate(self.policy.optimizer)
 

stable_baselines3/ppo/ppo.py

@@ -166,6 +166,8 @@ class PPO(OnPolicyAlgorithm):
         """
         Update policy using the currently gathered rollout buffer.
         """
+        # Switch to train mode (this affects batch norm / dropout)
+        self.policy.train()
         # Update optimizer learning rate
         self._update_learning_rate(self.policy.optimizer)
         # Compute current clip range

stable_baselines3/sac/sac.py

@@ -180,6 +180,8 @@ class SAC(OffPolicyAlgorithm):
         self.critic_target = self.policy.critic_target
 
     def train(self, gradient_steps: int, batch_size: int = 64) -> None:
+        # Switch to train mode (this affects batch norm / dropout)
+        self.policy.train()
         # Update optimizers learning rate
         optimizers = [self.actor.optimizer, self.critic.optimizer]
         if self.ent_coef_optimizer is not None:

stable_baselines3/td3/td3.py

@@ -132,6 +132,8 @@ class TD3(OffPolicyAlgorithm):
         self.critic_target = self.policy.critic_target
 
     def train(self, gradient_steps: int, batch_size: int = 100) -> None:
+        # Switch to train mode (this affects batch norm / dropout)
+        self.policy.train()
 
         # Update learning rate according to lr schedule
         self._update_learning_rate([self.actor.optimizer, self.critic.optimizer])

stable_baselines3/version.txt

@@ -1 +1 @@
-1.2.0a1
+1.2.0a2

tests/test_predict.py

@@ -1,8 +1,12 @@
 import gym
+import numpy as np
 import pytest
 import torch as th
+import torch.nn as nn
 
 from stable_baselines3 import A2C, DQN, PPO, SAC, TD3
+from stable_baselines3.common.preprocessing import get_flattened_obs_dim
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
 from stable_baselines3.common.utils import get_device
 from stable_baselines3.common.vec_env import DummyVecEnv
 
@@ -69,3 +73,74 @@ def test_predict(model_class, env_id, device):
 
         action, _ = model.predict(vec_env_obs, deterministic=False)
         assert action.shape[0] == vec_env_obs.shape[0]
+
+
+class FlattenBatchNormExtractor(BaseFeaturesExtractor):
+    """
+    Feature extract that flatten the input and uses batch normalization.
+    Used as a placeholder when feature extraction is not needed.
+
+    :param observation_space:
+    """
+
+    def __init__(self, observation_space: gym.Space):
+        super(FlattenBatchNormExtractor, self).__init__(observation_space, get_flattened_obs_dim(observation_space))
+        self.flatten = nn.Flatten()
+        self.batch_norm = nn.BatchNorm1d(self._features_dim)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, observations: th.Tensor) -> th.Tensor:
+        result = self.flatten(observations)
+        result = self.batch_norm(result)
+        result = self.dropout(result)
+        return result
+
+
+@pytest.mark.parametrize("model_class", MODEL_LIST)
+@pytest.mark.parametrize("env_id", ["Pendulum-v0", "CartPole-v1"])
+def test_batch_norm_dropout(model_class, env_id):
+
+    if env_id == "CartPole-v1":
+        if model_class in [SAC, TD3]:
+            return
+    elif model_class in [DQN]:
+        return
+
+    model_kwargs = dict(seed=1)
+
+    if model_class in [DQN, TD3, SAC]:
+        model_kwargs["learning_starts"] = 0
+    else:
+        model_kwargs["n_steps"] = 64
+
+    policy_kwargs = dict(
+        features_extractor_class=FlattenBatchNormExtractor,
+        net_arch=[16, 16],
+    )
+    model = model_class("MlpPolicy", env_id, policy_kwargs=policy_kwargs, verbose=1, **model_kwargs)
+
+    if model_class in [SAC, TD3]:
+        batch_norm = model.policy.actor.features_extractor.batch_norm
+    elif model_class in [PPO, A2C]:
+        batch_norm = model.policy.features_extractor.batch_norm
+    else:
+        # DQN
+        batch_norm = model.policy.q_net.features_extractor.batch_norm
+
+    # batch norm param before training
+    bias_before_learn = batch_norm.bias.detach().cpu().numpy().copy()
+    running_mean_before_learn = batch_norm.running_mean.detach().cpu().numpy().copy()
+    model.learn(100)
+    env = model.get_env()
+    observation = env.reset()
+
+    bias_after_learn = batch_norm.bias.detach().cpu().numpy()
+    running_mean_after_learn = batch_norm.running_mean.detach().cpu().numpy().copy()
+
+    # Run twice on the same observation to test if it is deterministic
+    first_prediction, _ = model.predict(observation, deterministic=True)
+    second_prediction, _ = model.predict(observation, deterministic=True)
+
+    np.testing.assert_allclose(first_prediction, second_prediction)
+    assert not np.allclose(bias_before_learn, bias_after_learn)
+    assert not np.allclose(running_mean_before_learn, running_mean_after_learn)