File size: 689 Bytes
b100cf9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dd8faff
b100cf9
 
 
 
 
 
 
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
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()