diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 35b9ec9..af3b66d 100644 --- a/docs/misc/changelog.rst +++ b/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: ^^^^^^^^^^ diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index c5e53e5..e0d326c 100644 --- a/stable_baselines3/common/callbacks.py +++ b/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): """ diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 167f34d..45c735d 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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) diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 5189991..9230dd7 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/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 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index b1a3339..55c7f61 100644 --- a/tests/test_callbacks.py +++ b/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])