From 3efab0d267e74cb03264411d4500ddde0c163404 Mon Sep 17 00:00:00 2001 From: David Blom <63226970+davidblom603@users.noreply.github.com> Date: Sat, 14 Aug 2021 14:08:27 +0200 Subject: [PATCH] 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 Co-authored-by: Antonin Raffin --- docs/misc/changelog.rst | 5 +- stable_baselines3/a2c/a2c.py | 3 + stable_baselines3/common/base_class.py | 1 + .../common/off_policy_algorithm.py | 3 + .../common/on_policy_algorithm.py | 3 + stable_baselines3/common/policies.py | 2 + stable_baselines3/dqn/dqn.py | 2 + stable_baselines3/ppo/ppo.py | 2 + stable_baselines3/sac/sac.py | 2 + stable_baselines3/td3/td3.py | 2 + stable_baselines3/version.txt | 2 +- tests/test_predict.py | 75 +++++++++++++++++++ 12 files changed, 99 insertions(+), 3 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 6513f5c..6a1965a 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index 03b1fc8..88c8992 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/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) diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index f8531d6..6b0df32 100644 --- a/stable_baselines3/common/base_class.py +++ b/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) diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index dce4209..ae8569f 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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 diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 5d872a9..eb3417c 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/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 diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 8e5394f..377b548 100644 --- a/stable_baselines3/common/policies.py +++ b/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): diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index d68a643..9f5214b 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/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) diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index 28f8777..b3a8e99 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/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 diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 9fcaff9..f53e399 100644 --- a/stable_baselines3/sac/sac.py +++ b/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: diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index ef652c0..d8ad25d 100644 --- a/stable_baselines3/td3/td3.py +++ b/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]) diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 4597d21..7ce60fd 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.2.0a1 +1.2.0a2 diff --git a/tests/test_predict.py b/tests/test_predict.py index 2927796..689c75d 100644 --- a/tests/test_predict.py +++ b/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)