Tetris-RL / evaluate.py
BaljinderH's picture
Update evaluate.py
dd8faff verified
raw
history blame contribute delete
689 Bytes
import gym
from stable_baselines3 import DQN
from tetris_env import TetrisEnv
import pygame
import time
def main():
env = TetrisEnv()
model = DQN.load("models/dqn_tetris")
episodes = 5
for ep in range(1, episodes + 1):
obs = env.reset()
done = False
total_reward = 0
while not done:
env.render(mode='human')
action, _states = model.predict(obs, deterministic=True)
obs, reward, done, info = env.step(action)
total_reward += reward
pygame.time.wait(100)
print(f"Episode {ep}: Total Reward = {total_reward}")
env.close()
if __name__ == "__main__":
main()