Fix n_updates logging in MaskablePPO and RecurrentPPO when using early stopping (#325)

* Initial plan

* Fix MaskablePPO inaccurate n_updates counting when target_kl early exits

Co-authored-by: araffin <1973948+araffin@users.noreply.github.com>

* Fix for recurrent ppo too

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: araffin <1973948+araffin@users.noreply.github.com>
Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

Author: Copilot

docs/misc/changelog.md

@@ -2,7 +2,7 @@
 
 # Changelog
 
-## Release 2.8.0a1 (WIP)
+## Release 2.8.0a4 (WIP)
 
 ### Breaking Changes:
 
@@ -16,7 +16,8 @@
 
 ### Bug Fixes:
 
-- Fix RecurrentPPO and MaskablePPO forward and predict do not reshape action before clip it (@immortal-boy)
+- Fix `MaskablePPO` and `RecurrentPPO` inaccurate `n_updates` counting when `target_kl` early exits the training loop
+- Fix `RecurrentPPO` and `MaskablePPO` forward and predict do not reshape action before clip it (@immortal-boy)
 - Do not call `forward()` method directly in `RecurrentPPO` (@immortal-boy)
 
 ### Deprecations:

sb3_contrib/ppo_mask/ppo_mask.py

@@ -407,10 +407,9 @@ def train(self) -> None:
                 th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
                 self.policy.optimizer.step()
 
+            self._n_updates += 1
             if not continue_training:
                 break
-
-        self._n_updates += self.n_epochs
         explained_var = explained_variance(self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten())
 
         # Logs

sb3_contrib/ppo_recurrent/ppo_recurrent.py

@@ -416,10 +416,10 @@ def train(self) -> None:
                 th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
                 self.policy.optimizer.step()
 
+            self._n_updates += 1
             if not continue_training:
                 break
 
-        self._n_updates += self.n_epochs
         explained_var = explained_variance(self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten())
 
         # Logs

sb3_contrib/version.txt

@@ -1 +1 @@
-2.8.0a1
+2.8.0a4

setup.py

@@ -67,7 +67,7 @@
     packages=[package for package in find_packages() if package.startswith("sb3_contrib")],
     package_data={"sb3_contrib": ["py.typed", "version.txt"]},
     install_requires=[
-        "stable_baselines3>=2.8.0a0,<3.0",
+        "stable_baselines3>=2.8.0a4,<3.0",
     ],
     description="Contrib package of Stable Baselines3, experimental code.",
     author="Antonin Raffin",