What are the variables that need to be saved and loaded, so that a DQN model starts where it left off?

Question

TensorFlow allows users to save the weights and the model architecture, however, that will be insufficient unless the values of certain other variables are also stored. For instance, in DQN, if $\epsilon$ is not stored the model will start exploring from scratch and a new model will have to be trained.

What are the variables that need to be saved and loaded, so that a DQN model starts where it left off? Some pseudocode will be highly appreciated!

Here is my current model with code

## Slightly modified from the following repository - https://github.com/gsurma/cartpole

from __future__ import absolute_import, division, print_function, unicode_literals

import os
import random
import gym
import numpy as np
import tensorflow as tf

from collections import deque
from tensorflow.models import Sequential
from tensorflow.layers import Dense
from tensorflow.optimizers import Adam


ENV_NAME = "CartPole-v1"

GAMMA = 0.95
LEARNING_RATE = 0.001

MEMORY_SIZE = 1000000
BATCH_SIZE = 20

EXPLORATION_MAX = 1.0
EXPLORATION_MIN = 0.01
EXPLORATION_DECAY = 0.995

checkpoint_path = "training_1/cp.ckpt"


class DQNSolver:

    def __init__(self, observation_space, action_space):
        self.exploration_rate = EXPLORATION_MAX

        self.action_space = action_space
        self.memory = deque(maxlen=MEMORY_SIZE)

        self.model = Sequential()
        self.model.add(Dense(24, input_shape=(observation_space,), activation="relu"))
        self.model.add(Dense(24, activation="relu"))
        self.model.add(Dense(self.action_space, activation="linear"))
        self.model.compile(loss="mse", optimizer=Adam(lr=LEARNING_RATE))

    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def act(self, state):
        if np.random.rand() < self.exploration_rate:
            return random.randrange(self.action_space)
        q_values = self.model.predict(state)
        return np.argmax(q_values[0])

    def experience_replay(self):
        if len(self.memory) < BATCH_SIZE:
            return
        batch = random.sample(self.memory, BATCH_SIZE)
        for state, action, reward, state_next, terminal in batch:
            q_update = reward
            if not terminal:
                q_update = (reward + GAMMA * np.amax(self.model.predict(state_next)[0]))
            q_values = self.model.predict(state)
            q_values[0][action] = q_update
            self.model.fit(state, q_values, verbose=0)
        self.exploration_rate *= EXPLORATION_DECAY
        self.exploration_rate = max(EXPLORATION_MIN, self.exploration_rate)


def cartpole():
    env = gym.make(ENV_NAME)
    #score_logger = ScoreLogger(ENV_NAME)
    observation_space = env.observation_space.shape[0]
    action_space = env.action_space.n
    dqn_solver = DQNSolver(observation_space, action_space)
    checkpoint = tf.train.get_checkpoint_state(os.getcwd()+"/saved_networks")
    print('checkpoint:', checkpoint)
    if checkpoint and checkpoint.model_checkpoint_path:
        dqn_solver.model = keras.models.load_model('cartpole.h5')
        dqn_solver.model = model.load_weights('cartpole_weights.h5')
        
    run = 0
    i = 0
    while i<5:
        i = i + 1
        #total = 0
        run += 1
        state = env.reset()
        state = np.reshape(state, [1, observation_space])
        step = 0
        while True:
            step += 1
            #env.render()
            action = dqn_solver.act(state)
            state_next, reward, terminal, info = env.step(action)
            #total += reward
            reward = reward if not terminal else -reward
            state_next = np.reshape(state_next, [1, observation_space])
            dqn_solver.remember(state, action, reward, state_next, terminal)
            state = state_next
            dqn_solver.model.save('cartpole.h5')
            dqn_solver.model.save_weights('cartpole_weights.h5')
            if terminal:
                print("Run: " + str(run) + ", exploration: " + str(dqn_solver.exploration_rate) + ", score: " + str(step))
                #score_logger.add_score(step, run)
                break
            dqn_solver.experience_replay()


if __name__ == "__main__":
    cartpole()

Answer 1

Typically you would need to save the network weights, hyper-parameters and the replay buffer if you wanted to stop training and then come back at a later date and carry on training. Usually, I do this by writing it all as a class in Python (the agent, the memory buffer, hyper-parameters etc.) and saving the final object with Pickle.

Looking at your code, the only thing I would personally have done different would be to define the model outside of the class and have the class take as input a network; however I usually use PyTorch as opposed to Keras/Tensorflow so I'm not sure which method works better.

As per OP's request in the comments, here is a snippet of code I used for Car-pool.

class DQN:

def __init__(self, observation_space, action_space):
    self.exploration_rate = epsilon_max
    self.observation_space = observation_space
    self.batch_size = batch_size
    self.gamma = gamma

    self.action_space = action_space

    self.memory = deque(maxlen=max_memory)

    self.model = Net()
    self.loss_fn = nn.MSELoss()
    self.optimizer = optim.Adam(self.model.parameters(), lr=0.0001)

def remember(self, state, action, reward, next_state, done):
    self.memory.append((state, action, reward, next_state, done))

def act(self, state):
    self.model.eval()
    state = torch.from_numpy(state).type(torch.FloatTensor)
    self.exploration_rate *= exploration_decay
    self.exploration_rate = max(epsilon_min, self.exploration_rate)
    if np.random.uniform() < self.exploration_rate:
        return np.random.randint(0, self.action_space)
    q_values = self.model(state)
    action = torch.argmax(q_values, dim=1)
    return int(action)

def experience_replay(self):
    if len(self.memory) < batch_size:
        return

    self.optimizer.zero_grad()
    states,actions,rewards,next_states,dones = self.get_batch()
    states = torch.FloatTensor(states).reshape((self.batch_size,4))
    next_states = torch.FloatTensor(next_states).reshape((self.batch_size,4))
    rewards = torch.FloatTensor(rewards).reshape((self.batch_size,1))
    dones = torch.FloatTensor(dones).reshape((self.batch_size,1))

    # get the q update ready
    # first take max and set up target
    max_vals,argmax = torch.max(self.model(next_states),axis=1)
    q_target = rewards + self.gamma * max_vals.reshape((self.batch_size,1)) * (1-dones)

    q_values = self.model(states)
    for _ in range(self.batch_size):
        q_values[_][actions[_]] = q_target[_][0]
    input = self.model(states)
    q_values = q_values.detach()
    loss = self.loss_fn(input, q_values)
    loss.backward()
    self.optimizer.step()

def get_batch(self):
    batch = random.sample(self.memory, self.batch_size)

    states = []
    actions = []
    rewards = []
    next_states = []
    dones = []
    for state,action,reward,next_state,done in batch:
        states.append(state)
        actions.append(action)
        rewards.append(reward)
        next_states.append(next_state)
        dones.append(done)
    return states,actions,rewards,next_states,dones


def tensor_max(self, ten):
    ten = ten.numpy()
    ten = np.amax(ten, 1).reshape(1, 1)
    ten = torch.from_numpy(ten).type(torch.FloatTensor)
    return ten