Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| from moviepy.editor import * | |
| def replay(option): | |
| path = "" | |
| # Get the correct model | |
| if (option == "LunarLander-v2 🚀👩🚀"): | |
| path = "./LunarLander-v2.mp4" | |
| elif(option == "CartPole-v1 🕹️"): | |
| path = "./CartPole-v1.mp4" | |
| elif(option == "Atari Space Invaders 👾"): | |
| path = "./SpaceInvadersNoFrameskip-v4.mp4" | |
| # The only turnaround I found (Base64 video pb) | |
| videoclip = VideoFileClip(path) | |
| videoclip.write_videofile("new_filename.mp4") | |
| return 'new_filename.mp4' | |
| iface = gr.Interface( | |
| replay, | |
| [ | |
| gr.inputs.Dropdown(["Atari Space Invaders 👾", "CartPole-v1 🕹️", "LunarLander-v2 🚀👩🚀"]), | |
| ], | |
| "video", | |
| title = 'Stable Baselines 3 with 🤗', | |
| description = '', | |
| article = | |
| '''<div> | |
| <p style="text-align: center">This version of the RL library allows you to load models directly from the Hugging Face Hub</p> | |
| <p style="text-align: center"> Select the trained agent you want to watch perform. | |
| These models are from <a href="https://github.com/araffin/rl-baselines-zoo">Stable Baseline Zoo</a></p> | |
| <p> | |
| There are currently 3 models: | |
| <ul> | |
| <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4">PPO SpaceInvadersNoFrameskip-v4</a></li> | |
| <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-LunarLander-v2">PPO LunarLander-v2</a></li> | |
| <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-CartPole-v1">PPO CartPole-v1</a></li> | |
| </ul> | |
| </div>''' | |
| ) | |
| iface.launch() | |
| """ | |
| TODO: Next version with live video generation | |
| import gradio as gr | |
| import os | |
| from Recorder import Recorder | |
| from stable_baselines3 import PPO | |
| #The Agent plays and we generate the video | |
| def replay(option): | |
| video_path = "" | |
| # Get the correct model | |
| if (option == "LunarLander-v2 🚀👩🚀"): | |
| env_name = "Lunar Lander v2" | |
| agent_name = "PPO" | |
| print("TEST") | |
| hf_model_filename = "LunarLander-v2" | |
| hf_model_id = "ThomasSimonini/stable-baselines3-ppo-LunarLander-v2" | |
| video_path = replay_gym(hf_model_filename, hf_model_id) | |
| elif(option == "CartPole-v1 🕹️"): | |
| hf_model_filename = "CartPole-v1" | |
| hf_model_id = "ThomasSimonini/stable-baselines3-ppo-CartPole-v1" | |
| video_path = replay_gym(hf_model_filename, hf_model_id) | |
| elif(option == "Atari Space Invaders 👾"): | |
| hf_model_filename = "SpaceInvadersNoFrameskip-v4" | |
| hf_model_id = "ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4" | |
| video_path = replay_atari(hf_model_filename, hf_model_id) | |
| #video_path = "./SpaceInvadersNoFrameskip-v4.mp4" | |
| return video_path | |
| def replay_gym(hf_model_filename, hf_model_id): | |
| import gym | |
| from stable_baselines3.common.evaluation import evaluate_policy | |
| model = PPO.load_from_huggingface(hf_model_id,hf_model_filename) | |
| eval_env = gym.make(hf_model_filename) | |
| directory = './video' | |
| env = Recorder(eval_env, directory) | |
| obs = env.reset() | |
| done = False | |
| while not done: | |
| action, _state = model.predict(obs) | |
| obs, reward, done, info = env.step(action) | |
| clip = env.play() | |
| return clip | |
| def replay_atari(hf_model_filename, hf_model_id): | |
| os.system("python -m atari_py.import_roms \"content/atari_roms\"") | |
| import gym | |
| from stable_baselines3.common.env_util import make_atari_env | |
| from stable_baselines3.common.vec_env import VecFrameStack | |
| from stable_baselines3.common.evaluation import evaluate_policy | |
| model = PPO.load_from_huggingface(hf_model_id, hf_model_filename) | |
| eval_env = make_atari_env(hf_model_filename, n_envs=1, seed=0) | |
| eval_env = VecFrameStack(eval_env, n_stack=4) | |
| model = PPO.load_from_huggingface(hf_model_id, hf_model_filename) | |
| import gym | |
| directory = './video' | |
| env = Recorder(eval_env, directory) | |
| obs = env.reset() | |
| done = False | |
| while not done: | |
| action, _state = model.predict(obs) | |
| obs, reward, done, info = env.step(action) | |
| clip = env.play() | |
| return clip | |
| iface = gr.Interface( | |
| replay, | |
| [ | |
| gr.inputs.Dropdown(["Atari Space Invaders 👾", "CartPole-v1 🕹️", "LunarLander-v2 🚀👩🚀"]), | |
| ], | |
| "video", | |
| title = 'Stable Baselines 3 with 🤗', | |
| description = '', | |
| article = | |
| '''<div> | |
| <p style="text-align: center">This version of the RL library allows you to load models directly from the Hugging Face Hub</p> | |
| <p style="text-align: center"> Select the trained agent you want to watch perform. We record your agent playing. | |
| <p style="text-align: center"> Don't forget to <b>click on clear between each record.</b> </p> | |
| These models are from <a href="https://github.com/araffin/rl-baselines-zoo">Stable Baseline Zoo</a></p> | |
| <p> | |
| There are currently 3 models: | |
| <ul> | |
| <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4">PPO SpaceInvadersNoFrameskip-v4</a></li> | |
| <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-LunarLander-v2">PPO LunarLander-v2</a></li> | |
| <li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-CartPole-v1">PPO CartPole-v1</a></li> | |
| </ul> | |
| </div>''' | |
| ) | |
| iface.launch() | |
| """ | |