[Feature Request] Add logger.close to StopTrainingOnMaxEpisodes

[vcadillog]

🚀 Feature

Add logger.close() to StopTrainingOnMaxEpisodes class.

Motivation

While I was working around with this amazing tool had some problems training models with large timesteps so the StopTrainingOnMaxEpisodes callback worked really well, but it didn't closed the tensorboard log properly, so it was empty after training.

Pitch

Add self.logger.close() at the end of the stop condition in _on_step function is found in StopTrainingOnMaxEpisodes module or a variable able to control it.

Alternatives

def _on_step(self) -> bool:
    # Check that the `dones` local variable is defined
    assert "dones" in self.locals, "`dones` variable is not defined, please check your code next to `callback.on_step()`"
    self.n_episodes += np.sum(self.locals["dones"]).item()

    continue_training = self.n_episodes < self._total_max_episodes

    if self.verbose > 0 and not continue_training:
        mean_episodes_per_env = self.n_episodes / self.training_env.num_envs
        mean_ep_str = (
            f"with an average of {mean_episodes_per_env:.2f} episodes per env" if self.training_env.num_envs > 1 else ""
        )

        print(
            f"Stopping training with a total of {self.num_timesteps} steps because the "
            f"{self.locals.get('tb_log_name')} model reached max_episodes={self.max_episodes}, "
            f"by playing for {self.n_episodes} episodes "
            f"{mean_ep_str}"
        )

####### This line added ########
self.logger.close()
###########################

    return continue_training

### Checklist

• I have checked that there is no similar issue in the repo (required)

[araffin]

Hello,
logger.close seems be call nowhere indeed... but that's weird we didn't get the problem earlier (I never experienced that, even when using early stopping), could you share your configuration too? (using sb3.get_system_info())
I would welcome a PR that solves this issue =) (we need to check other callbacks do not needed that neither).

[vcadillog]

Hi, this is the result of system_info

OS: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic #1 SMP Sun Apr 24 10:03:06 PDT 2022
Python: 3.7.13
Stable-Baselines3: 1.6.0
PyTorch: 1.12.1+cu113
GPU Enabled: True
Numpy: 1.21.6
Gym: 0.21.0

({'OS': 'Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic #1 SMP Sun Apr 24 10:03:06 PDT 2022',
'Python': '3.7.13',
'Stable-Baselines3': '1.6.0',
'PyTorch': '1.12.1+cu113',
'GPU Enabled': 'True',
'Numpy': '1.21.6',
'Gym': '0.21.0'},
'OS: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic #1 SMP Sun Apr 24 10:03:06 PDT 2022\nPython: 3.7.13\nStable-Baselines3: 1.6.0\nPyTorch: 1.12.1+cu113\nGPU Enabled: True\nNumpy: 1.21.6\nGym: 0.21.0\n')