Requirements

• jax (https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-via-pip-easier)
• mujoco (https://mujoco.readthedocs.io/en/stable/python.html#installation)
• mujoco-mjx (https://mujoco.readthedocs.io/en/stable/mjx.html#installation)
• pillow (https://pypi.org/project/pillow/)

Quick start

python3 -m venv .venv
source .venv/bin/activate
python main.py  # runs the demo