Backport changes from PyTorch version of DQN Atari into the TF notebook

[Dmitry Nikulin]

There are a lot of things in the PyTorch notebook that make this assignment a lot easier. They should be backported to the TF version.

One thing to note is that n_lives exists in the PyTorch version, but not in TF.