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("PongNoFrameskip-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-pong-v4-ppo", filename="best-model.zip")
11
12# Load the model
13model = PPO.load(model_path)
14
15# Create the environment
16env = make_atari_env("PongNoFrameskip-v4", n_envs=1)
17env = VecFrameStack(env, n_stack=4)
18env = VecTransposeImage(env)
19
20# Enjoy the trained agent
21obs = env.reset()
22for i in range(1000):
23 action, _states = model.predict(obs, deterministic=True)
24 obs, rewards, dones, info = env.step(action)
25 env.render("human")