DLR-RM / DLR-RM/stable-baselines3
[Feature Request] Add logger.close to StopTrainingOnMaxEpisodes
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### 🚀 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
- [x] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo (**required**)
Contributor guide
Research direction
Search for the StopTrainingOnMaxEpisodes class and inspect its _on_step method, especially the stop condition and logger lifecycle. Confirm how the tensorboard logger is created and closed elsewhere before deciding on the change. Done means training stopped by max episodes no longer leaves the tensorboard log empty.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100