[feature request] tqdm for training

It would be convenient to have tqdm wrapped around the main training loop — especially if it was from tqdm.auto import tqdm, as it autodetects whether it runs in a Jupyter notebook or in a terminal and uses either stderr or a widget.

[dniku]

For those finding this issue via search, here is a version that works with Stable Baselines 3:

from tqdm.notebook import tqdm
from stable_baselines3.common.callbacks import BaseCallback

class TqdmCallback(BaseCallback):
    def __init__(self):
        super().__init__()
        self.progress_bar = None
    
    def _on_training_start(self):
        self.progress_bar = tqdm(total=self.locals['total_timesteps'])
    
    def _on_step(self):
        self.progress_bar.update(1)
        return True

    def _on_training_end(self):
        self.progress_bar.close()
        self.progress_bar = None

[dniku]

It is a good idea to implement that using callback, thank you.

I’ve experimented a little and this seems to work with PPO2:

from tqdm.auto import tqdm

class TqdmCallback(object):
    def __init__(self):
        self.pbar = None
    
    def __call__(self, _locals, _globals):
        if self.pbar is None:
            self.pbar = tqdm(total=_locals['nupdates'])
            
        self.pbar.update(1)
        
        if _locals['update'] == _locals['nupdates']:
            self.pbar.close()
            self.pbar = None

        return True

One thought is that perhaps stable-baselines should accept multiple callbacks (Keras style), but that would be a topic for another issue.