Upload unit1.py

unit1.py

@@ -0,0 +1,102 @@
+# %%
+# 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",
+)
+# %%