Views
No views yet
1!pip install stable-baselines3[extra] gymnasium[box2d] huggingface_hub
2
3import gymnasium as gym
4from stable_baselines3 import PPO
5from stable_baselines3.common.monitor import Monitor
6from stable_baselines3.common.evaluation import evaluate_policy
7from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder
8from huggingface_hub import hf_hub_download
9import base64
10from IPython import display
11import os
12
13# Repository details
14REPO_ID = "Srgreen/ppo-LunarLander-v3"
15FILENAME = "ppo-LunarLander-v3.zip"
16
17print("Downloading the trained model from the Hugging Face Hub...")
18checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
19
20# 1. Create a vectorized environment (required for the video recorder)
21video_folder = "./videos"
22env = DummyVecEnv([lambda: Monitor(gym.make("LunarLander-v3", render_mode="rgb_array"))])
23
24# 2. Wrap the environment to record the video of the simulation
25env = VecVideoRecorder(
26 env,
27 video_folder,
28 record_video_trigger=lambda x: x == 0, # Records the episode
29 video_length=1000,
30 name_prefix="lunar-lander-eval"
31)
32
33# 3. Load the trained PPO model
34print("Loading model weights into the PPO agent...")
35model = PPO.load(checkpoint_path, env=env)
36
37# 4. Evaluate the agent over 10 episodes to get the official metrics
38print("Evaluating the agent over 10 episodes...")
39mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=10, deterministic=True)
40
41print("-" * 40)
42print(f"Mean Reward: {mean_reward:.2f} +/- {std_reward:.2f}")
43print("-" * 40)
44if mean_reward >= 200:
45 print("Result: Successful pilot! Perfect landing on the Moon surface. 🌛🥳")
46else:
47 print("Result: The agent could use more training steps.")
48print("-" * 40)
49
50# 5. Run one more episode just to show the recorded landing video
51print("Preparing the video playback...")
52obs = env.reset()
53done = False
54while not done:
55 action, _states = model.predict(obs, deterministic=True)
56 obs, rewards, dones, infos = env.step(action)
57 done = dones[0]
58
59# Close the environment to save the video file properly
60env.close()
61
62# 6. Helper function to render the recorded MP4 video inside Google Colab
63def show_video(directory):
64 html = []
65 for filename in os.listdir(directory):
66 if filename.endswith(".mp4"):
67 video_path = os.path.join(directory, filename)
68 video_b64 = base64.b64encode(open(video_path, 'rb').read()).decode('ascii')
69 html.append(f'''
70 <video controls width="600" autoplay loop muted>
71 <source src="data:video/mp4;base64,{video_b64}" type="video/mp4" />
72 </video>
73 ''')
74 return "".join(html)
75
76# Display the video in the notebook
77print("Here is your agent's landing simulation:")
78display.display(display.HTML(show_video(video_folder)))