From 84163b468c99538f2c98a3ebcc6124974ec631fd Mon Sep 17 00:00:00 2001
From: PatrickHelm <91552435+PatrickHelm@users.noreply.github.com>
Date: Wed, 30 Aug 2023 17:02:41 +0200
Subject: [PATCH] Fixes `update_locals()` in  `collect_rollouts()` of
 `OnPolicyAlgorithm` (#1660)

* calls update_locals() before on_rollout_end()

* update changelog
---
 docs/misc/changelog.rst                         | 1 +
 stable_baselines3/common/on_policy_algorithm.py | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
index 84ec850..ae9ffcd 100644
--- a/docs/misc/changelog.rst
+++ b/docs/misc/changelog.rst
@@ -22,6 +22,7 @@ Bug Fixes:
 ^^^^^^^^^^
 - Moves VectorizedActionNoise into ``_setup_learn()`` in OffPolicyAlgorithm (@PatrickHelm)
 - Prevents out of bound error on Windows if no seed is passed (@PatrickHelm)
+- Calls ``callback.update_locals()`` before ``callback.on_rollout_end()`` in OnPolicyAlgorithm (@PatrickHelm)
 
 Deprecations:
 ^^^^^^^^^^^^^
diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py
index 87e1929..f038df6 100644
--- a/stable_baselines3/common/on_policy_algorithm.py
+++ b/stable_baselines3/common/on_policy_algorithm.py
@@ -221,6 +221,8 @@ class OnPolicyAlgorithm(BaseAlgorithm):
 
         rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)
 
+        callback.update_locals(locals())
+
         callback.on_rollout_end()
 
         return True