ModuleNotFoundError: No module named 'baselines'

Hello.
When running [jax_agent_visualizer.ipynb](https://colab.research.google.com/github/google/dopamine/blob/master/dopamine/colab/jax_agent_visualizer.ipynb) notebook in colab this error is raised on `from dopamine.discrete_domains import atari_lib`
```python
ModuleNotFoundError Traceback (most recent call last)
[](https://localhost:8080/#) in ()
1 # @title Imports
----> 2 from dopamine.discrete_domains import atari_lib
3 from dopamine.jax.agents.dqn import dqn_agent
4 from dopamine.jax.agents.implicit_quantile import implicit_quantile_agent
5 from dopamine.jax.agents.quantile import quantile_agent

[/usr/local/lib/python3.11/dist-packages/dopamine/discrete_domains/atari_lib.py](https://localhost:8080/#) in
43
44 import ale_py
---> 45 from baselines.common import atari_wrappers
46 import cv2
47 import gin

ModuleNotFoundError: No module named 'baselines'
```