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 <antonin.raffin@ensta.org>

* Change to minibatch language

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

* Cleaning up language describing rollout buffer requirements

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

* 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 <antonin.raffin@ensta.org>

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

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

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)