File size: 2,412 Bytes
8097aca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | # %%
# Import required packages
import gymnasium as gym
from huggingface_sb3 import package_to_hub
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder
# %%
# Test random environment
env_id = "LunarLander-v3"
env = gym.make(env_id)
observation, info = env.reset()
for _ in range(20):
action = env.action_space.sample()
print("Action taken:", action)
observation, reward, terminated, truncated, info = env.step(action)
if terminated or truncated:
print("Environment is reset")
observation, info = env.reset()
env.close()
# %%
# Check observation and action spaces
env.reset()
print("_____OBSERVATION SPACE_____ \n")
print("Observation Space Shape", env.observation_space.shape)
print("Sample observation", env.observation_space.sample()) # Get a random observation
print("\n _____ACTION SPACE_____ \n")
print("Action Space Shape", env.action_space.n)
print("Action Space Sample", env.action_space.sample()) # Take a random action
# %%
# Check SB3 model device
model = PPO("MlpPolicy", env, device="auto")
print(model.device)
# %%
# Train PPO agent
model = PPO(
policy="MlpPolicy",
env=env,
n_steps=1024,
batch_size=64,
n_epochs=4,
gamma=0.999,
gae_lambda=0.98,
ent_coef=0.01,
verbose=1,
)
model.learn(total_timesteps=500_000)
# %%
# Train agent for 1M timesteps
model.learn(total_timesteps=1_000_000)
model.save("ppo-lunar-lander")
# %%
# Evaluate the agent
model = PPO.load("ppo-lunar-lander", env=env)
eval_env = Monitor(gym.make(env_id))
mean_reward, std_reward = evaluate_policy(
model, eval_env, n_eval_episodes=100, deterministic=True
)
print(f"mean_reward={mean_reward:.2f} +/- {std_reward}")
# %%
# Publish the trained agent
eval_env = DummyVecEnv(
[lambda: Monitor(gym.make(env_id, render_mode="rgb_array"))]
)
eval_env = VecVideoRecorder(
eval_env,
"videos/",
record_video_trigger=lambda x: x == 0,
video_length=1000,
name_prefix="ppo-lunar-lander-demo",
)
package_to_hub(
model=model,
model_name="ppo-lunar-lander-v2",
model_architecture="PPO",
env_id=env_id,
eval_env=eval_env,
repo_id="pabloramesc/ppo-lunar-lander-v2",
commit_message="Upload PPO agent for LunarLander-v2",
)
# %%
|