Skip to content

๐Ÿค— Model Zoo

Open In Colab

CleanRL now has ๐Ÿงช experimental support for saving and loading models from ๐Ÿค— HuggingFace's Model Hub. We are rolling out this feature in phases, and currently only support saving and loading models from the following algorithm varaints:

Algorithm Variants Implemented
โœ… Deep Q-Learning (DQN) dqn.py, docs
dqn_atari.py, docs
dqn_jax.py, docs
dqn_atari_jax.py, docs
โœ… Categorical DQN (C51) c51.py, docs
c51_atari.py, docs
c51_jax.py, docs
c51_atari_jax.py, docs

Load models from the Model Hub

We have a simple utility enjoy.py to load models from the hub and run them in an environment. We currently support the following commands:

poetry install -E dqn
poetry run python -m cleanrl_utils.enjoy --exp-name dqn --env-id CartPole-v1
poetry install -E dqn_jax
poetry run python -m cleanrl_utils.enjoy --exp-name dqn_jax --env-id CartPole-v1

poetry install -E dqn_atari
poetry run python -m cleanrl_utils.enjoy --exp-name dqn_atari --env-id BreakoutNoFrameskip-v4
poetry install -E dqn_atari_jax
poetry run python -m cleanrl_utils.enjoy --exp-name dqn_atari_jax --env-id BreakoutNoFrameskip-v4

To see a list of supported models, please visit ๐Ÿค— https://huggingface.co/cleanrl.

What happens under the hood?

The cleanrl_utils.enjoy is a simple wrapper to load the models from the hub and run them in an environment. A minimal version of the script can be found at cleanrl_utils/evals/dqn_eval.py, which may give you a more fine-grained control and access to the model.

Save model to Model Hub

In the supported algorithm variants, you can run the script with the --save-model flag, which saves a model to the runs folder, and the --upload-model flag, which upload the model to huggingface under your default entity (username). Optionally, you may override the default entity with --hf-entity flag.

poetry run python cleanrl/dqn_jax.py --env-id CartPole-v1 --save-model --upload-model # --hf-entity cleanrl
poetry run python cleanrl/dqn_atari_jax.py --env-id SeaquestNoFrameskip-v4  --save-model --upload-model # --hf-entity cleanrl