[Bug] Callbacks can't access local variables

Description From the discussion here #756, callbacks in DQN should have access to local variables, e.g. done through BaseCallback.locals / self.locals. However, this isn’t the case.

Code example

from stable_baselines.common.callbacks import BaseCallback
import gym
import stable_baselines as sb
env = gym.make("MountainCar-v0")
class TestCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
    def _on_step(self):
        assert ('done' in self.locals)

agent = sb.DQN('MlpPolicy', env, verbose=1)
agent.learn(100, callback=TestCallback())

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/stelios/anaconda3/envs/thesis37/lib/python3.7/site-packages/stable_baselines/deepq/dqn.py", line 224, in learn
    if callback.on_step() is False:
  File "/Users/stelios/anaconda3/envs/thesis37/lib/python3.7/site-packages/stable_baselines/common/callbacks.py", line 89, in on_step
    return self._on_step()
  File "<stdin>", line 5, in _on_step
AssertionError

Workaround

In DQN.learn, line 219, the following results in assertion error.

new_obs, rew, done, info = self.env.step(env_action)

self.num_timesteps += 1
assert 'done' in callback.locals  # fails
# Stop training if return value is False
if callback.on_step() is False:
    break

If we call locals(), then the assertion succeeds.

new_obs, rew, done, info = self.env.step(env_action)

self.num_timesteps += 1
locals()
assert 'done' in callback.locals  # succeeds
# Stop training if return value is False
if callback.on_step() is False:
    break

[Miffyli]

I was just about to test this, but you got ahead of me 👍.

I think it would be nice to have same treatment for other algorithms as well, where locals has access to things like done and reward (can be quite relevant for on_step callback). This could include all env.step related variables (obs, reward, done, info and chosen action). I am not 100% sure if this can be done with all algorithms, but I can see use-cases for this kind of information.

@araffin Quick thoughts on this, with you being the main author of this code?

[PartiallyTyped]

I have solved this by initializing the variables before the callback.on_training_start call.

done = False
callback.on_training_start(locals(), globals())

Should I make a PR with this change? Which other variables should I initialise?