🤗 Model Zoo
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