import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
import os
import random
import gym

import matplotlib.pyplot as plt


from experience_replay_buffer import ExpBuffer
from utils import preprocessImage, FrameBuffer



tf.reset_default_graph()

#####################
#  HYPERPARAMETERS  #
#####################
env_name = "Pong-v0"
ep_max = int(1e6)          # max number of episodes to learn from
step_tot_max = int(200e6)
step_max = int(1e6)                # max steps in an episode
gamma = 0.99                   # future reward discount
learning_rate = float(1e-4)         # Q-network learning rate
seed = 42
action_repeat = 4
# Exploration parameters
explore_start = 1.0            # exploration probability at start
explore_stop = 0.1            # minimum exploration probability
explore_steps = int(1e6)
explore_decay_rate = (explore_start - explore_stop) / explore_steps            # exponential decay rate for exploration prob
# Network parameters
frame_buffer_size = 4
hidden_size = 512               # number of units in each Q-network hidden layer
QNet_update_freq = int(1e4)
# Experience buffer parameters
exp_buffer_size = int(1e6)            # ExpBuffer capacity
minibatch_size = 32                # experience mini-batch size
pretrain_step_max = minibatch_size+frame_buffer_size   # number experiences to pretrain the ExpBuffer
# Save parameters
model_save_interval = 100
model_save_dir = "./saves"
save_model = False
# TensorBoard
tensorboard_update_freq = 50


# Initialize environment
env = gym.make(env_name)
env_actions = env.action_space.n
# Set random seeds
def setAllRandomSeeds(environment, seed=42):
    random.seed(seed)
    np.random.seed(seed)
    tf.set_random_seed(seed)
    environment.seed(seed)
setAllRandomSeeds(env, seed)


print("\n*********************")
print("*  Run information  *")
print("*********************")
print("  Environment name   : {}".format(env_name))
print("  Action space       : {}".format(env.action_space))
print("  Observation space  : {}".format(env.observation_space))
print("  ep_max             : {}".format(ep_max))
print("  step_tot_max       : {}".format(step_tot_max))
print("  step_max           : {}".format(step_max))
print("  gamma              : {}".format(gamma))
print("  learning_rate      : {}".format(learning_rate))
print("  seed               : {}".format(seed))
print("  action_repeat      : {}".format(action_repeat))
print("  explore_start      : {}".format(explore_start))
print("  explore_stop       : {}".format(explore_stop))
print("  explore_steps      : {}".format(explore_steps))
print("  explore_decay_rate : {}".format(explore_decay_rate))
print("  frame_buffer_size  : {}".format(frame_buffer_size))
print("  QNet_update_freq   : {}".format(QNet_update_freq))
print("  pretrain_step_max  : {}".format(pretrain_step_max))
print("  exp_buffer_size    : {}".format(exp_buffer_size))
print("  minibatch_size     : {}".format(minibatch_size))
print("")
print("  model_save_interval : {}".format(model_save_interval))
print("  model_save_dir      : {}".format(model_save_dir))
print("  save_model          : {}".format(save_model))
print("")
print("  tensorboard_update_freq : {}".format(tensorboard_update_freq))
print("")


# Initialize experience replay buffer object
ExpBuffer = ExpBuffer(max_size=exp_buffer_size)

# Initialize frame buffer objects for stacking frames when doing inference (i.e. predicting Q-value for 'a')
buff_f = FrameBuffer(size=frame_buffer_size)
buff_f_next = FrameBuffer(size=frame_buffer_size)


