Update app.py

app.py

@@ -1,8 +1,48 @@
 import gradio as gr
+from huggingface_sb3 import load_from_hub
+
+import gym
+import os 
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import VecNormalize
+
+from stable_baselines3.common.env_util import make_atari_env
+from stable_baselines3.common.vec_env import VecFrameStack
+
+from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv
 
 
 def replay(model_id, filename, environment, evaluate):
-    pass
+	# Load the model
+	checkpoint = load_from_hub(model_id, filename)
+    
+	# Because we using 3.7 on Colab and this agent was trained with 3.8 to avoid Pickle errors:
+	custom_objects = {
+            "learning_rate": 0.0,
+            "lr_schedule": lambda _: 0.0,
+            "clip_range": lambda _: 0.0,}
+	
+	model= PPO.load(checkpoint, custom_objects=custom_objects)
+
+	eval_env = make_atari_env(environment, n_envs=1)
+	eval_env = VecFrameStack(env, n_stack=4)
+
+	video_folder = 'logs/videos/'
+	video_length = 100
+
+
+	# Record the video starting at the first step
+	env = VecVideoRecorder(eval_env, video_folder,
+                       record_video_trigger=lambda x: x == 0, video_length=video_length,
+                       name_prefix=f"random-agent-{env_id}")
+
+	env.reset()
+	for _ in range(video_length + 1):
+		action, _states = model.predict(obs)
+    	obs, rewards, dones, info = env.step(action)
+	# Save the video
+	env.close()
+
 
 iface = gr.Interface(fn=replay, inputs=[
                                         gr.inputs.Textbox(lines=1, placeholder=None, default="", label="Model Id: "),