ThomasSimonini/Stable-Baselines3

1import gradio as gr2import os3from moviepy.editor import *4 5def replay(option):6	path = ""7	# Get the correct model8	if (option == "LunarLander-v2 🚀👩‍🚀"):9  		path = "./LunarLander-v2.mp4"10	elif(option == "CartPole-v1 🕹️"):11  		path = "./CartPole-v1.mp4"12	elif(option == "Atari Space Invaders 👾"):13  		path = "./SpaceInvadersNoFrameskip-v4.mp4"14	15	16  	# The only turnaround I found (Base64 video pb)17	videoclip = VideoFileClip(path)18	videoclip.write_videofile("new_filename.mp4")19	return 'new_filename.mp4'20 21iface = gr.Interface(22    replay,23    [24        gr.inputs.Dropdown(["Atari Space Invaders 👾", "CartPole-v1 🕹️", "LunarLander-v2 🚀👩‍🚀"]),25    ],26    "video",27     title = 'Stable Baselines 3 with 🤗',28            description = '',29             article = 30                        '''<div>31                        	<p style="text-align: center">This version of the RL library allows you to load models directly from the Hugging Face Hub</p>32                            <p style="text-align: center"> Select the trained agent you want to watch perform.33                            These models are from <a href="https://github.com/araffin/rl-baselines-zoo">Stable Baseline Zoo</a></p>34                            <p>35                            There are currently 3 models:36                            <ul>37                            	<li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4">PPO SpaceInvadersNoFrameskip-v4</a></li>38                            	<li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-LunarLander-v2">PPO LunarLander-v2</a></li>39                            	<li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-CartPole-v1">PPO CartPole-v1</a></li>40                            </ul>41                        </div>'''42            )43   44 45iface.launch()46 47"""48TODO: Next version with live video generation49import gradio as gr50import os51 52from Recorder import Recorder53 54from stable_baselines3 import PPO55 56 57#The Agent plays and we generate the video58def replay(option):59  video_path = ""60  # Get the correct model61  if (option == "LunarLander-v2 🚀👩‍🚀"):62    env_name = "Lunar Lander v2"63    agent_name = "PPO"64    print("TEST")65    hf_model_filename = "LunarLander-v2"66    hf_model_id = "ThomasSimonini/stable-baselines3-ppo-LunarLander-v2"67    video_path = replay_gym(hf_model_filename, hf_model_id)68  elif(option == "CartPole-v1 🕹️"):69      hf_model_filename = "CartPole-v1"70      hf_model_id = "ThomasSimonini/stable-baselines3-ppo-CartPole-v1"71      video_path = replay_gym(hf_model_filename, hf_model_id)72  elif(option == "Atari Space Invaders 👾"):73    hf_model_filename = "SpaceInvadersNoFrameskip-v4"74    hf_model_id = "ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4"75    video_path = replay_atari(hf_model_filename, hf_model_id)76      #video_path = "./SpaceInvadersNoFrameskip-v4.mp4"77 78  return video_path79 80 81def replay_gym(hf_model_filename, hf_model_id):82  import gym83  from stable_baselines3.common.evaluation import evaluate_policy84 85 86  model = PPO.load_from_huggingface(hf_model_id,hf_model_filename)87 88  eval_env = gym.make(hf_model_filename)89 90  directory = './video'91  env = Recorder(eval_env, directory)92 93  obs = env.reset()94  done = False95  while not done:96      action, _state = model.predict(obs)97      obs, reward, done, info = env.step(action)98  clip = env.play()99  return clip100 101 102def replay_atari(hf_model_filename, hf_model_id):103  os.system("python -m atari_py.import_roms \"content/atari_roms\"")104  import gym105  from stable_baselines3.common.env_util import make_atari_env106  from stable_baselines3.common.vec_env import VecFrameStack107 108  from stable_baselines3.common.evaluation import evaluate_policy109 110  model = PPO.load_from_huggingface(hf_model_id, hf_model_filename)111 112 113  eval_env = make_atari_env(hf_model_filename, n_envs=1, seed=0)114  eval_env = VecFrameStack(eval_env, n_stack=4)115 116  model = PPO.load_from_huggingface(hf_model_id, hf_model_filename)117 118  import gym119  directory = './video'120  env = Recorder(eval_env, directory)121 122  obs = env.reset()123  done = False124  while not done:125      action, _state = model.predict(obs)126      obs, reward, done, info = env.step(action)127  clip = env.play()128  return clip129 130 131 132iface = gr.Interface(133    replay,134    [135        gr.inputs.Dropdown(["Atari Space Invaders 👾", "CartPole-v1 🕹️", "LunarLander-v2 🚀👩‍🚀"]),136    ],137    "video",138     title = 'Stable Baselines 3 with 🤗',139            description = '',140             article = 141                        '''<div>142                          <p style="text-align: center">This version of the RL library allows you to load models directly from the Hugging Face Hub</p>143                            <p style="text-align: center"> Select the trained agent you want to watch perform. We record your agent playing.144                            <p style="text-align: center"> Don't forget to <b>click on clear between each record.</b> </p>145                            These models are from <a href="https://github.com/araffin/rl-baselines-zoo">Stable Baseline Zoo</a></p>146                            <p>147                            There are currently 3 models:148                            <ul>149                              <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4">PPO SpaceInvadersNoFrameskip-v4</a></li>150                              <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-LunarLander-v2">PPO LunarLander-v2</a></li>151                              <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-CartPole-v1">PPO CartPole-v1</a></li>152                            </ul>153                        </div>'''154            )155   156 157iface.launch()158"""159 160