🤗 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
✅ Deep Deterministic Policy Gradient (DDPG)	`ddpg_continuous_action.py`, docs
	`ddpg_continuous_action_jax.py`, docs
✅ Twin Delayed Deep Deterministic Policy Gradient (TD3)	`td3_continuous_action.py`, docs
	`td3_continuous_action_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