예제 소스

cartpol.py

import gymnasium as gym

import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.initializers import RandomUniform

class DQN(tf.keras.Model):
    def __init__(self, action_size):
        super(DQN, self).__init__()
        self.fc1 = Dense(24, activation='relu')
        self.fc2 = Dense(24, activation='relu')
        self.fc_out = Dense(action_size, kernel_initializer=RandomUniform(-1e-3, 1e-3))

    def call(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        q = self.fc_out(x)
        return q
    
env = gym.make('CartPole-v1', render_mode="rgb_array")
env.reset()

def render():
    _, _, terminated, truncated, _ = env.step(env.action_space.sample())
    if terminated or truncated:
        env.reset()
    img = env.render()
    return img

while True:
    img = render()

    if img is not None:
        # 이미지를 시각화합니다.
        plt.imshow(img)
        plt.axis('off')  # 축을 숨깁니다.
        plt.show(block=False)  # Non-blocking
        # plt.pause(0.01)  # 짧은 딜레이

    if plt.waitforbuttonpress(0.1):  # 키 입력 대기 (0.1초 동안)
        key = plt.gcf().canvas.get_tk_widget().tk.call("focus_get")
        if key == 'q':  # 'q' 키를 누르면 종료
            break

env.close()
plt.close()  # 모든 플롯을 닫습니다.