Integrating Hugging Face Hub

[simoninithomas]

Hey team 👋,

Description

We work on the integration of Hugging Face Hub with rl-baselines-zoo.

The Hugging Face Hub works as a central place where anyone can share and explore saved models.

The idea is to allow people to share and save their models on the hub.

The PR is a draft for now because there are some points to discuss and code to change given the decisions + some bugs I need to resolve. So not all the elements of the checklist are done.

What does our implementation for now:

• load_from_hf.py: download the model from Hugging Face Hub and place it in the correct folder (rl-saved-agents by default).
• save_to_hub.py: upload the model from rl-baselines-zoo to Hugging Face Hub.

Case 1: I load a sb3 saved model to test it

# I want to get the sb3-zoo ppo-CartPole v-1 (since I don't add --organization we consider that is --organization="sb3" by default)
# This download the model and put it in rl-saved-agents/ppo/CartPole-v1_1
!python utils/load_from_hf.py --algo ppo --env CartPole-v1

# I try the environment
!python enjoy.py --algo ppo --env CartPole-v1 --no-render

Case 2: I train a model and push to hub (some bugs for now in terms of git that I need to resolve)

!python train.py --algo ppo --env CartPole-v1 -n 100 --f logs

# Will change to save_to_hf
!python utils/save_to_hub.py --organization "ThomasSimonini" --repo-name "ppo-CartPole-v1" --model-dir "logs/ppo/CartPole-v1_1" --commit-message "Add Cartpole-v1_1 to my repo")

Question 1: —load_from_hf and —save_to_hf args

Instead of having this procedure of calling save_to_hf.py and load_from_hf as described above we could add a parameter —load_to_hf and —save_to_hf with a Subcommand register as we done with AllenNLP.

https://github.com/allenai/allennlp/blob/82b1f4f80899c7eab31fa1ee0949be2eb12fc184/allennlp/commands/push_to_hf.py#L16

@Subcommand.register("push-to-hf")
class PushToHf(Subcommand):
    def add_subparser(self, parser: argparse._SubParsersAction) -> argparse.ArgumentParser:
        description = """Push a model to the Hugging Face Hub.
        Pushing your models to the Hugging Face Hub ([hf.co](https://hf.co/))
        allows you to share your models with others. On top of that, you can try
        the models directly in the browser with the available widgets.
        Before running this command, login to Hugging Face with `huggingface-cli login`.
        You can specify either a `serialization_dir` or an `archive_path`, but using the
        first option is recommended since the `serialization_dir` contains more useful
        information such as metrics and TensorBoard traces.
        """
        subparser = parser.add_parser(self.name, description=description, help=description)
        subparser.set_defaults(func=push)

        subparser.add_argument(
            "-n",
            "--repo-name",
            required=True,
            type=str,
            default="Name of the repository",
            help="Name of the repository",
        )
				...
			

This way a user will just need to do something like this:

!python train.py --algo ppo --env CartPole-v1 -n 100 --f logs --save-to-hub --organization "ThomasSimonini" --repo-name "PPO-SpaceInvaders"

Question 2: Repos names

For the repos from rl-trained-agents we wanted to know what is your preferences for the name of the repos in the Hugging Face Hub ?

# 1 
sb3/{algo}-{env_name}
# => sb3/ppo-CartPole-v1

# 2
sb3/{ALGO}-{env_name}
# => sb3/PPO-CartPole-v1

# 3
sb3/{env_name}-{algo}

Motivation and Context

Allowing people to be able to easily share their saved models, using the Hugging Face Hub that allows you to host and share for free your models.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist:

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code using make format (required)
• I have checked the codestyle using make check-codestyle and make lint (required)
• I have ensured make pytest and make type both pass. (required)

Note: we are using a maximum length of 127 characters per line

[araffin]

Hello,
thanks for the PR.
As a quick remark, please reformat the code using make format and use typing everywhere ;)

Case 1: I load a sb3 saved model to test it

I would even add automatic download to enjoy if the model is not present (with user prompt that can be bypassed for testing)

Case 2: I train a model and push to hub (some bugs for now in terms of git that I need to resolve)

Looks good but I would keep the api of the zoo if possible or at least allow it. Also, do you save the entire folder or only the last model? (a folder can contain any checkpoints and the best model too)

we could add a parameter —load_to_hf and —save_to_hf with a Subcommand register

not sure if it makes sense for training (you usually know only late if you want to save the model to the hub or not depending on the results) but for enjoy yes it totally makes sense ;)

For the repos from rl-trained-agents we wanted to know what is your preferences for the name of the repos in the Hugging Face Hub ?

I would go for option one as it is consistent with how you train the models, what do you think?

[araffin]

please add this to the requirements

also missing: updating the README

[simoninithomas]

Hey Antonin, thanks for your feedback,

I would even add automatic download to enjoy if the model is not present (with user prompt that can be bypassed for testing)

