From b1aee717729e84d0e5a05e37324a6b1521db6e1b Mon Sep 17 00:00:00 2001 From: Cody Wild Date: Mon, 11 Jan 2021 08:03:32 -0800 Subject: [PATCH] Improve error messages when PPO effective batch size is 1 and when last mini-batch is truncated (#270) * Add warning about total_env_steps not dividing neatly into batch size * Stylistic cleanup * Black reformatting * Add clearer documentation and update changelog * Update changelog.rst * Use specific RolloutBuffer terminology Co-authored-by: Antonin RAFFIN * Change to minibatch language Co-authored-by: Antonin RAFFIN * Cleaning up language describing rollout buffer requirements Co-authored-by: Antonin RAFFIN * Switch to using env.num_envs * Working tests * Black and isort still fighting each other * codestyle finally happy * Basic test exists, possibly in the wrong file * Update phrasing Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 6 +++++- stable_baselines3/ppo/ppo.py | 24 ++++++++++++++++++++++-- tests/test_utils.py | 22 ++++++++++++++++++++-- 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 401f414..9862ad4 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -35,6 +35,9 @@ Bug Fixes: - Fixed bug that the arguments order of ``explained_variance()`` in ``ppo.py`` and ``a2c.py`` is not correct (@thisray) - Fixed bug where full ``HerReplayBuffer`` leads to an index error. (@megan-klaiber) - Fixed bug where replay buffer could not be saved if it was too big (> 4 Gb) for python<3.8 (thanks @hn2) +- Added informative ``PPO`` construction error in edge-case scenario where ``n_steps * n_envs = 1`` (size of rollout buffer), + which otherwise causes downstream breaking errors in training (@decodyng) + Deprecations: ^^^^^^^^^^^^^ @@ -48,6 +51,7 @@ Others: - Renamed variables in the ``train()`` method of ``SAC``, ``TD3`` and ``DQN`` to match SB3-Contrib. - Updated docker base image to Ubuntu 18.04 - Set tensorboard min version to 2.2.0 (earlier version are apparently not working with PyTorch) +- Added warning for ``PPO`` when ``n_steps * n_envs`` is not a multiple of ``batch_size`` (last mini-batch truncated) (@decodyng) Documentation: ^^^^^^^^^^^^^^ @@ -535,4 +539,4 @@ And all the contributors: @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray -@tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio +@tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @decodyng diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index 52579b8..ca22da6 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -1,3 +1,4 @@ +import warnings from typing import Any, Dict, Optional, Type, Union import numpy as np @@ -28,7 +29,9 @@ class PPO(OnPolicyAlgorithm): :param learning_rate: The learning rate, it can be a function of the current progress remaining (from 1 to 0) :param n_steps: The number of steps to run for each environment per update - (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel) + (i.e. rollout buffer size is n_steps * n_envs where n_envs is number of environment copies running in parallel) + NOTE: n_steps * n_envs must be greater than 1 (because of the advantage normalization) + See https://github.com/pytorch/pytorch/issues/29372 :param batch_size: Minibatch size :param n_epochs: Number of epoch when optimizing the surrogate loss :param gamma: Discount factor @@ -115,7 +118,24 @@ class PPO(OnPolicyAlgorithm): spaces.MultiBinary, ), ) - + if self.env is not None: + # Check that `n_steps * n_envs > 1` to avoid NaN + # when doing advantage normalization + buffer_size = self.env.num_envs * self.n_steps + assert ( + buffer_size > 1 + ), f"`n_steps * n_envs` must be greater than 1. Currently n_steps={self.n_steps} and n_envs={self.env.num_envs}" + # Check that the rollout buffer size is a multiple of the mini-batch size + untruncated_batches = buffer_size // batch_size + if buffer_size % batch_size > 0: + warnings.warn( + f"You have specified a mini-batch size of {batch_size}," + f" but because the `RolloutBuffer` is of size `n_steps * n_envs = {buffer_size}`," + f" after every {untruncated_batches} untruncated mini-batches," + f" there will be a truncated mini-batch of size {buffer_size % batch_size}\n" + f"We recommend using a `batch_size` that is a multiple of `n_steps * n_envs`.\n" + f"Info: (n_steps={self.n_steps} and n_envs={self.env.num_envs})" + ) self.batch_size = batch_size self.n_epochs = n_epochs self.clip_range = clip_range diff --git a/tests/test_utils.py b/tests/test_utils.py index 3684900..1e03089 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,7 +6,7 @@ import numpy as np import pytest import torch as th -from stable_baselines3 import A2C +from stable_baselines3 import A2C, PPO from stable_baselines3.common.atari_wrappers import ClipRewardEnv from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_vec_env, unwrap_wrapper from stable_baselines3.common.evaluation import evaluate_policy @@ -80,7 +80,12 @@ def test_vec_env_monitor_kwargs(): env = make_vec_env("MountainCarContinuous-v0", n_envs=1, seed=0, monitor_kwargs={"allow_early_resets": True}) assert env.get_attr("allow_early_resets")[0] is True - env = make_atari_env("BreakoutNoFrameskip-v4", n_envs=1, seed=0, monitor_kwargs={"allow_early_resets": True}) + env = make_atari_env( + "BreakoutNoFrameskip-v4", + n_envs=1, + seed=0, + monitor_kwargs={"allow_early_resets": True}, + ) assert env.get_attr("allow_early_resets")[0] is True @@ -333,3 +338,16 @@ def test_is_wrapped(): assert is_wrapped(env, Monitor) # Test that unwrap works as expected assert unwrap_wrapper(env, Monitor) == monitor_env + + +def test_ppo_warnings(): + """Test that PPO warns and errors correctly on + problematic rollour buffer sizes""" + + # Only 1 step: advantage normalization will return NaN + with pytest.raises(AssertionError): + PPO("MlpPolicy", "Pendulum-v0", n_steps=1) + + # Truncated mini-batch + with pytest.warns(UserWarning): + PPO("MlpPolicy", "Pendulum-v0", n_steps=6, batch_size=8)