Views
No views yet
1import gymnasium as gym
2from stable_baselines3 import PPO
3from stable_baselines3.common.env_util import make_atari_env
4from stable_baselines3.common.vec_env import VecTransposeImage
5from stable_baselines3.common.atari_wrappers import WarpFrame
6
7# Load the trained model
8model = PPO.load("best-model.zip")
9
10# Create the environment
11env = make_atari_env("BreakoutNoFrameskip-v4", n_envs=1)
12env = VecFrameStack(env, n_stack=4)
13env = VecTransposeImage(env)
14
15# Reset the environment
16obs, info = env.reset()
17
18# Enjoy the trained agent
19for _ in range(1000):
20 action, _states = model.predict(obs, deterministic=True)
21 obs, rewards, terminated, truncated, info = env.step(action)
22 if terminated or truncated:
23 obs, info = env.reset()
24 env.render()
25env.close()pip install huggingface_hub1from huggingface_hub import hf_hub_download
2import torch as th
3import gymnasium as gym
4from stable_baselines3 import PPO
5from stable_baselines3.common.env_util import make_atari_env
6from stable_baselines3.common.vec_env import VecTransposeImage
7from stable_baselines3.common.atari_wrappers import WarpFrame
8
9# Download the model from the Hub
10model_path = hf_hub_download(repo_id="kuds/atari-breakout-v4-ppo", filename="best-model.zip")
11
12
13
14
15# Load the model
16model = PPO.load(model_path)
17
18# Create the environment
19env = make_atari_env("BreakoutNoFrameskip-v4", n_envs=1)
20env = VecFrameStack(env, n_stack=4)
21env = VecTransposeImage(env)
22
23# Enjoy the trained agent
24obs = env.reset()
25for i in range(1000):
26 action, _states = model.predict(obs, deterministic=True)
27 obs, rewards, dones, info = env.step(action)
28 env.render("human")
29
30env.close()