Fixes update_locals() in collect_rollouts() of OnPolicyAlgorithm (

#1660) * calls update_locals() before on_rollout_end() * update changelog
DLR-RM · Aug 30, 2023 · 84163b4 · 84163b4
1 parent c99d65c
commit 84163b4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git 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
@@ -221,6 +221,8 @@ def collect_rollouts(
 
         rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)
 
+        callback.update_locals(locals())
+
         callback.on_rollout_end()
 
         return True