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

[kvrban]

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' (<class 'dopamine.discrete_domains.run_experiment.Runner'>)

[arpit15]

I am seeing the same issue. @kvrban were you able to find a way around it?