DLR-RM / DLR-RM/stable-baselines3

[Feature Request] Add logger.close to StopTrainingOnMaxEpisodes

Open
#1,049 2 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.