import os import argparse import yaml import json from pathlib import Path from huggingface_hub import HfApi from huggingface_hub.repocard import metadata_save from typing import Tuple from mlagents_envs import logging_util from mlagents_envs.logging_util import get_logger logger = get_logger(__name__) logging_util.set_log_level(logging_util.INFO) def _generate_config(local_dir: Path, configfile_name: str) -> None: """ Generate a config.json file from configuration.yaml To do that we convert yaml to json :param local_dir: path of the local directory :param configfile_name: name of the yaml config file (by default configuration.yaml) """ # Read the YAML file and generate a config.json with open(os.path.join(local_dir, configfile_name)) as yaml_in: yaml_object = yaml.safe_load(yaml_in) with open(os.path.join(local_dir, "config.json"), "w") as json_out: json.dump(yaml_object, json_out) def _generate_metadata(model_name: str, env_id: str) -> dict: """ Define the tags for the model card :param model_name: name of the model :param env_id: name of the environment """ env_tag = "ML-Agents-" + env_id metadata = {} metadata["library_name"] = "ml-agents" metadata["tags"] = [ # type: ignore env_id, "deep-reinforcement-learning", "reinforcement-learning", env_tag, ] return metadata def _generate_model_card( local_dir: Path, configfile_name: str, repo_id: str ) -> Tuple[str, dict]: """ Generate the model card :param local_dir: local path of the directory :param configfile_name: name of the yaml config file (by default configuration.yaml) :param repo_id: id of the model repository from the Hugging Face Hub """ # Step 1: Read the config.json with open(os.path.join(local_dir, "config.json")) as f: data = json.load(f) # Get env_id env_id = list(data["behaviors"].keys())[0] # Get trainer_type model_name = data["behaviors"][env_id]["trainer_type"] # Step 2: Create the metadata metadata = _generate_metadata(model_name, env_id) # Step 3: Generate the model card model_card = f""" # **{model_name}** Agent playing **{env_id}** This is a trained model of a **{model_name}** agent playing **{env_id}** using the [Unity ML-Agents Library](https://github.com/Unity-Technologies/ml-agents). ## Usage (with ML-Agents) The Documentation: https://unity-technologies.github.io/ml-agents/ML-Agents-Toolkit-Documentation/ We wrote a complete tutorial to learn to train your first agent using ML-Agents and publish it to the Hub: - A *short tutorial* where you teach Huggy the Dog 🐶 to fetch the stick and then play with him directly in your browser: https://huggingface.co/learn/deep-rl-course/unitbonus1/introduction - A *longer tutorial* to understand how works ML-Agents: https://huggingface.co/learn/deep-rl-course/unit5/introduction ### Resume the training ```bash mlagents-learn --run-id= --resume ``` ### Watch your Agent play You can watch your agent **playing directly in your browser** 1. If the environment is part of ML-Agents official environments, go to https://huggingface.co/unity 2. Step 1: Find your model_id: {repo_id} 3. Step 2: Select your *.nn /*.onnx file 4. Click on Watch the agent play 👀 """ return model_card, metadata def _save_model_card( local_dir: Path, generated_model_card: str, metadata: dict ) -> None: """Save a model card to the directory. :param local_dir: local directory path :param generated_model_card: model card generated by _generate_model_card() method :param metadata: metadata """ readme_path = local_dir / "README.md" with readme_path.open("w", encoding="utf-8") as f: f.write(generated_model_card) # Save our metrics to Readme metadata metadata_save(readme_path, metadata) def package_to_hub( run_id: str, path_of_run_id: Path, repo_id: str, commit_message: str, configfile_name: str, ) -> None: """ This method generates the model card and upload the run_id folder with all his files into the Hub :param run_id : name of the run :param path_of_run_id: path of the run_id folder that contains the onnx model. :param repo_id: id of the model repository from the Hugging Face Hub :param commit_message: commit message :param configfile_name: name of the yaml config file (by default configuration.yaml) """ logger.info( f"This function will create a model card and upload your {run_id} " f"into HuggingFace Hub. This is a work in progress: If you encounter a bug, " f"please send open an issue" ) _, repo_name = repo_id.split("/") # Step 1: Create the repo api = HfApi() repo_url = api.create_repo( repo_id=repo_id, exist_ok=True, ) local_path = Path(path_of_run_id) # Step 2: Create a config file _generate_config(local_path, configfile_name) # Step 3: Generate and save the model card generated_model_card, metadata = _generate_model_card( local_path, configfile_name, repo_id ) _save_model_card(local_path, generated_model_card, metadata) logger.info(f"Pushing repo {run_id} to the Hugging Face Hub") # Step 4. Push everything to the Hub api.upload_folder( repo_id=repo_id, folder_path=local_path, commit_message=commit_message ) logger.info( f"Your model is pushed to the hub. You can view your model here: {repo_url}" ) def main(): parser = argparse.ArgumentParser() parser.add_argument("--run-id", help="Name of the run-id folder", type=str) parser.add_argument( "--local-dir", help="Path of the run_id folder that contains the trained model", type=str, default="./", ) parser.add_argument( "--repo-id", help="Repo id of the model repository from the Hugging Face Hub", type=str, ) parser.add_argument( "--commit-message", help="Commit message", type=str, default="Push to Hub" ) parser.add_argument( "--configfile-name", help="Name of the configuration yaml file", type=str, default="configuration.yaml", ) args = parser.parse_args() # Push model to hub package_to_hub( args.run_id, args.local_dir, args.repo_id, args.commit_message, args.configfile_name, ) # For python debugger to directly run this script if __name__ == "__main__": main()