#########################
#  COMPUTATIONAL GRAPH  #
#########################\
def createTargetNetworkUpdateOperations(tf_trainable_vars):
    '''Create a list where each element is a tensorflow assign operation which overwrites
       trainable variables of 'target' with 'moving'.

    First HALF of all variables correpsond to the 'moving' Q-network.
    Later HALF are the 'target' Q-network.
    [ ... moving ... | ... target ... ]

    Note: The Q-network objects must be initialized in THIS ORDER !!!

    Args:
      tf_trainable_vars (tf var list) : List containing all trainable variables.
    Returns:
      op_list (list) : List containing tf operations to copy 'moving' -> 'target'.
    '''
    tot_vars = len(tf_trainable_vars)
    op_list = []
    # Get the 'moving' Q-network variables from the first half ('//'' is integer division)
    for i, moving_var in enumerate(tf_trainable_vars[0:tot_vars//2]):
        # Define a tf operation which copies the 'moving' variable into 'target' variable
        op_list.append( tf_trainable_vars[i+tot_vars//2].assign(moving_var.value()) )

    return op_list


def updateTargetNetwork(op_list, sess):
    '''Run previously defined tf operations to overwrite 'target' Q-network with 'moving' Q-network.'''
    for op in op_list:
        sess.run(op)


def saveModel(ep, save_dir, tf_saver, sess, save_model=False):
    '''Create a save model folder within the specified save directory.

    Args:
      ep (int)             : Episode number where the model is saved.
      save_dir (str)       : Path to the directory where the model will be saved
      tf_saver (tf object) : 
      sess : (tf object)   : 
      save_model : (bool)  : 
    '''
    if(save_model == False):
        return
    # Create the save directory if it does not already exist
    if(os.path.isdir(save_dir) == False):
        os.mkdir(save_dir)
    # Create save folder
    save_folder_path = os.path.join(save_dir, "save_{}".format(ep))
    os.mkdir(save_folder_path)
    # Define save name
    file_name = "save_{}.ckpt".format(ep)
    save_path = os.path.join(save_folder_path, file_name)

    print("Saving model : {}".format(save_path))

    tf_saver.save(sess, save_path)


class QNetwork:
    '''Class which represent the network which outputs a Q-value based on an input state.
    '''
    def __init__(self, learning_rate, img_w, img_h, frame_num, action_size, hidden_size, name='QNetwork'):
        '''Initialize all tensorflow operation objects when initializing Q-network object.

        Use the CNN function approximation to generate and store Q-value arrays.
        These Q-value arrays are then inputted into 'Q_input' placeholder in order to compute other things.
        '''
        # state inputs to the Q-network
        with tf.variable_scope(name):

            ##################
            #  PLACEHOLDERS  #
            ##################
            # Array with an integer for every sample
            self.actions_ = tf.placeholder(tf.int32, [None], name='actions')
            # One hot encode the actions to later choose the Q-value for the action
            one_hot_actions = tf.one_hot(self.actions_, action_size)
            # Target Q values for training
            self.targetQs_ = tf.placeholder(tf.float32, [None], name='target')

            ###########
            #  INPUT  #
            ###########
            self.inputs = tf.placeholder(tf.float32, [None, img_w, img_h, frame_num], name='inputs')
            
            ############
            #  LAYERS  #
            ############

            self.conv1 = slim.conv2d(inputs=self.inputs,
                                     num_outputs=32,
                                     kernel_size=[8,8],
                                     stride=[4,4],
                                     padding="VALID",
                                     biases_initializer=None)

            self.conv2 = slim.conv2d(inputs=self.conv1,
                                     num_outputs=64,
                                     kernel_size=[4,4],
                                     stride=[2,2],
                                     padding="VALID",
                                     biases_initializer=None)

            self.conv3 = slim.conv2d(inputs=self.conv2,
                                     num_outputs=64,
                                     kernel_size=[3,3],
                                     stride=[1,1],
                                     padding="VALID",
                                     biases_initializer=None)

            self.flatten = slim.flatten(self.conv3)

            self.fc4 = slim.fully_connected(self.flatten, hidden_size, activation_fn=tf.nn.relu, biases_initializer=None)

            self.output = slim.fully_connected(self.fc4, action_size, activation_fn=None, biases_initializer=None)

            # ReLU hidden layers
            #self.fc1 = tf.contrib.layers.fully_connected(self.inputs, hidden_size)
            #self.fc2 = tf.contrib.layers.fully_connected(self.fc1, hidden_size)
            # Linear output layer
            #self.output = tf.contrib.layers.fully_connected(self.fc2, action_size, activation_fn=None)
            
            ##############
            #  TRAINING  #
            ##############
            # Value of selected Q-value          
            self.Q = tf.reduce_sum(tf.multiply(self.output, one_hot_actions), axis=1)
            # L2 loss (targetQ - Q)^2
            self.loss = tf.reduce_mean(tf.square(self.targetQs_ - self.Q))
            tf.summary.scalar("loss", self.loss)
            # Training operations
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss)
            # TensorBoard operation
            self.merged = tf.summary.merge_all()

QNet_moving = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size, name='QNet_moving')
QNet_target = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size, name='QNet_target')


tf_init = tf.global_variables_initializer()
tf_saver = tf.train.Saver()
tf_trainable_vars = tf.trainable_variables()
# List with operations to overwrite the 'target' Q-network variables
QNet_target_update_op_list = createTargetNetworkUpdateOperations(tf_trainable_vars)


############################
#  INITIALIZE ENVIRONMENT  #
############################
# Fill frame buffers
buff_f.reset(env)
buff_f_next.reset(env)


####################
#  PRETRAIN STEPS  #
####################
# Fill replay buffer with randomly selected experiences before doing training steps
for pretrain_step in range(pretrain_step_max):
    # Make a random action
    a = env.action_space.sample()
    f_next, r, done, _ = env.step(a)
    f_next = preprocessImage(f_next)
    # Store experience to buffer
    exp = (f_next, a, r, done)
    ExpBuffer.store(exp)
    if done:      
        # Start new episode
        env.reset()


###############
#  MAIN LOOP  #
###############
saver = tf.train.Saver()
rewards_list = []
with tf.Session() as sess:
    # Initialize variables
    sess.run(tf_init)
    updateTargetNetwork(QNet_target_update_op_list, sess)
    # TensorBoard storage directory
    tb_train_writer = tf.summary.FileWriter('./logs/1/train', sess.graph)

    step_tot = 0

    for ep in range(1, ep_max):

        # Reset buffer at start of episode
        buff_f.reset(env)
        buff_f_next.reset(env)

        # Reset environment at start of episode
        env.reset()
        done = False
        step = 0
        # Store total reward for the episode
        ep_rew = 0    

        while step < step_max:

            step_tot += 1
            # env.render() 
            
            ###################
            #  SELECT ACTION  #
            ###################
            # Reduce exploration
            explore_p = explore_stop + (explore_start - explore_stop)*np.exp(-explore_decay_rate*step_tot) 
            if explore_p > np.random.rand():
                # Select exploratory random action
                a = env.action_space.sample()
            else:
                # Select optimal action according to policy (i.e. Q-network)
                Qs = sess.run(QNet_moving.output, feed_dict={QNet_moving.inputs: np.reshape(buff_f.sample(), (1,84,84,4))})
                a = np.argmax(Qs)
            
            #################
            #  TAKE ACTION  #
            #################
            # Do 'action repeat' -> Train only on every 4th frame
            r_repeat_tot = 0
            for _ in range(action_repeat):
                f_next, r_repeat, done, _ = env.step(a)
                r_repeat_tot += r_repeat
                if(done == True):
                    break
            r = r_repeat_tot
            
            f_next = preprocessImage(f_next)
            buff_f_next.store(f_next)

            ######################
            #  STORE EXPERIENCE  #
            ######################
            exp = (f_next, a, r, done)
            ExpBuffer.store(exp)
            
            ##############
            #  TRAINING  #
            ##############
            # Sample one minibatch from buffer into separate arrays
            # Each array index correspond to the same experience
            s_minibatch, a_minibatch, r_minibatch, s_next_minibatch, done_minibatch = ExpBuffer.sample(minibatch_size)
            
            # Compute 'target term'
            ep_termination = 1 - done_minibatch
            target_Q_array = sess.run(QNet_target.output, feed_dict={QNet_target.inputs: s_next_minibatch})
            targets = r_minibatch + gamma * np.max(target_Q_array, axis=1)*ep_termination

            # Update moving Q-network
            _ = sess.run(QNet_moving.opt, feed_dict={QNet_moving.inputs: s_minibatch,
                                                     QNet_moving.targetQs_: targets,
                                                     QNet_moving.actions_: a_minibatch})

            ###########################
            #  UPDATE TARGET NETWORK  #
            ###########################
            if(step_tot % QNet_update_freq == 0):
                updateTargetNetwork(QNet_target_update_op_list, sess) 

            # Update TensorBoard output
            if(step_tot % tensorboard_update_freq == 0):
                tb_summary = sess.run(QNet_moving.merged, feed_dict={QNet_moving.inputs: s_minibatch,
                                                                     QNet_moving.targetQs_: targets,
                                                                     QNet_moving.actions_: a_minibatch})
                tb_train_writer.add_summary(tb_summary, step_tot)

            # End of step
            f = f_next
            buff_f.store(f)
            step += 1
            ep_rew += r

            # End episode
            if(done == True):
                print('Episode: {}'.format(ep),
                      'Steps : {}'.format(step_tot),
                      'Total reward: {}'.format(ep_rew),
                      'Explore P: {:.4f}'.format(explore_p))
                rewards_list.append((ep,ep_rew,step_tot))
                break

    
        # Save model checkpoint
        if(ep % model_save_interval == 0):
            saveModel(ep, model_save_dir, tf_saver, sess, save_model)



eps, rews, steps = np.array(rewards_list).T
eps = np.reshape(eps, [-1,1])
rews = np.reshape(rews, [-1,1])
steps = np.reshape(steps, [-1,1])
out_array = np.concatenate((steps, rews),axis=1)
print(out_array.shape)
# Output rewards per episode
np.savetxt("rew_new.dat", out_array, fmt="%d")
#plt.plot(rewards_list)
#plt.show()

env.close()

80c84a36c10e3624b96a9e5ef42189e6239eaa8f
<mask token>


def createTargetNetworkUpdateOperations(tf_trainable_vars):
    """Create a list where each element is a tensorflow assign operation which overwrites
       trainable variables of 'target' with 'moving'.

    First HALF of all variables correpsond to the 'moving' Q-network.
    Later HALF are the 'target' Q-network.
    [ ... moving ... | ... target ... ]

    Note: The Q-network objects must be initialized in THIS ORDER !!!

    Args:
      tf_trainable_vars (tf var list) : List containing all trainable variables.
    Returns:
      op_list (list) : List containing tf operations to copy 'moving' -> 'target'.
    """
    tot_vars = len(tf_trainable_vars)
    op_list = []
    for i, moving_var in enumerate(tf_trainable_vars[0:tot_vars // 2]):
        op_list.append(tf_trainable_vars[i + tot_vars // 2].assign(
            moving_var.value()))
    return op_list


def updateTargetNetwork(op_list, sess):
    """Run previously defined tf operations to overwrite 'target' Q-network with 'moving' Q-network."""
    for op in op_list:
        sess.run(op)


<mask token>


class QNetwork:
    """Class which represent the network which outputs a Q-value based on an input state.
    """

    def __init__(self, learning_rate, img_w, img_h, frame_num, action_size,
        hidden_size, name='QNetwork'):
        """Initialize all tensorflow operation objects when initializing Q-network object.

        Use the CNN function approximation to generate and store Q-value arrays.
        These Q-value arrays are then inputted into 'Q_input' placeholder in order to compute other things.
        """
        with tf.variable_scope(name):
            self.actions_ = tf.placeholder(tf.int32, [None], name='actions')
            one_hot_actions = tf.one_hot(self.actions_, action_size)
            self.targetQs_ = tf.placeholder(tf.float32, [None], name='target')
            self.inputs = tf.placeholder(tf.float32, [None, img_w, img_h,
                frame_num], name='inputs')
            self.conv1 = slim.conv2d(inputs=self.inputs, num_outputs=32,
                kernel_size=[8, 8], stride=[4, 4], padding='VALID',
                biases_initializer=None)
            self.conv2 = slim.conv2d(inputs=self.conv1, num_outputs=64,
                kernel_size=[4, 4], stride=[2, 2], padding='VALID',
                biases_initializer=None)
            self.conv3 = slim.conv2d(inputs=self.conv2, num_outputs=64,
                kernel_size=[3, 3], stride=[1, 1], padding='VALID',
                biases_initializer=None)
            self.flatten = slim.flatten(self.conv3)
            self.fc4 = slim.fully_connected(self.flatten, hidden_size,
                activation_fn=tf.nn.relu, biases_initializer=None)
            self.output = slim.fully_connected(self.fc4, action_size,
                activation_fn=None, biases_initializer=None)
            self.Q = tf.reduce_sum(tf.multiply(self.output, one_hot_actions
                ), axis=1)
            self.loss = tf.reduce_mean(tf.square(self.targetQs_ - self.Q))
            tf.summary.scalar('loss', self.loss)
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss
                )
            self.merged = tf.summary.merge_all()


<mask token>

<mask token>
tf.reset_default_graph()
<mask token>


def setAllRandomSeeds(environment, seed=42):
    random.seed(seed)
    np.random.seed(seed)
    tf.set_random_seed(seed)
    environment.seed(seed)


setAllRandomSeeds(env, seed)
print("""
*********************""")
print('*  Run information  *')
print('*********************')
print('  Environment name   : {}'.format(env_name))
print('  Action space       : {}'.format(env.action_space))
print('  Observation space  : {}'.format(env.observation_space))
print('  ep_max             : {}'.format(ep_max))
print('  step_tot_max       : {}'.format(step_tot_max))
print('  step_max           : {}'.format(step_max))
print('  gamma              : {}'.format(gamma))
print('  learning_rate      : {}'.format(learning_rate))
print('  seed               : {}'.format(seed))
print('  action_repeat      : {}'.format(action_repeat))
print('  explore_start      : {}'.format(explore_start))
print('  explore_stop       : {}'.format(explore_stop))
print('  explore_steps      : {}'.format(explore_steps))
print('  explore_decay_rate : {}'.format(explore_decay_rate))
print('  frame_buffer_size  : {}'.format(frame_buffer_size))
print('  QNet_update_freq   : {}'.format(QNet_update_freq))
print('  pretrain_step_max  : {}'.format(pretrain_step_max))
print('  exp_buffer_size    : {}'.format(exp_buffer_size))
print('  minibatch_size     : {}'.format(minibatch_size))
print('')
print('  model_save_interval : {}'.format(model_save_interval))
print('  model_save_dir      : {}'.format(model_save_dir))
print('  save_model          : {}'.format(save_model))
print('')
print('  tensorboard_update_freq : {}'.format(tensorboard_update_freq))
print('')
<mask token>


def createTargetNetworkUpdateOperations(tf_trainable_vars):
    """Create a list where each element is a tensorflow assign operation which overwrites
       trainable variables of 'target' with 'moving'.

    First HALF of all variables correpsond to the 'moving' Q-network.
    Later HALF are the 'target' Q-network.
    [ ... moving ... | ... target ... ]

    Note: The Q-network objects must be initialized in THIS ORDER !!!

    Args:
      tf_trainable_vars (tf var list) : List containing all trainable variables.
    Returns:
      op_list (list) : List containing tf operations to copy 'moving' -> 'target'.
    """
    tot_vars = len(tf_trainable_vars)
    op_list = []
    for i, moving_var in enumerate(tf_trainable_vars[0:tot_vars // 2]):
        op_list.append(tf_trainable_vars[i + tot_vars // 2].assign(
            moving_var.value()))
    return op_list


def updateTargetNetwork(op_list, sess):
    """Run previously defined tf operations to overwrite 'target' Q-network with 'moving' Q-network."""
    for op in op_list:
        sess.run(op)


def saveModel(ep, save_dir, tf_saver, sess, save_model=False):
    """Create a save model folder within the specified save directory.

    Args:
      ep (int)             : Episode number where the model is saved.
      save_dir (str)       : Path to the directory where the model will be saved
      tf_saver (tf object) : 
      sess : (tf object)   : 
      save_model : (bool)  : 
    """
    if save_model == False:
        return
    if os.path.isdir(save_dir) == False:
        os.mkdir(save_dir)
    save_folder_path = os.path.join(save_dir, 'save_{}'.format(ep))
    os.mkdir(save_folder_path)
    file_name = 'save_{}.ckpt'.format(ep)
    save_path = os.path.join(save_folder_path, file_name)
    print('Saving model : {}'.format(save_path))
    tf_saver.save(sess, save_path)


class QNetwork:
    """Class which represent the network which outputs a Q-value based on an input state.
    """

    def __init__(self, learning_rate, img_w, img_h, frame_num, action_size,
        hidden_size, name='QNetwork'):
        """Initialize all tensorflow operation objects when initializing Q-network object.

        Use the CNN function approximation to generate and store Q-value arrays.
        These Q-value arrays are then inputted into 'Q_input' placeholder in order to compute other things.
        """
        with tf.variable_scope(name):
            self.actions_ = tf.placeholder(tf.int32, [None], name='actions')
            one_hot_actions = tf.one_hot(self.actions_, action_size)
            self.targetQs_ = tf.placeholder(tf.float32, [None], name='target')
            self.inputs = tf.placeholder(tf.float32, [None, img_w, img_h,
                frame_num], name='inputs')
            self.conv1 = slim.conv2d(inputs=self.inputs, num_outputs=32,
                kernel_size=[8, 8], stride=[4, 4], padding='VALID',
                biases_initializer=None)
            self.conv2 = slim.conv2d(inputs=self.conv1, num_outputs=64,
                kernel_size=[4, 4], stride=[2, 2], padding='VALID',
                biases_initializer=None)
            self.conv3 = slim.conv2d(inputs=self.conv2, num_outputs=64,
                kernel_size=[3, 3], stride=[1, 1], padding='VALID',
                biases_initializer=None)
            self.flatten = slim.flatten(self.conv3)
            self.fc4 = slim.fully_connected(self.flatten, hidden_size,
                activation_fn=tf.nn.relu, biases_initializer=None)
            self.output = slim.fully_connected(self.fc4, action_size,
                activation_fn=None, biases_initializer=None)
            self.Q = tf.reduce_sum(tf.multiply(self.output, one_hot_actions
                ), axis=1)
            self.loss = tf.reduce_mean(tf.square(self.targetQs_ - self.Q))
            tf.summary.scalar('loss', self.loss)
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss
                )
            self.merged = tf.summary.merge_all()


<mask token>
buff_f.reset(env)
buff_f_next.reset(env)
for pretrain_step in range(pretrain_step_max):
    a = env.action_space.sample()
    f_next, r, done, _ = env.step(a)
    f_next = preprocessImage(f_next)
    exp = f_next, a, r, done
    ExpBuffer.store(exp)
    if done:
        env.reset()
<mask token>
with tf.Session() as sess:
    sess.run(tf_init)
    updateTargetNetwork(QNet_target_update_op_list, sess)
    tb_train_writer = tf.summary.FileWriter('./logs/1/train', sess.graph)
    step_tot = 0
    for ep in range(1, ep_max):
        buff_f.reset(env)
        buff_f_next.reset(env)
        env.reset()
        done = False
        step = 0
        ep_rew = 0
        while step < step_max:
            step_tot += 1
            explore_p = explore_stop + (explore_start - explore_stop) * np.exp(
                -explore_decay_rate * step_tot)
            if explore_p > np.random.rand():
                a = env.action_space.sample()
            else:
                Qs = sess.run(QNet_moving.output, feed_dict={QNet_moving.
                    inputs: np.reshape(buff_f.sample(), (1, 84, 84, 4))})
                a = np.argmax(Qs)
            r_repeat_tot = 0
            for _ in range(action_repeat):
                f_next, r_repeat, done, _ = env.step(a)
                r_repeat_tot += r_repeat
                if done == True:
                    break
            r = r_repeat_tot
            f_next = preprocessImage(f_next)
            buff_f_next.store(f_next)
            exp = f_next, a, r, done
            ExpBuffer.store(exp)
            (s_minibatch, a_minibatch, r_minibatch, s_next_minibatch,
                done_minibatch) = ExpBuffer.sample(minibatch_size)
            ep_termination = 1 - done_minibatch
            target_Q_array = sess.run(QNet_target.output, feed_dict={
                QNet_target.inputs: s_next_minibatch})
            targets = r_minibatch + gamma * np.max(target_Q_array, axis=1
                ) * ep_termination
            _ = sess.run(QNet_moving.opt, feed_dict={QNet_moving.inputs:
                s_minibatch, QNet_moving.targetQs_: targets, QNet_moving.
                actions_: a_minibatch})
            if step_tot % QNet_update_freq == 0:
                updateTargetNetwork(QNet_target_update_op_list, sess)
            if step_tot % tensorboard_update_freq == 0:
                tb_summary = sess.run(QNet_moving.merged, feed_dict={
                    QNet_moving.inputs: s_minibatch, QNet_moving.targetQs_:
                    targets, QNet_moving.actions_: a_minibatch})
                tb_train_writer.add_summary(tb_summary, step_tot)
            f = f_next
            buff_f.store(f)
            step += 1
            ep_rew += r
            if done == True:
                print('Episode: {}'.format(ep), 'Steps : {}'.format(
                    step_tot), 'Total reward: {}'.format(ep_rew),
                    'Explore P: {:.4f}'.format(explore_p))
                rewards_list.append((ep, ep_rew, step_tot))
                break
        if ep % model_save_interval == 0:
            saveModel(ep, model_save_dir, tf_saver, sess, save_model)
<mask token>
print(out_array.shape)
np.savetxt('rew_new.dat', out_array, fmt='%d')
env.close()

<mask token>
tf.reset_default_graph()
env_name = 'Pong-v0'
ep_max = int(1000000.0)
step_tot_max = int(200000000.0)
step_max = int(1000000.0)
gamma = 0.99
learning_rate = float(0.0001)
seed = 42
action_repeat = 4
explore_start = 1.0
explore_stop = 0.1
explore_steps = int(1000000.0)
explore_decay_rate = (explore_start - explore_stop) / explore_steps
frame_buffer_size = 4
hidden_size = 512
QNet_update_freq = int(10000.0)
exp_buffer_size = int(1000000.0)
minibatch_size = 32
pretrain_step_max = minibatch_size + frame_buffer_size
model_save_interval = 100
model_save_dir = './saves'
save_model = False
tensorboard_update_freq = 50
env = gym.make(env_name)
env_actions = env.action_space.n


def setAllRandomSeeds(environment, seed=42):
    random.seed(seed)
    np.random.seed(seed)
    tf.set_random_seed(seed)
    environment.seed(seed)


setAllRandomSeeds(env, seed)
print("""
*********************""")
print('*  Run information  *')
print('*********************')
print('  Environment name   : {}'.format(env_name))
print('  Action space       : {}'.format(env.action_space))
print('  Observation space  : {}'.format(env.observation_space))
print('  ep_max             : {}'.format(ep_max))
print('  step_tot_max       : {}'.format(step_tot_max))
print('  step_max           : {}'.format(step_max))
print('  gamma              : {}'.format(gamma))
print('  learning_rate      : {}'.format(learning_rate))
print('  seed               : {}'.format(seed))
print('  action_repeat      : {}'.format(action_repeat))
print('  explore_start      : {}'.format(explore_start))
print('  explore_stop       : {}'.format(explore_stop))
print('  explore_steps      : {}'.format(explore_steps))
print('  explore_decay_rate : {}'.format(explore_decay_rate))
print('  frame_buffer_size  : {}'.format(frame_buffer_size))
print('  QNet_update_freq   : {}'.format(QNet_update_freq))
print('  pretrain_step_max  : {}'.format(pretrain_step_max))
print('  exp_buffer_size    : {}'.format(exp_buffer_size))
print('  minibatch_size     : {}'.format(minibatch_size))
print('')
print('  model_save_interval : {}'.format(model_save_interval))
print('  model_save_dir      : {}'.format(model_save_dir))
print('  save_model          : {}'.format(save_model))
print('')
print('  tensorboard_update_freq : {}'.format(tensorboard_update_freq))
print('')
ExpBuffer = ExpBuffer(max_size=exp_buffer_size)
buff_f = FrameBuffer(size=frame_buffer_size)
buff_f_next = FrameBuffer(size=frame_buffer_size)


def createTargetNetworkUpdateOperations(tf_trainable_vars):
    """Create a list where each element is a tensorflow assign operation which overwrites
       trainable variables of 'target' with 'moving'.

    First HALF of all variables correpsond to the 'moving' Q-network.
    Later HALF are the 'target' Q-network.
    [ ... moving ... | ... target ... ]

    Note: The Q-network objects must be initialized in THIS ORDER !!!

    Args:
      tf_trainable_vars (tf var list) : List containing all trainable variables.
    Returns:
      op_list (list) : List containing tf operations to copy 'moving' -> 'target'.
    """
    tot_vars = len(tf_trainable_vars)
    op_list = []
    for i, moving_var in enumerate(tf_trainable_vars[0:tot_vars // 2]):
        op_list.append(tf_trainable_vars[i + tot_vars // 2].assign(
            moving_var.value()))
    return op_list


def updateTargetNetwork(op_list, sess):
    """Run previously defined tf operations to overwrite 'target' Q-network with 'moving' Q-network."""
    for op in op_list:
        sess.run(op)


def saveModel(ep, save_dir, tf_saver, sess, save_model=False):
    """Create a save model folder within the specified save directory.

    Args:
      ep (int)             : Episode number where the model is saved.
      save_dir (str)       : Path to the directory where the model will be saved
      tf_saver (tf object) : 
      sess : (tf object)   : 
      save_model : (bool)  : 
    """
    if save_model == False:
        return
    if os.path.isdir(save_dir) == False:
        os.mkdir(save_dir)
    save_folder_path = os.path.join(save_dir, 'save_{}'.format(ep))
    os.mkdir(save_folder_path)
    file_name = 'save_{}.ckpt'.format(ep)
    save_path = os.path.join(save_folder_path, file_name)
    print('Saving model : {}'.format(save_path))
    tf_saver.save(sess, save_path)


class QNetwork:
    """Class which represent the network which outputs a Q-value based on an input state.
    """

    def __init__(self, learning_rate, img_w, img_h, frame_num, action_size,
        hidden_size, name='QNetwork'):
        """Initialize all tensorflow operation objects when initializing Q-network object.

        Use the CNN function approximation to generate and store Q-value arrays.
        These Q-value arrays are then inputted into 'Q_input' placeholder in order to compute other things.
        """
        with tf.variable_scope(name):
            self.actions_ = tf.placeholder(tf.int32, [None], name='actions')
            one_hot_actions = tf.one_hot(self.actions_, action_size)
            self.targetQs_ = tf.placeholder(tf.float32, [None], name='target')
            self.inputs = tf.placeholder(tf.float32, [None, img_w, img_h,
                frame_num], name='inputs')
            self.conv1 = slim.conv2d(inputs=self.inputs, num_outputs=32,
                kernel_size=[8, 8], stride=[4, 4], padding='VALID',
                biases_initializer=None)
            self.conv2 = slim.conv2d(inputs=self.conv1, num_outputs=64,
                kernel_size=[4, 4], stride=[2, 2], padding='VALID',
                biases_initializer=None)
            self.conv3 = slim.conv2d(inputs=self.conv2, num_outputs=64,
                kernel_size=[3, 3], stride=[1, 1], padding='VALID',
                biases_initializer=None)
            self.flatten = slim.flatten(self.conv3)
            self.fc4 = slim.fully_connected(self.flatten, hidden_size,
                activation_fn=tf.nn.relu, biases_initializer=None)
            self.output = slim.fully_connected(self.fc4, action_size,
                activation_fn=None, biases_initializer=None)
            self.Q = tf.reduce_sum(tf.multiply(self.output, one_hot_actions
                ), axis=1)
            self.loss = tf.reduce_mean(tf.square(self.targetQs_ - self.Q))
            tf.summary.scalar('loss', self.loss)
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss
                )
            self.merged = tf.summary.merge_all()


QNet_moving = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size,
    name='QNet_moving')
QNet_target = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size,
    name='QNet_target')
tf_init = tf.global_variables_initializer()
tf_saver = tf.train.Saver()
tf_trainable_vars = tf.trainable_variables()
QNet_target_update_op_list = createTargetNetworkUpdateOperations(
    tf_trainable_vars)
buff_f.reset(env)
buff_f_next.reset(env)
for pretrain_step in range(pretrain_step_max):
    a = env.action_space.sample()
    f_next, r, done, _ = env.step(a)
    f_next = preprocessImage(f_next)
    exp = f_next, a, r, done
    ExpBuffer.store(exp)
    if done:
        env.reset()
saver = tf.train.Saver()
rewards_list = []
with tf.Session() as sess:
    sess.run(tf_init)
    updateTargetNetwork(QNet_target_update_op_list, sess)
    tb_train_writer = tf.summary.FileWriter('./logs/1/train', sess.graph)
    step_tot = 0
    for ep in range(1, ep_max):
        buff_f.reset(env)
        buff_f_next.reset(env)
        env.reset()
        done = False
        step = 0
        ep_rew = 0
        while step < step_max:
            step_tot += 1
            explore_p = explore_stop + (explore_start - explore_stop) * np.exp(
                -explore_decay_rate * step_tot)
            if explore_p > np.random.rand():
                a = env.action_space.sample()
            else:
                Qs = sess.run(QNet_moving.output, feed_dict={QNet_moving.
                    inputs: np.reshape(buff_f.sample(), (1, 84, 84, 4))})
                a = np.argmax(Qs)
            r_repeat_tot = 0
            for _ in range(action_repeat):
                f_next, r_repeat, done, _ = env.step(a)
                r_repeat_tot += r_repeat
                if done == True:
                    break
            r = r_repeat_tot
            f_next = preprocessImage(f_next)
            buff_f_next.store(f_next)
            exp = f_next, a, r, done
            ExpBuffer.store(exp)
            (s_minibatch, a_minibatch, r_minibatch, s_next_minibatch,
                done_minibatch) = ExpBuffer.sample(minibatch_size)
            ep_termination = 1 - done_minibatch
            target_Q_array = sess.run(QNet_target.output, feed_dict={
                QNet_target.inputs: s_next_minibatch})
            targets = r_minibatch + gamma * np.max(target_Q_array, axis=1
                ) * ep_termination
            _ = sess.run(QNet_moving.opt, feed_dict={QNet_moving.inputs:
                s_minibatch, QNet_moving.targetQs_: targets, QNet_moving.
                actions_: a_minibatch})
            if step_tot % QNet_update_freq == 0:
                updateTargetNetwork(QNet_target_update_op_list, sess)
            if step_tot % tensorboard_update_freq == 0:
                tb_summary = sess.run(QNet_moving.merged, feed_dict={
                    QNet_moving.inputs: s_minibatch, QNet_moving.targetQs_:
                    targets, QNet_moving.actions_: a_minibatch})
                tb_train_writer.add_summary(tb_summary, step_tot)
            f = f_next
            buff_f.store(f)
            step += 1
            ep_rew += r
            if done == True:
                print('Episode: {}'.format(ep), 'Steps : {}'.format(
                    step_tot), 'Total reward: {}'.format(ep_rew),
                    'Explore P: {:.4f}'.format(explore_p))
                rewards_list.append((ep, ep_rew, step_tot))
                break
        if ep % model_save_interval == 0:
            saveModel(ep, model_save_dir, tf_saver, sess, save_model)
eps, rews, steps = np.array(rewards_list).T
eps = np.reshape(eps, [-1, 1])
rews = np.reshape(rews, [-1, 1])
steps = np.reshape(steps, [-1, 1])
out_array = np.concatenate((steps, rews), axis=1)
print(out_array.shape)
np.savetxt('rew_new.dat', out_array, fmt='%d')
env.close()

import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
import os
import random
import gym
import matplotlib.pyplot as plt
from experience_replay_buffer import ExpBuffer
from utils import preprocessImage, FrameBuffer
tf.reset_default_graph()
env_name = 'Pong-v0'
ep_max = int(1000000.0)
step_tot_max = int(200000000.0)
step_max = int(1000000.0)
gamma = 0.99
learning_rate = float(0.0001)
seed = 42
action_repeat = 4
explore_start = 1.0
explore_stop = 0.1
explore_steps = int(1000000.0)
explore_decay_rate = (explore_start - explore_stop) / explore_steps
frame_buffer_size = 4
hidden_size = 512
QNet_update_freq = int(10000.0)
exp_buffer_size = int(1000000.0)
minibatch_size = 32
pretrain_step_max = minibatch_size + frame_buffer_size
model_save_interval = 100
model_save_dir = './saves'
save_model = False
tensorboard_update_freq = 50
env = gym.make(env_name)
env_actions = env.action_space.n


def setAllRandomSeeds(environment, seed=42):
    random.seed(seed)
    np.random.seed(seed)
    tf.set_random_seed(seed)
    environment.seed(seed)


setAllRandomSeeds(env, seed)
print("""
*********************""")
print('*  Run information  *')
print('*********************')
print('  Environment name   : {}'.format(env_name))
print('  Action space       : {}'.format(env.action_space))
print('  Observation space  : {}'.format(env.observation_space))
print('  ep_max             : {}'.format(ep_max))
print('  step_tot_max       : {}'.format(step_tot_max))
print('  step_max           : {}'.format(step_max))
print('  gamma              : {}'.format(gamma))
print('  learning_rate      : {}'.format(learning_rate))
print('  seed               : {}'.format(seed))
print('  action_repeat      : {}'.format(action_repeat))
print('  explore_start      : {}'.format(explore_start))
print('  explore_stop       : {}'.format(explore_stop))
print('  explore_steps      : {}'.format(explore_steps))
print('  explore_decay_rate : {}'.format(explore_decay_rate))
print('  frame_buffer_size  : {}'.format(frame_buffer_size))
print('  QNet_update_freq   : {}'.format(QNet_update_freq))
print('  pretrain_step_max  : {}'.format(pretrain_step_max))
print('  exp_buffer_size    : {}'.format(exp_buffer_size))
print('  minibatch_size     : {}'.format(minibatch_size))
print('')
print('  model_save_interval : {}'.format(model_save_interval))
print('  model_save_dir      : {}'.format(model_save_dir))
print('  save_model          : {}'.format(save_model))
print('')
print('  tensorboard_update_freq : {}'.format(tensorboard_update_freq))
print('')
ExpBuffer = ExpBuffer(max_size=exp_buffer_size)
buff_f = FrameBuffer(size=frame_buffer_size)
buff_f_next = FrameBuffer(size=frame_buffer_size)


def createTargetNetworkUpdateOperations(tf_trainable_vars):
    """Create a list where each element is a tensorflow assign operation which overwrites
       trainable variables of 'target' with 'moving'.

    First HALF of all variables correpsond to the 'moving' Q-network.
    Later HALF are the 'target' Q-network.
    [ ... moving ... | ... target ... ]

    Note: The Q-network objects must be initialized in THIS ORDER !!!

    Args:
      tf_trainable_vars (tf var list) : List containing all trainable variables.
    Returns:
      op_list (list) : List containing tf operations to copy 'moving' -> 'target'.
    """
    tot_vars = len(tf_trainable_vars)
    op_list = []
    for i, moving_var in enumerate(tf_trainable_vars[0:tot_vars // 2]):
        op_list.append(tf_trainable_vars[i + tot_vars // 2].assign(
            moving_var.value()))
    return op_list


def updateTargetNetwork(op_list, sess):
    """Run previously defined tf operations to overwrite 'target' Q-network with 'moving' Q-network."""
    for op in op_list:
        sess.run(op)


def saveModel(ep, save_dir, tf_saver, sess, save_model=False):
    """Create a save model folder within the specified save directory.

    Args:
      ep (int)             : Episode number where the model is saved.
      save_dir (str)       : Path to the directory where the model will be saved
      tf_saver (tf object) : 
      sess : (tf object)   : 
      save_model : (bool)  : 
    """
    if save_model == False:
        return
    if os.path.isdir(save_dir) == False:
        os.mkdir(save_dir)
    save_folder_path = os.path.join(save_dir, 'save_{}'.format(ep))
    os.mkdir(save_folder_path)
    file_name = 'save_{}.ckpt'.format(ep)
    save_path = os.path.join(save_folder_path, file_name)
    print('Saving model : {}'.format(save_path))
    tf_saver.save(sess, save_path)


class QNetwork:
    """Class which represent the network which outputs a Q-value based on an input state.
    """

    def __init__(self, learning_rate, img_w, img_h, frame_num, action_size,
        hidden_size, name='QNetwork'):
        """Initialize all tensorflow operation objects when initializing Q-network object.

        Use the CNN function approximation to generate and store Q-value arrays.
        These Q-value arrays are then inputted into 'Q_input' placeholder in order to compute other things.
        """
        with tf.variable_scope(name):
            self.actions_ = tf.placeholder(tf.int32, [None], name='actions')
            one_hot_actions = tf.one_hot(self.actions_, action_size)
            self.targetQs_ = tf.placeholder(tf.float32, [None], name='target')
            self.inputs = tf.placeholder(tf.float32, [None, img_w, img_h,
                frame_num], name='inputs')
            self.conv1 = slim.conv2d(inputs=self.inputs, num_outputs=32,
                kernel_size=[8, 8], stride=[4, 4], padding='VALID',
                biases_initializer=None)
            self.conv2 = slim.conv2d(inputs=self.conv1, num_outputs=64,
                kernel_size=[4, 4], stride=[2, 2], padding='VALID',
                biases_initializer=None)
            self.conv3 = slim.conv2d(inputs=self.conv2, num_outputs=64,
                kernel_size=[3, 3], stride=[1, 1], padding='VALID',
                biases_initializer=None)
            self.flatten = slim.flatten(self.conv3)
            self.fc4 = slim.fully_connected(self.flatten, hidden_size,
                activation_fn=tf.nn.relu, biases_initializer=None)
            self.output = slim.fully_connected(self.fc4, action_size,
                activation_fn=None, biases_initializer=None)
            self.Q = tf.reduce_sum(tf.multiply(self.output, one_hot_actions
                ), axis=1)
            self.loss = tf.reduce_mean(tf.square(self.targetQs_ - self.Q))
            tf.summary.scalar('loss', self.loss)
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss
                )
            self.merged = tf.summary.merge_all()


QNet_moving = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size,
    name='QNet_moving')
QNet_target = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size,
    name='QNet_target')
tf_init = tf.global_variables_initializer()
tf_saver = tf.train.Saver()
tf_trainable_vars = tf.trainable_variables()
QNet_target_update_op_list = createTargetNetworkUpdateOperations(
    tf_trainable_vars)
buff_f.reset(env)
buff_f_next.reset(env)
for pretrain_step in range(pretrain_step_max):
    a = env.action_space.sample()
    f_next, r, done, _ = env.step(a)
    f_next = preprocessImage(f_next)
    exp = f_next, a, r, done
    ExpBuffer.store(exp)
    if done:
        env.reset()
saver = tf.train.Saver()
rewards_list = []
with tf.Session() as sess:
    sess.run(tf_init)
    updateTargetNetwork(QNet_target_update_op_list, sess)
    tb_train_writer = tf.summary.FileWriter('./logs/1/train', sess.graph)
    step_tot = 0
    for ep in range(1, ep_max):
        buff_f.reset(env)
        buff_f_next.reset(env)
        env.reset()
        done = False
        step = 0
        ep_rew = 0
        while step < step_max:
            step_tot += 1
            explore_p = explore_stop + (explore_start - explore_stop) * np.exp(
                -explore_decay_rate * step_tot)
            if explore_p > np.random.rand():
                a = env.action_space.sample()
            else:
                Qs = sess.run(QNet_moving.output, feed_dict={QNet_moving.
                    inputs: np.reshape(buff_f.sample(), (1, 84, 84, 4))})
                a = np.argmax(Qs)
            r_repeat_tot = 0
            for _ in range(action_repeat):
                f_next, r_repeat, done, _ = env.step(a)
                r_repeat_tot += r_repeat
                if done == True:
                    break
            r = r_repeat_tot
            f_next = preprocessImage(f_next)
            buff_f_next.store(f_next)
            exp = f_next, a, r, done
            ExpBuffer.store(exp)
            (s_minibatch, a_minibatch, r_minibatch, s_next_minibatch,
                done_minibatch) = ExpBuffer.sample(minibatch_size)
            ep_termination = 1 - done_minibatch
            target_Q_array = sess.run(QNet_target.output, feed_dict={
                QNet_target.inputs: s_next_minibatch})
            targets = r_minibatch + gamma * np.max(target_Q_array, axis=1
                ) * ep_termination
            _ = sess.run(QNet_moving.opt, feed_dict={QNet_moving.inputs:
                s_minibatch, QNet_moving.targetQs_: targets, QNet_moving.
                actions_: a_minibatch})
            if step_tot % QNet_update_freq == 0:
                updateTargetNetwork(QNet_target_update_op_list, sess)
            if step_tot % tensorboard_update_freq == 0:
                tb_summary = sess.run(QNet_moving.merged, feed_dict={
                    QNet_moving.inputs: s_minibatch, QNet_moving.targetQs_:
                    targets, QNet_moving.actions_: a_minibatch})
                tb_train_writer.add_summary(tb_summary, step_tot)
            f = f_next
            buff_f.store(f)
            step += 1
            ep_rew += r
            if done == True:
                print('Episode: {}'.format(ep), 'Steps : {}'.format(
                    step_tot), 'Total reward: {}'.format(ep_rew),
                    'Explore P: {:.4f}'.format(explore_p))
                rewards_list.append((ep, ep_rew, step_tot))
                break
        if ep % model_save_interval == 0:
            saveModel(ep, model_save_dir, tf_saver, sess, save_model)
eps, rews, steps = np.array(rewards_list).T
eps = np.reshape(eps, [-1, 1])
rews = np.reshape(rews, [-1, 1])
steps = np.reshape(steps, [-1, 1])
out_array = np.concatenate((steps, rews), axis=1)
print(out_array.shape)
np.savetxt('rew_new.dat', out_array, fmt='%d')
env.close()

import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
import os
import random
import gym

import matplotlib.pyplot as plt


from experience_replay_buffer import ExpBuffer
from utils import preprocessImage, FrameBuffer



tf.reset_default_graph()

#####################
#  HYPERPARAMETERS  #
#####################
env_name = "Pong-v0"
ep_max = int(1e6)          # max number of episodes to learn from
step_tot_max = int(200e6)
step_max = int(1e6)                # max steps in an episode
gamma = 0.99                   # future reward discount
learning_rate = float(1e-4)         # Q-network learning rate
seed = 42
action_repeat = 4
# Exploration parameters
explore_start = 1.0            # exploration probability at start
explore_stop = 0.1            # minimum exploration probability
explore_steps = int(1e6)
explore_decay_rate = (explore_start - explore_stop) / explore_steps            # exponential decay rate for exploration prob
# Network parameters
frame_buffer_size = 4
hidden_size = 512               # number of units in each Q-network hidden layer
QNet_update_freq = int(1e4)
# Experience buffer parameters
exp_buffer_size = int(1e6)            # ExpBuffer capacity
minibatch_size = 32                # experience mini-batch size
pretrain_step_max = minibatch_size+frame_buffer_size   # number experiences to pretrain the ExpBuffer
# Save parameters
model_save_interval = 100
model_save_dir = "./saves"
save_model = False
# TensorBoard
tensorboard_update_freq = 50


# Initialize environment
env = gym.make(env_name)
env_actions = env.action_space.n
# Set random seeds
def setAllRandomSeeds(environment, seed=42):
    random.seed(seed)
    np.random.seed(seed)
    tf.set_random_seed(seed)
    environment.seed(seed)
setAllRandomSeeds(env, seed)


print("\n*********************")
print("*  Run information  *")
print("*********************")
print("  Environment name   : {}".format(env_name))
print("  Action space       : {}".format(env.action_space))
print("  Observation space  : {}".format(env.observation_space))
print("  ep_max             : {}".format(ep_max))
print("  step_tot_max       : {}".format(step_tot_max))
print("  step_max           : {}".format(step_max))
print("  gamma              : {}".format(gamma))
print("  learning_rate      : {}".format(learning_rate))
print("  seed               : {}".format(seed))
print("  action_repeat      : {}".format(action_repeat))
print("  explore_start      : {}".format(explore_start))
print("  explore_stop       : {}".format(explore_stop))
print("  explore_steps      : {}".format(explore_steps))
print("  explore_decay_rate : {}".format(explore_decay_rate))
print("  frame_buffer_size  : {}".format(frame_buffer_size))
print("  QNet_update_freq   : {}".format(QNet_update_freq))
print("  pretrain_step_max  : {}".format(pretrain_step_max))
print("  exp_buffer_size    : {}".format(exp_buffer_size))
print("  minibatch_size     : {}".format(minibatch_size))
print("")
print("  model_save_interval : {}".format(model_save_interval))
print("  model_save_dir      : {}".format(model_save_dir))
print("  save_model          : {}".format(save_model))
print("")
print("  tensorboard_update_freq : {}".format(tensorboard_update_freq))
print("")


# Initialize experience replay buffer object
ExpBuffer = ExpBuffer(max_size=exp_buffer_size)

# Initialize frame buffer objects for stacking frames when doing inference (i.e. predicting Q-value for 'a')
buff_f = FrameBuffer(size=frame_buffer_size)
buff_f_next = FrameBuffer(size=frame_buffer_size)


#########################
#  COMPUTATIONAL GRAPH  #
#########################\
def createTargetNetworkUpdateOperations(tf_trainable_vars):
    '''Create a list where each element is a tensorflow assign operation which overwrites
       trainable variables of 'target' with 'moving'.

    First HALF of all variables correpsond to the 'moving' Q-network.
    Later HALF are the 'target' Q-network.
    [ ... moving ... | ... target ... ]

    Note: The Q-network objects must be initialized in THIS ORDER !!!

    Args:
      tf_trainable_vars (tf var list) : List containing all trainable variables.
    Returns:
      op_list (list) : List containing tf operations to copy 'moving' -> 'target'.
    '''
    tot_vars = len(tf_trainable_vars)
    op_list = []
    # Get the 'moving' Q-network variables from the first half ('//'' is integer division)
    for i, moving_var in enumerate(tf_trainable_vars[0:tot_vars//2]):
        # Define a tf operation which copies the 'moving' variable into 'target' variable
        op_list.append( tf_trainable_vars[i+tot_vars//2].assign(moving_var.value()) )

    return op_list


def updateTargetNetwork(op_list, sess):
    '''Run previously defined tf operations to overwrite 'target' Q-network with 'moving' Q-network.'''
    for op in op_list:
        sess.run(op)


def saveModel(ep, save_dir, tf_saver, sess, save_model=False):
    '''Create a save model folder within the specified save directory.

    Args:
      ep (int)             : Episode number where the model is saved.
      save_dir (str)       : Path to the directory where the model will be saved
      tf_saver (tf object) : 
      sess : (tf object)   : 
      save_model : (bool)  : 
    '''
    if(save_model == False):
        return
    # Create the save directory if it does not already exist
    if(os.path.isdir(save_dir) == False):
        os.mkdir(save_dir)
    # Create save folder
    save_folder_path = os.path.join(save_dir, "save_{}".format(ep))
    os.mkdir(save_folder_path)
    # Define save name
    file_name = "save_{}.ckpt".format(ep)
    save_path = os.path.join(save_folder_path, file_name)

    print("Saving model : {}".format(save_path))

    tf_saver.save(sess, save_path)


class QNetwork:
    '''Class which represent the network which outputs a Q-value based on an input state.
    '''
    def __init__(self, learning_rate, img_w, img_h, frame_num, action_size, hidden_size, name='QNetwork'):
        '''Initialize all tensorflow operation objects when initializing Q-network object.

        Use the CNN function approximation to generate and store Q-value arrays.
        These Q-value arrays are then inputted into 'Q_input' placeholder in order to compute other things.
        '''
        # state inputs to the Q-network
        with tf.variable_scope(name):

            ##################
            #  PLACEHOLDERS  #
            ##################
            # Array with an integer for every sample
            self.actions_ = tf.placeholder(tf.int32, [None], name='actions')
            # One hot encode the actions to later choose the Q-value for the action
            one_hot_actions = tf.one_hot(self.actions_, action_size)
            # Target Q values for training
            self.targetQs_ = tf.placeholder(tf.float32, [None], name='target')

            ###########
            #  INPUT  #
            ###########
            self.inputs = tf.placeholder(tf.float32, [None, img_w, img_h, frame_num], name='inputs')
            
            ############
            #  LAYERS  #
            ############

            self.conv1 = slim.conv2d(inputs=self.inputs,
                                     num_outputs=32,
                                     kernel_size=[8,8],
                                     stride=[4,4],
                                     padding="VALID",
                                     biases_initializer=None)

            self.conv2 = slim.conv2d(inputs=self.conv1,
                                     num_outputs=64,
                                     kernel_size=[4,4],
                                     stride=[2,2],
                                     padding="VALID",
                                     biases_initializer=None)

            self.conv3 = slim.conv2d(inputs=self.conv2,
                                     num_outputs=64,
                                     kernel_size=[3,3],
                                     stride=[1,1],
                                     padding="VALID",
                                     biases_initializer=None)

            self.flatten = slim.flatten(self.conv3)

            self.fc4 = slim.fully_connected(self.flatten, hidden_size, activation_fn=tf.nn.relu, biases_initializer=None)

            self.output = slim.fully_connected(self.fc4, action_size, activation_fn=None, biases_initializer=None)

            # ReLU hidden layers
            #self.fc1 = tf.contrib.layers.fully_connected(self.inputs, hidden_size)
            #self.fc2 = tf.contrib.layers.fully_connected(self.fc1, hidden_size)
            # Linear output layer
            #self.output = tf.contrib.layers.fully_connected(self.fc2, action_size, activation_fn=None)
            
            ##############
            #  TRAINING  #
            ##############
            # Value of selected Q-value          
            self.Q = tf.reduce_sum(tf.multiply(self.output, one_hot_actions), axis=1)
            # L2 loss (targetQ - Q)^2
            self.loss = tf.reduce_mean(tf.square(self.targetQs_ - self.Q))
            tf.summary.scalar("loss", self.loss)
            # Training operations
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss)
            # TensorBoard operation
            self.merged = tf.summary.merge_all()

QNet_moving = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size, name='QNet_moving')
QNet_target = QNetwork(learning_rate, 84, 84, 4, env_actions, hidden_size, name='QNet_target')


tf_init = tf.global_variables_initializer()
tf_saver = tf.train.Saver()
tf_trainable_vars = tf.trainable_variables()
# List with operations to overwrite the 'target' Q-network variables
QNet_target_update_op_list = createTargetNetworkUpdateOperations(tf_trainable_vars)


############################
#  INITIALIZE ENVIRONMENT  #
############################
# Fill frame buffers
buff_f.reset(env)
buff_f_next.reset(env)


####################
#  PRETRAIN STEPS  #
####################
# Fill replay buffer with randomly selected experiences before doing training steps
for pretrain_step in range(pretrain_step_max):
    # Make a random action
    a = env.action_space.sample()
    f_next, r, done, _ = env.step(a)
    f_next = preprocessImage(f_next)
    # Store experience to buffer
    exp = (f_next, a, r, done)
    ExpBuffer.store(exp)
    if done:      
        # Start new episode
        env.reset()


###############
#  MAIN LOOP  #
###############
saver = tf.train.Saver()
rewards_list = []
with tf.Session() as sess:
    # Initialize variables
    sess.run(tf_init)
    updateTargetNetwork(QNet_target_update_op_list, sess)
    # TensorBoard storage directory
    tb_train_writer = tf.summary.FileWriter('./logs/1/train', sess.graph)

    step_tot = 0

    for ep in range(1, ep_max):

        # Reset buffer at start of episode
        buff_f.reset(env)
        buff_f_next.reset(env)

        # Reset environment at start of episode
        env.reset()
        done = False
        step = 0
        # Store total reward for the episode
        ep_rew = 0    

        while step < step_max:

            step_tot += 1
            # env.render() 
            
            ###################
            #  SELECT ACTION  #
            ###################
            # Reduce exploration
            explore_p = explore_stop + (explore_start - explore_stop)*np.exp(-explore_decay_rate*step_tot) 
            if explore_p > np.random.rand():
                # Select exploratory random action
                a = env.action_space.sample()
            else:
                # Select optimal action according to policy (i.e. Q-network)
                Qs = sess.run(QNet_moving.output, feed_dict={QNet_moving.inputs: np.reshape(buff_f.sample(), (1,84,84,4))})
                a = np.argmax(Qs)
            
            #################
            #  TAKE ACTION  #
            #################
            # Do 'action repeat' -> Train only on every 4th frame
            r_repeat_tot = 0
            for _ in range(action_repeat):
                f_next, r_repeat, done, _ = env.step(a)
                r_repeat_tot += r_repeat
                if(done == True):
                    break
            r = r_repeat_tot
            
            f_next = preprocessImage(f_next)
            buff_f_next.store(f_next)

            ######################
            #  STORE EXPERIENCE  #
            ######################
            exp = (f_next, a, r, done)
            ExpBuffer.store(exp)
            
            ##############
            #  TRAINING  #
            ##############
            # Sample one minibatch from buffer into separate arrays
            # Each array index correspond to the same experience
            s_minibatch, a_minibatch, r_minibatch, s_next_minibatch, done_minibatch = ExpBuffer.sample(minibatch_size)
            
            # Compute 'target term'
            ep_termination = 1 - done_minibatch
            target_Q_array = sess.run(QNet_target.output, feed_dict={QNet_target.inputs: s_next_minibatch})
            targets = r_minibatch + gamma * np.max(target_Q_array, axis=1)*ep_termination

            # Update moving Q-network
            _ = sess.run(QNet_moving.opt, feed_dict={QNet_moving.inputs: s_minibatch,
                                                     QNet_moving.targetQs_: targets,
                                                     QNet_moving.actions_: a_minibatch})

            ###########################
            #  UPDATE TARGET NETWORK  #
            ###########################
            if(step_tot % QNet_update_freq == 0):
                updateTargetNetwork(QNet_target_update_op_list, sess) 

            # Update TensorBoard output
            if(step_tot % tensorboard_update_freq == 0):
                tb_summary = sess.run(QNet_moving.merged, feed_dict={QNet_moving.inputs: s_minibatch,
                                                                     QNet_moving.targetQs_: targets,
                                                                     QNet_moving.actions_: a_minibatch})
                tb_train_writer.add_summary(tb_summary, step_tot)

            # End of step
            f = f_next
            buff_f.store(f)
            step += 1
            ep_rew += r

            # End episode
            if(done == True):
                print('Episode: {}'.format(ep),
                      'Steps : {}'.format(step_tot),
                      'Total reward: {}'.format(ep_rew),
                      'Explore P: {:.4f}'.format(explore_p))
                rewards_list.append((ep,ep_rew,step_tot))
                break

    
        # Save model checkpoint
        if(ep % model_save_interval == 0):
            saveModel(ep, model_save_dir, tf_saver, sess, save_model)



eps, rews, steps = np.array(rewards_list).T
eps = np.reshape(eps, [-1,1])
rews = np.reshape(rews, [-1,1])
steps = np.reshape(steps, [-1,1])
out_array = np.concatenate((steps, rews),axis=1)
print(out_array.shape)
# Output rewards per episode
np.savetxt("rew_new.dat", out_array, fmt="%d")
#plt.plot(rewards_list)
#plt.show()

env.close()
