"""Score the hand-edited TQC hopper on Hopper-v4.

    pip install "stable-baselines3>=2.9" "sb3-contrib>=2.9" "gymnasium[mujoco]>=1.0"
    python example.py tqc-Hopper-v4-hand-edited.zip 0 100
"""
import sys
import numpy as np
import gymnasium as gym
from sb3_contrib import TQC

path, first, n = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
model = TQC.load(path, device="cpu")
env = gym.make("Hopper-v4")
returns, falls = [], 0
for seed in range(first, first + n):
    obs, _ = env.reset(seed=seed)
    total = 0.0
    for _ in range(1000):
        action, _ = model.predict(obs, deterministic=True)
        obs, reward, terminated, truncated, _ = env.step(action)
        total += reward
        if terminated or truncated:
            break
    returns.append(total)
    falls += int(terminated)
print(f"{n} episodes: mean return {np.mean(returns):.1f}, falls {falls} ({100 * falls / n:.1f}%)")