Callbacks have access to locals (#115)

* callbacks have access to locals

* changeloc

* doc

* callbacks have access to locals

* changeloc

* doc

* Added update function for child callbacks

* Pre-Release 0.8.0 (#134)

* Fix double reset and improve typing coverage (#136)

* Fix double reset and improve typing coverage

* Revert minor edit

* Add doc about types

* Update child callbacks

* cleaned imports

* format

* import order

* Simplify tests and add comments

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

docs/misc/changelog.rst

@@ -53,6 +53,7 @@ New Features:
 - Refactored opening paths for saving and loading to use strings, pathlib or io.BufferedIOBase (@PartiallyTyped)
 - Added ``DDPG`` algorithm as a special case of ``TD3``.
 - Introduced ``BaseModel`` abstract parent for ``BasePolicy``, which critics inherit from.
+- Callbacks have access to rollout collection locals as in SB2. (@PartiallyTyped)
 
 Bug Fixes:
 ^^^^^^^^^^

stable_baselines3/common/callbacks.py

@@ -33,8 +33,8 @@ class BaseCallback(ABC):
         # n_envs * n times env.step() was called
         self.num_timesteps = 0  # type: int
         self.verbose = verbose
-        self.locals = None  # type: Optional[Dict[str, Any]]
-        self.globals = None  # type: Optional[Dict[str, Any]]
+        self.locals: Dict[str, Any] = {}
+        self.globals: Dict[str, Any] = {}
         self.logger = None
         # Sometimes, for event callback, it is useful
         # to have access to the parent object
@@ -103,6 +103,23 @@ class BaseCallback(ABC):
     def _on_rollout_end(self) -> None:
         pass
 
+    def update_locals(self, locals_: Dict[str, Any]) -> None:
+        """
+        Update the references to the local variables.
+
+        :param locals_: (Dict[str, Any]) the local variables during rollout collection
+        """
+        self.locals.update(locals_)
+        self.update_child_locals(locals_)
+
+    def update_child_locals(self, locals_: Dict[str, Any]) -> None:
+        """
+        Update the references to the local variables on sub callbacks.
+
+        :param locals_: (Dict[str, Any]) the local variables during rollout collection
+        """
+        pass
+
 
 class EventCallback(BaseCallback):
     """
@@ -137,6 +154,15 @@ class EventCallback(BaseCallback):
     def _on_step(self) -> bool:
         return True
 
+    def update_child_locals(self, locals_: Dict[str, Any]) -> None:
+        """
+        Update the references to the local variables.
+
+        :param locals_: (Dict[str, Any]) the local variables during rollout collection
+        """
+        if self.callback is not None:
+            self.callback.update_locals(locals_)
+
 
 class CallbackList(BaseCallback):
     """
@@ -178,6 +204,15 @@ class CallbackList(BaseCallback):
         for callback in self.callbacks:
             callback.on_training_end()
 
+    def update_child_locals(self, locals_: Dict[str, Any]) -> None:
+        """
+        Update the references to the local variables.
+
+        :param locals_: (Dict[str, Any]) the local variables during rollout collection
+        """
+        for callback in self.callbacks:
+            callback.update_locals(locals_)
+
 
 class CheckpointCallback(BaseCallback):
     """
@@ -343,6 +378,15 @@ class EvalCallback(EventCallback):
 
         return True
 
+    def update_child_locals(self, locals_: Dict[str, Any]) -> None:
+        """
+        Update the references to the local variables.
+
+        :param locals_: (Dict[str, Any]) the local variables during rollout collection
+        """
+        if self.callback:
+            self.callback.update_locals(locals_)
+
 
 class StopTrainingOnRewardThreshold(BaseCallback):
     """

stable_baselines3/common/off_policy_algorithm.py

@@ -401,6 +401,8 @@ class OffPolicyAlgorithm(BaseAlgorithm):
                 # Rescale and perform action
                 new_obs, reward, done, infos = env.step(action)
 
+                # Give access to local variables
+                callback.update_locals(locals())
                 # Only stop training if return value is False, not when it is None.
                 if callback.on_step() is False:
                     return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False)

stable_baselines3/common/on_policy_algorithm.py

@@ -162,6 +162,8 @@ class OnPolicyAlgorithm(BaseAlgorithm):
 
             new_obs, rewards, dones, infos = env.step(clipped_actions)
 
+            # Give access to local variables
+            callback.update_locals(locals())
             if callback.on_step() is False:
                 return False
 

tests/test_callbacks.py

@@ -33,15 +33,22 @@ def test_callbacks(tmp_path, model_class):
     eval_callback = EvalCallback(
         eval_env, callback_on_new_best=callback_on_best, best_model_save_path=log_folder, log_path=log_folder, eval_freq=100
     )
-
     # Equivalent to the `checkpoint_callback`
     # but here in an event-driven manner
     checkpoint_on_event = CheckpointCallback(save_freq=1, save_path=log_folder, name_prefix="event")
+
     event_callback = EveryNTimesteps(n_steps=500, callback=checkpoint_on_event)
 
     callback = CallbackList([checkpoint_callback, eval_callback, event_callback])
-
     model.learn(500, callback=callback)
+
+    # Check access to local variables
+    assert model.env.observation_space.contains(callback.locals["new_obs"][0])
+    # Check that the child callback was called
+    assert checkpoint_callback.locals["new_obs"] is callback.locals["new_obs"]
+    assert event_callback.locals["new_obs"] is callback.locals["new_obs"]
+    assert checkpoint_on_event.locals["new_obs"] is callback.locals["new_obs"]
+
     model.learn(500, callback=None)
     # Transform callback into a callback list automatically
     model.learn(500, callback=[checkpoint_callback, eval_callback])