예제 소스

LunarLander-dqn.py

import gymnasium as gym
import numpy as np
import tensorflow as tf

from tensorflow.keras import models, layers, optimizers
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.initializers import RandomUniform
from tensorflow.keras.optimizers import Adam

from collections import deque
import random
import matplotlib.pyplot as plt
import cv2
import os

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

class DQNAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        self.discount_factor = 0.99
        self.learning_rate = 0.001
        self.epsilon = 1.0
        self.epsilon_decay = 0.999
        self.epsilon_min = 0.01
        self.batch_size = 64
        self.train_start = 1000
        self.memory = deque(maxlen=2000)
        self.model = DQN(action_size)
        self.target_model = DQN(action_size)
        self.optimizer = Adam(learning_rate=self.learning_rate)
        
        # 모델 빌드
        self.model.build(input_shape=(None, state_size))
        self.target_model.build(input_shape=(None, state_size))
        
        self.load_model()
        self.update_target_model()

    def update_target_model(self):
        self.target_model.set_weights(self.model.get_weights())
    
    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def choose_action(self, state):
        return random.randrange(self.action_size) if (np.random.rand() <= self.epsilon) else np.argmax(self.model.predict(state))

    def train_model(self):
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

        mini_batch = random.sample(self.memory, self.batch_size)

        states = np.array([sample[0][0] for sample in mini_batch])
        actions = np.array([sample[1] for sample in mini_batch])
        rewards = np.array([sample[2] for sample in mini_batch])
        next_states = np.array([sample[3][0] for sample in mini_batch])
        dones = np.array([sample[4] for sample in mini_batch])

        # 입력 데이터 형태 확인
        print(f"States shape: {states.shape}")
        print(f"Next states shape: {next_states.shape}")

        with tf.GradientTape() as tape:
            predicts = self.model(states)
            one_hot_action = tf.one_hot(actions, self.action_size)
            predicts = tf.reduce_sum(one_hot_action * predicts, axis=1)

            target_predicts = self.target_model(next_states)
            target_predicts = tf.stop_gradient(target_predicts)

            max_q = np.amax(target_predicts, axis=-1)
            targets = rewards + (1 - dones) * self.discount_factor * max_q

            # 손실 함수와 출력 형태 확인
            print(f"Predicts shape: {predicts.shape}")
            print(f"Targets shape: {targets.shape}")

            loss = tf.reduce_mean(tf.square(targets - predicts))

        trainable_vars = self.model.trainable_variables
        grads = tape.gradient(loss, trainable_vars)

        if grads is None:
            print("Warning: gradients are None. Check your model and data.")
            return
        grads_and_vars = [(g, v) for g, v in zip(grads, trainable_vars) if g is not None]
        if not grads_and_vars:
            print("Warning: No gradients to apply.")
            return
        try:
            self.optimizer.apply_gradients(grads_and_vars)
        except Exception as e:
            print(f"Error applying gradients: {e}")
    
    def load_model(self):
        if os.path.exists('./save_model_LunarLander/model.index'):
            self.model.load_weights('./save_model_LunarLander/model')
            print("Model loaded successfully!")

if __name__ == "__main__":
    # Create the environment
    env = gym.make('LunarLander-v2', render_mode="rgb_array")
    state_size = env.observation_space.shape[0]
    action_size = env.action_space.n

    agent = DQNAgent(state_size, action_size)

    scores, episodes = [], []
    score_avg = 0

    num_episode = 300
    for e in range(num_episode):
        done = False
        score = 0
        max_score = -np.inf

        state, _ = env.reset()
        state = state.reshape(1, -1)

        while not done:
            img = env.render()

            if img is not None:
                cv2.imshow('LunarLander', img)
                cv2.waitKeyEx(delay=1)

            action = agent.choose_action(state)

            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            next_state = next_state.reshape(1, -1)

            # Modify reward
            # Add penalty for moving away from the target
            if abs(next_state[0][0]) > abs(state[0][0]):
                reward -= 10

            # Add penalty for moving too horizontally
            if abs(next_state[0][1]) < 0.05 and abs(next_state[0][3]) < 0.05:
                reward -= 10

            score += reward

            agent.remember(state, action, reward, next_state, done)

            if len(agent.memory) >= agent.train_start:
                agent.train_model()

            state = next_state

            if done:
                if score > max_score :
                    max_score = score
                    agent.update_target_model()
                    agent.model.save_weights('./save_model_LunarLander/model.weights.h5')

                cv2.destroyAllWindows()

                score_avg = score / len(agent.memory)
                print('episode: {:3d} | score avg {:3.2f} | memory length: {:4d} | epsilon: {:.4f}'.format(e, score_avg, len(agent.memory), agent.epsilon))

                scores.append(score_avg)
                episodes.append(e)
                plt.plot(episodes, scores, 'b')
                plt.xlabel('episode')
                plt.ylabel('average score')
                # plt.savefig('LunarLander_graph.png')
                plt.show(block = False)
                plt.pause(0.000001)
    env.close()