I suppose you mean for rl-trained-agents repository, this is something possible.

For instance if I type:
!python enjoy.py --algo ppo --env CartPole-v1 --no-render

And if the model is not in the user rl-trained-agents folder, I'll retrieve it (repo_id = sb3/ppo-CartPole-v1) from the hub.

But for custom trained agents, that would imply arguments (repo_id) so I'm not sure it's something you're agree on. Or we can have a user prompt as you mentioned (not the final one, this is quite dirty it's just for illustration purposes)

print(f"No model found for ")#{algo} on {env_id}, path: {model_path}, do you want to retrieve it from the Hugging Face Hub?")
answer = None 
while answer not in ("yes", "no"): 
  answer = input("Enter yes or no: ") 
  if answer == "yes":
    # Check if custom input
    answer = input("Is it a custom trained agent (not one from rl-trained-agents repository)?: ")
    if ("yes"):
      repo_id = input("What's the Repo id?: ")
      # Load the saved model 
      print("Load the saved model")
    elif answer == "no": 
      repo_id = input("What's the Repo id?: ")
      # Load the saved model
  elif answer == "no": 
    print("No")
  else: 
    print("Please enter yes or no.") 

The problem is where to bypass it? because we have 2 checks:

• enjoy.py Line 83: but it only checks if the folder exists
• enjoy.py Line 114: we could add here but again it would imply to run 85 -> 112 (so put this in a function?)

---

Looks good but I would keep the api of the zoo if possible or at least allow it. ".

What do you mean? Not having a save_to_hub.py and directly save on train.py?

---

Also, do you save the entire folder or only the last model? (a folder can contain any checkpoints and the best model too)

That's the question, what do you think is the best? In terms of usage

• Have a repo with multiple folders

Repo_id (does not exist it's an example): ThomasSimonini/SpaceInvadersNoFrameskip-v4

├── ppo-SpaceInvadersNoFrameskip-v4_1/
│   ├── SpaceInvadersNoFrameskip-v4.zip
│   ├── monitor.csv
│   ├── evaluations.npz        
│   ├── SpaceInvadersNoFrameskip-v4/
│     ├── args.yml
│     ├── config.yml
│     ├── vecnormalize.pkl
├── ppo-SpaceInvadersNoFrameskip-v4_2/
├── ppo-SpaceInvadersNoFrameskip-v4_3/
├── README.md

• Have a repo with one folder

├── ppo-SpaceInvadersNoFrameskip-v4_1/
├── README.md

• Have a repo with "no" folder

├── SpaceInvadersNoFrameskip-v4.zip
├── monitor.csv
├── evaluations.npz        
├── SpaceInvadersNoFrameskip-v4/
│     ├── args.yml
│     ├── config.yml
│     ├── vecnormalize.pkl

=> In this case that implies that I need to generate a folder when download it with exp_id = 1. Since it's the way enjoy.py works.

[araffin]

And if the model is not in the user rl-trained-agents folder, I'll retrieve it (repo_id = sb3/ppo-CartPole-v1) from the hub.

yes

But for custom trained agents, that would imply arguments (repo_id) so I'm not sure it's something you're agree on.

I would do only for "official" pre-trained agents, there is your download script for the download of custom models.

The problem is where to bypass it? because we have 2 checks:

I think
https://github.com/DLR-RM/rl-baselines3-zoo/blob/master/enjoy.py#L83

would make more sense (and only if folder == "rl-trained-agents").

What do you mean? Not having a save_to_hub.py and directly save on train.py?

I meant allowing to use --algo, --env, --folder and --exp-id instead of the path to a folder, to be consistent with the rest of the scripts (but we can allow both if you think passing the path to the model folder is more convenient).

In this case that implies that I need to generate a folder when download it with exp_id = 1

it should create it but with exp_id = get_latest_run_id() + 1 (as done in the train).
I think no folder makes more sense as the repo in the hub is the folder, no?

That's the question, what do you think is the best? In terms of usage

If there is no restrictions on space for huggingface hub, it should save everything but we can probably add an option to filter what is saved (for instance do not save checkpoints or do not save best model).

[simoninithomas]

Hey Antonin 👋 ,

So for the first step in the integration of RL-Zoo, we can automatically retrieve an "official" model from the hub, as you mentioned the idea is when you run:

python  enjoy.py --algo ppo --env CartPole-v1

And this model is not in trained-agents/ folder, it automatically retrieve it from the hub:

Line 113:

    if not found:
        local_dir = f"rl-trained-agents/{algo}/{env_id}_1"
        clone_from = f"TestSB3/{algo}-{env_id}"
        repo = Repository(local_dir, clone_from) # This function clone a Hugging Face Repo and place it into the local_dir path

About the model card, here's a simple version, wdyt (there will be on the right a video widget of the agent performing):
https://huggingface.co/TestSB3/ppo-CartPole-v1