Colab examples need to install distrax gymnax and brax?

[StoneT2000]

Tried running the colab example to look at the speed but it seems like colab doesn't come installed with distrax brax gymnax

https://colab.research.google.com/github/luchris429/purejaxrl/blob/main/examples/example_1.ipynb

ModuleNotFoundError                       Traceback (most recent call last)

[<ipython-input-1-91b777b64c88>](https://localhost:8080/#) in <cell line: 9>()
      7 from typing import Sequence, NamedTuple, Any
      8 from flax.training.train_state import TrainState
----> 9 import distrax
     10 import gymnax
     11 from gymnax.wrappers.purerl import LogWrapper, FlattenObservationWrapper

ModuleNotFoundError: No module named 'distrax'

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.

[sunnyswag]

Yes, I have also encountered a similar problem. just pip install those

[luchris429]

ahh thanks! should be resolved now!