Colab demo for jax agents not working. Error: "no attribute 'online_network'"

Hi, i tried to study the colab sample linked from: https://github.com/google/dopamine/tree/master/dopamine/jax

This colab sample:
https://colab.research.google.com/github/google/dopamine/blob/master/dopamine/colab/jax_agent_visualizer.ipynb

First issue is easy to fix, you need additional:

```
!pip install gym[accept-rom-license]
!pip install gym[atari]
```
to get the atari environments working.

But this issue then exists:

```

AttributeError: 'MyDQNAgent' object has no attribute 'online_network'
In call to configurable 'Runner' ()
```