Upload train_and_push.py

train_and_push.py

@@ -0,0 +1,91 @@
+import argparse
+
+# Virtual display
+from pyvirtualdisplay import Display
+virtual_display = Display(visible=0, size=(1400, 900))
+virtual_display.start()
+
+# Import package
+import gym
+from huggingface_sb3 import load_from_hub, package_to_hub, push_to_hub
+from huggingface_hub import notebook_login # To log to our Hugging Face account to be able to upload models to the Hub.
+from stable_baselines3 import PPO
+from stable_baselines3.common.evaluation import evaluate_policy
+from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.vec_env import DummyVecEnv
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--model_name', dest='model_name',
+                        default="ppo-LunarLander-v2", type=str, help='model name') 
+parser.add_argument('--total_timesteps', dest='total_timesteps',
+                        default=1000000, type=int, help='total timesteps') 
+parser.add_argument('--n_envs', dest='n_envs',
+                        default=16, type=int, help='n_envs') 
+parser.add_argument('--repo_id', dest='repo_id',
+                        default="thien1892/LunarLander-v2-ppo", type=str, help='repo_id') 
+parser.add_argument('--commit_message', dest='commit_message',
+                        default="Upload PPO LunarLander-v2 trained agent", type=str, help='commit_message')
+parser.add_argument('--re_train', dest='re_train',
+                        default = True, type= str2bool, help='commit_message')
+parser.add_argument('--id_retrain', dest='id_retrain',
+                        default="thien1892/LunarLander-v2-ppo-5m", type=str, help='id_retrain')
+parser.add_argument('--filename_retrain', dest='filename_retrain',
+                        default="ppo-LunarLander-v2-5m.zip", type=str, help='filename_retrain')
+parser.add_argument('--learning_rate', dest='learning_rate',
+                        default=1e-4, type=float, help='learning_rate') 
+args = parser.parse_args()
+
+if __name__ == '__main__':
+    # Create the environment
+    env = make_vec_env('LunarLander-v2', n_envs= args.n_envs)
+
+    # Model
+    if not args.re_train :
+        model = PPO(
+            policy = 'MlpPolicy',
+            env = env,
+            n_steps = 1024,
+            batch_size = 64,
+            n_epochs = 4,
+            gamma = 0.999,
+            gae_lambda = 0.98,
+            ent_coef = 0.01,
+            learning_rate = args.learning_rate,
+            verbose=1)
+    else:
+        checkpoint = load_from_hub(args.id_retrain, args.filename_retrain)
+        model = PPO.load(checkpoint, reset_num_timesteps=True, print_system_info=True, env = env, learning_rate = args.learning_rate)
+
+    # Train
+    model.learn(total_timesteps = args.total_timesteps)
+
+    # Save the model
+    model.save(args.model_name)
+
+    # Evaluate model
+    eval_env = gym.make("LunarLander-v2") # create new environment
+    mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10, deterministic=True)
+    print(f"mean_reward={mean_reward:.2f} +/- {std_reward}")
+
+    # Push to HF hub
+    env_id = "LunarLander-v2"
+    eval_env = DummyVecEnv([lambda: gym.make(env_id)])
+    model_architecture = "PPO"
+
+    package_to_hub(model = model, # Our trained model
+                  model_name = args.model_name, # The name of our trained model 
+                  model_architecture = model_architecture, # The model architecture we used: in our case PPO
+                  env_id = env_id, # Name of the environment
+                  eval_env = eval_env, # Evaluation Environment
+                  repo_id = args.repo_id, # id of the model repository from the Hugging Face Hub (repo_id = {organization}/{repo_name} for instance ThomasSimonini/ppo-LunarLander-v2
+                  commit_message = args.commit_message)