import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.evaluation import evaluate_policy
from huggingface_hub import hf_hub_download

# -------------------------
# Environment Setup
# -------------------------
env = gym.make("LunarLander-v2", render_mode="human")       # Human render
eval_env = Monitor(gym.make("LunarLander-v2"))              # Evaluation (no render)

# -------------------------
# Load pretrained model
# -------------------------
model_path = hf_hub_download("Vishand03/lunarlander-ppo", "model.zip")
model = PPO.load(model_path)

# -------------------------
# Run one episode
# -------------------------
obs, _ = env.reset()
done = False
while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, _ = env.step(action)
    done = terminated or truncated

# -------------------------
# Evaluate policy
# -------------------------
mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10, deterministic=True)
print(f"Mean Reward: {mean_reward:.2f} +/- {std_reward:.2f}")