예제 소스

cartpol-origin.py

import os
import pickle
import random
import sys

import gymnasium as gym
import numpy as np
import tensorflow as tf
from collections import deque

from tensorflow.keras.layers import Dense
from tensorflow.keras.initializers import RandomUniform
from tensorflow.keras.optimizers import Adam


SAVE_DIR = os.path.join(os.path.dirname(__file__), "save_model_cartpol_origin")
WEIGHTS_PATH = os.path.join(SAVE_DIR, "model.weights.h5")
TARGET_WEIGHTS_PATH = os.path.join(SAVE_DIR, "target_model.weights.h5")
REPLAY_PATH = os.path.join(SAVE_DIR, "replay.pkl")


class DQN(tf.keras.Model):
    def __init__(self, action_size):
        super().__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)
        return self.fc_out(x)


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.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):
        if np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)
        return int(np.argmax(self.model.predict(state, verbose=0)))

    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])

        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 = tf.stop_gradient(self.target_model(next_states))
            max_q = np.amax(target_predicts, axis=-1)
            targets = rewards + (1 - dones) * self.discount_factor * max_q
            loss = tf.reduce_mean(tf.square(targets - predicts))

        trainable_vars = self.model.trainable_variables
        grads = tape.gradient(loss, trainable_vars)
        grads_and_vars = [(g, v) for g, v in zip(grads, trainable_vars) if g is not None]
        self.optimizer.apply_gradients(grads_and_vars)


def _build_models(agent, flat_dim):
    dummy = np.zeros((1, flat_dim), dtype=np.float32)
    agent.model(dummy)
    agent.target_model(dummy)


def save_checkpoint(agent):
    os.makedirs(SAVE_DIR, exist_ok=True)
    _build_models(agent, agent.state_size)
    agent.model.save_weights(WEIGHTS_PATH)
    agent.target_model.save_weights(TARGET_WEIGHTS_PATH)
    payload = {
        "memory": list(agent.memory),
        "epsilon": float(agent.epsilon),
    }
    with open(REPLAY_PATH, "wb") as f:
        pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
    print(f"Saved checkpoint → {SAVE_DIR}")


def load_checkpoint(agent, flat_dim):
    if not (os.path.isfile(WEIGHTS_PATH) and os.path.isfile(TARGET_WEIGHTS_PATH) and os.path.isfile(REPLAY_PATH)):
        return False
    _build_models(agent, flat_dim)
    agent.model.load_weights(WEIGHTS_PATH)
    agent.target_model.load_weights(TARGET_WEIGHTS_PATH)
    with open(REPLAY_PATH, "rb") as f:
        payload = pickle.load(f)
    agent.memory.clear()
    for t in payload["memory"]:
        agent.memory.append(t)
    agent.epsilon = float(payload.get("epsilon", agent.epsilon))
    print(f"Loaded checkpoint from {SAVE_DIR} (memory len={len(agent.memory)}, epsilon={agent.epsilon:.4f})")
    return True


if __name__ == "__main__":
    env = gym.make("CartPole-v1")
    state_size = int(np.prod(env.observation_space.shape))
    action_size = int(env.action_space.n)
    agent = DQNAgent(state_size, action_size)

    load_checkpoint(agent, state_size)

    num_episodes = 15
    for _ in range(num_episodes):
        state, _ = env.reset()
        state = np.asarray(state, dtype=np.float32).reshape(1, -1)
        done = False

        while not done:
            action = agent.choose_action(state)
            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            next_state = np.asarray(next_state, dtype=np.float32).reshape(1, -1)
            agent.remember(state, action, reward, next_state, done)
            state = next_state
            if len(agent.memory) >= agent.train_start:
                agent.train_model()

        agent.update_target_model()

    save_checkpoint(agent)
    env.close()