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

MEMORY_SIZE = 1000000


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:
        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
            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
            if terminal:
                print("Run: " + str(run) + ", exploration: " + str(dqn_solver.exploration_rate) + ", score: " + str(step))
                #score_logger.add_score(step, run)

if __name__ == "__main__":
  • $\begingroup$ To answer your question, we need to know more details about your problem, learning algorithm, model, etc. There are many RL algorithms. $\endgroup$ – nbro Aug 23 '20 at 21:09
  • $\begingroup$ @nbro Thank you for your comment. I have included some code and the name of my model. Let me know if you have more questions. $\endgroup$ – desert_ranger Aug 23 '20 at 21:17
  • $\begingroup$ That's a lot better. Thanks! Just a clarification, you're not necessarily looking for code (although this is appreciated), but you mainly want to understand which variables are required to restarting training, right? $\endgroup$ – nbro Aug 23 '20 at 21:22
  • $\begingroup$ @nbro - Exactly. This is just some sample code that I got and modified from a source. Once I understand the variables/parameters that need to be saved, I will implement it on my research code (which mainly follows the DQN algorithm) $\endgroup$ – desert_ranger Aug 23 '20 at 21:35

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

    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)

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:
    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
  • $\begingroup$ Could you please post a sample of your code. That way, I can get an idea on how the pseudocode works. $\endgroup$ – desert_ranger Aug 31 '20 at 21:39
  • 1
    $\begingroup$ I've added this now $\endgroup$ – David Ireland Sep 1 '20 at 11:44

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.