🤗 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:
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?
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
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