# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.15.2
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# ## Florian BARRAL
# florian.barral@essec.edu

# **You may need to install [OpenCV](https://pypi.python.org/pypi/opencv-python) and [scikit-video](http://www.scikit-video.org/stable/).**

# +
import keras
import numpy as np
import io
import base64
from IPython.display import HTML
import skvideo.io

import cv2
import json
import random

from keras.models import Sequential,model_from_json
from keras.layers.core import Dense, Reshape, Flatten
from keras.optimizers import sgd
from keras.layers import Conv2D, MaxPooling2D, Activation, AveragePooling2D,Reshape,BatchNormalization


# -

# # MiniProject #3: Deep Reinforcement Learning

# __Notations__: $E_p$ is the expectation under probability $p$. Please justify each of your answer and widely comment your code.

# # Context

# In a reinforcement learning algorithm, we modelize each step $t$ as an action $a_t$ obtained from a state $s_t$, i.e. $\{(a_{t},s_{t})_{t\leq T}\}$ having the Markov property. We consider a discount factor $\gamma \in [0,1]$ that ensures convergence. The goal is to find among all the policies $\pi$, one that maximizes the expected reward:
#
# \begin{equation*}
# R(\pi)=\sum_{t\leq T}E_{p^{\pi}}[\gamma^t r(s_{t},a_{t})] \> ,
# \end{equation*}
#
# where: 
# \begin{equation*}p^{\pi}(a_{0},a_{1},s_{1},...,a_{T},s_{T})=p(a_{0})\prod_{t=1}^{T}\pi(a_{t}|s_{t})p(s_{t+1}|s_{t},a_{t}) \> .
# \end{equation*}
#
# We note the $Q$-function:
#
# \begin{equation*}Q^\pi(s,a)=E_{p^{\pi}}[\sum_{t\leq T}\gamma^{t}r(s_{t},a_{t})|s_{0}=s,a_{0}=a] \> .
# \end{equation*}
#
# Thus, the optimal Q function is:
# \begin{equation*}
# Q^*(s,a)=\max_{\pi}Q^\pi(s,a) \> .
# \end{equation*}
#
# In this project, we will apply the deep reinforcement learning techniques to a simple game: an agent will have to learn from scratch a policy that will permit it maximizing a reward.

# ## The environment, the agent and the game

# ### The environment

# ```Environment``` is an abstract class that represents the states, rewards, and actions to obtain the new state.

class Environment(object):
    def __init__(self):
        pass

    def act(self, act):
        """
        One can act on the environment and obtain its reaction:
        - the new state
        - the reward of the new state
        - should we continue the game?

        :return: state, reward, game_over
        """
        pass


    def reset(self):
        """
        Reinitialize the environment to a random state and returns
        the original state

        :return: state
        """
        pass
    
    def draw(self):
        """
        Visualize in the console or graphically the current state
        """
        pass


# The method ```act``` allows to act on the environment at a given state $s_t$ (stored internally), via action $a_t$. The method will return the new state $s_{t+1}$, the reward $r(s_{t},a_{t})$ and determines if $t\leq T$ (*game_over*).
#
# The method ```reset``` simply reinitializes the environment to a random state $s_0$.
#
# The method ```draw``` displays the current state $s_t$ (this is useful to check the behavior of the Agent).
#
# We modelize $s_t$ as a tensor, while $a_t$ is an integer.

# ### The Agent

# The goal of the ```Agent``` is to interact with the ```Environment``` by proposing actions $a_t$ obtained from a given state $s_t$ to attempt to maximize its __reward__ $r(s_t,a_t)$. We propose the following abstract class:

class Agent(object):
    def __init__(self, epsilon=0.1, n_action=4):
        self.epsilon = epsilon
        self.n_action = n_action
    
    def set_epsilon(self,e):
        self.epsilon = e

    def act(self,s,train=True):
        """ This function should return the next action to do:
        an integer between 0 and 4 (not included) with a random exploration of epsilon"""
        if train:
            if np.random.rand() <= self.epsilon:
                a = np.random.randint(0, self.n_action, size=1)[0]
            else:
                a = self.learned_act(s)
        else: # in some cases, this can improve the performance.. remove it if poor performances
            a = self.learned_act(s)

        return a

    def learned_act(self,s):
        """ Act via the policy of the agent, from a given state s
        it proposes an action a"""
        pass

    def reinforce(self, s, n_s, a, r, game_over_):
        """ This function is the core of the learning algorithm. 
        It takes as an input the current state s_, the next state n_s_
        the action a_ used to move from s_ to n_s_ and the reward r_.
        
        Its goal is to learn a policy.
        """
        pass

    def save(self):
        """ This function returns basic stats if applicable: the
        loss and/or the model"""
        pass

    def load(self):
        """ This function allows to restore a model"""
        pass


# ***
# __Question 1__:
# Explain the function act. Why is ```epsilon``` essential?

# The function ```act``` returns an action that the Agent needs to take according to a policy/tactic.  <br> 
# $\epsilon$ is essential because it allows the agent to explore. The algorithm would just return the learned tactic without this probability of random action. By exploring other states through random actions, the agent can optimize his learned policy/tactic.

# ***
# ### The Game

# The ```Agent``` and the ```Environment``` work in an interlaced way as in the following (take some time to understand this code as it is the core of the project)
#
# ```python
#
# epoch = 300
# env = Environment()
# agent = Agent()
#
#
# # Number of won games
# score = 0
# loss = 0
#
#
# for e in range(epoch):
#     # At each epoch, we restart to a fresh game and get the initial state
#     state = env.reset()
#     # This assumes that the games will end
#     game_over = False
#
#     win = 0
#     lose = 0
#     
#     while not game_over:
#         # The agent performs an action
#         action = agent.act(state)
#
#         # Apply an action to the environment, get the next state, the reward
#         # and if the games end
#         prev_state = state
#         state, reward, game_over = env.act(action)
#
#         # Update the counters
#         if reward > 0:
#             win = win + reward
#         if reward < 0:
#             lose = lose -reward
#
#         # Apply the reinforcement strategy
#         loss = agent.reinforce(prev_state, state,  action, reward, game_over)
#
#     # Save as a mp4
#     if e % 10 == 0:
#         env.draw(e)
#
#     # Update stats
#     score += win-lose
#
#     print("Epoch {:03d}/{:03d} | Loss {:.4f} | Win/lose count {}/{} ({})"
#           .format(e, epoch, loss, win, lose, win-lose))
#     agent.save()
# ```

# # The game, *eat cheese*

# A rat runs on an island and tries to eat as much as possible. The island is subdivided into $N\times N$ cells, in which there are cheese (+0.5) and poisonous cells (-1). The rat has a visibility of 2 cells (thus it can see $5^2$ cells). The rat is given a time $T$ to accumulate as much food as possible. It can perform 4 actions: going up, down, left, right. 
#
# The goal is to code an agent to solve this task that will learn by trial and error. We propose the following environment:

class Environment(object):
    def __init__(self, grid_size=10, max_time=500, temperature=0.1):
        grid_size = grid_size+4
        self.grid_size = grid_size
        self.max_time = max_time
        self.temperature = temperature

        #board on which one plays
        self.board = np.zeros((grid_size,grid_size))
        self.position = np.zeros((grid_size,grid_size))

        # coordinate of the rat
        self.x = 0
        self.y = 1

        # self time
        self.t = 0

        self.scale=16

        self.to_draw = np.zeros((max_time+2, grid_size*self.scale, grid_size*self.scale, 3))


    def draw(self,e):
        skvideo.io.vwrite(str(e) + '.mp4', self.to_draw)

    def get_frame(self,t):
        b = np.zeros((self.grid_size,self.grid_size,3))+128
        b[self.board>0,0] = 256
        b[self.board < 0, 2] = 256
        b[self.x,self.y,:]=256
        b[-2:,:,:]=0
        b[:,-2:,:]=0
        b[:2,:,:]=0
        b[:,:2,:]=0
        
        b =  cv2.resize(b, None, fx=self.scale, fy=self.scale, interpolation=cv2.INTER_NEAREST)

        self.to_draw[t,:,:,:]=b


    def act(self, action):
        """This function returns the new state, reward and decides if the
        game ends."""

        self.get_frame(int(self.t))

        self.position = np.zeros((self.grid_size, self.grid_size))

        self.position[0:2,:]= -1
        self.position[:,0:2] = -1
        self.position[-2:, :] = -1
        self.position[:,-2:] = -1

        self.position[self.x, self.y] = 1
        if action == 0:
            if self.x == self.grid_size-3:
                self.x = self.x-1
            else:
                self.x = self.x + 1
        elif action == 1:
            if self.x == 2:
                self.x = self.x+1
            else:
                self.x = self.x-1
        elif action == 2:
            if self.y == self.grid_size - 3:
                self.y = self.y - 1
            else:
                self.y = self.y + 1
        elif action == 3:
            if self.y == 2:
                self.y = self.y + 1
            else:
                self.y = self.y - 1
        else:
            RuntimeError('Error: action not recognized')

        self.t = self.t + 1
        reward = self.board[self.x, self.y]
        self.board[self.x, self.y] = 0
        game_over = self.t > self.max_time
        state = np.concatenate((self.board.reshape(self.grid_size, self.grid_size,1),
                        self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)
        state = state[self.x-2:self.x+3,self.y-2:self.y+3,:]

        return state, reward, game_over

    def reset(self):
        """This function resets the game and returns the initial state"""

        self.x = np.random.randint(3, self.grid_size-3, size=1)[0]
        self.y = np.random.randint(3, self.grid_size-3, size=1)[0]


        bonus = 0.5*np.random.binomial(1,self.temperature,size=self.grid_size**2)
        bonus = bonus.reshape(self.grid_size,self.grid_size)

        malus = -1.0*np.random.binomial(1,self.temperature,size=self.grid_size**2)
        malus = malus.reshape(self.grid_size, self.grid_size)

        self.to_draw = np.zeros((self.max_time+2, self.grid_size*self.scale, self.grid_size*self.scale, 3))


        malus[bonus>0]=0

        self.board = bonus + malus

        self.position = np.zeros((self.grid_size, self.grid_size))
        self.position[0:2,:]= -1
        self.position[:,0:2] = -1
        self.position[-2:, :] = -1
        self.position[:,-2:] = -1
        self.board[self.x,self.y] = 0
        self.t = 0

        state = np.concatenate((
                               self.board.reshape(self.grid_size, self.grid_size,1),
                        self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)

        state = state[self.x - 2:self.x + 3, self.y - 2:self.y + 3, :]
        return state

# The following elements are important because they correspond to the hyper parameters for this project:

# +
# parameters
size = 13
T=200
temperature=0.3
epochs_train=30 # set small when debugging
epochs_test=30 # set small when debugging

# display videos
def display_videos(name):
    video = io.open(name, 'r+b').read()
    encoded = base64.b64encode(video)
    return '''<video alt="test" controls>
                <source src="data:video/mp4;base64,{0}" type="video/mp4" />
             </video>'''.format(encoded.decode('ascii'))


# -

# __Question 2__ Explain the use of the arrays ```position``` and ```board```.

# The `position` is the currently tracked position of the rat on the island.
#
# The `board` is the representation of the island of all the cells, that means all the possible states.

# ## Random Agent

# ***
# __Question 3__ Implement a random Agent (only ```learned_act``` needs to be implemented):

class RandomAgent(Agent):
    def __init__(self,n_action=4):
        super(RandomAgent, self).__init__()
        self.n_action = n_action
        pass

    def learned_act(self, s):
        return np.random.randint(0 ,self.n_action ,size=1)
        pass


# ***
# ***
# __Question 4__ Visualize the game moves. You need to fill in the following function for the evaluation:

def test(agent,env,epochs,prefix=''):
    # Number of won games
    score = 0
        
    for e in range(1,epochs+1):
        
        ##### FILL IN HERE
        #At each epoch, we restart to a fresh game and get the initial state
        state = env.reset()
        # This assumes that the games will end
        game_over = False

        win = 0
        lose = 0
        
        while not game_over:
            # The agent performs an action
            action = agent.learned_act(state)
            
            # Apply an action to the environment, get the next state, the reward
            # And if the games end
            state, reward, game_over = env.act(action)
            
            # Update the counters
            if reward > 0:
                win += reward
            if reward < 0:
                lose -= reward
            
            # Apply the reinforcement strategy
            
        # Save as a mp4
        if e % 10 == 0:
            env.draw(prefix+str(e))

        # Update stats
        score += win-lose

        print("Win/lose count {}/{}. Average score ({})"
              .format(win, lose, score/(1+e)))
    print('Final score: '+str(score/epochs))


# +
# Initialize the game
env = Environment(grid_size=size, max_time=T,temperature=temperature)

# Initialize the agent!
agent = RandomAgent()

test(agent,env,epochs_test,prefix='random')
HTML(display_videos('random10.mp4'))


# -

# ***
# ## DQN

# Let us assume here that $T=\infty$.
#
# ***
# __Question 5__ Let $\pi$ be a policy, **(1)** show that:
#
# \begin{equation*}
# Q^{\pi}(s,a)=E_{(s',a')\sim p(.|s,a)}[r(s,a)+\gamma Q^{\pi}(s',a')]
# \end{equation*}
#
# Then, **(2)** show that for the optimal policy $\pi^*$ (we assume its existence), the following holds: 
#
# \begin{equation*}
# Q^{*}(s,a)=E_{s'\sim \pi^*(.|s,a)}[r(s,a)+\gamma\max_{a'}Q^{*}(s',a')].
# \end{equation*}
# Finally, **(3)** deduce that a plausible objective is:
#
# \begin{equation*}
# \mathcal{L}(\theta)=E_{s' \sim \pi^*(.|s,a)}\Vert r+\gamma\max\max_{a'}Q(s',a',\theta)-Q(s,a,\theta)\Vert^{2}.
# \end{equation*}
#
#
#

# **(1)**
#
# Given the Bellman equation:
#
# \begin{align*}
# Q^{\pi}(s,a) & = E_{\pi}[\Sigma^{\infty}_{t=0} \gamma^t r_t(s_t, \pi(s_t),s_{t+1}) | s_t = s,a_t = a] \\
# \Rightarrow Q^{\pi}(s_0,a_0) & = E_{\pi}[r_0(s_0,a_0,s_1)+\Sigma^{\infty}_{t=1} \gamma^t r_t (s_t,\pi(s_t),s_{t+1}) | s_0,a_0 ] \\
# & = E_{\pi}[r_0(s_0,a_0,s_1) + \gamma E_{\pi} \{ \Sigma_{t=1}^{\infty} \gamma^{t-1} r_t (s_t,\pi(s_t),s_{t+1} | s_1,a_1 \} s_0,a_0 ] \\
# & = E_{\pi} [r_0 (s_0,a_0,s_1)+\gamma Q^{\pi}(s_1,a_1)]
# \end{align*}
#
# Then:
# $$\boxed{Q^{\pi}(s,a) = E_{(s',a') \sim p(.|s,a)} [r(s,a)+\gamma Q^{\pi}(s',a')]}$$ 
# **(2)** 
#
# Given the Bellman equation:
#
# $$V^{\pi}(s_0) = E[\Sigma^{\infty}_{t=0} \gamma^t r_t (s_t,\pi(s_t),s_{t+1} | s_0] $$
# and 
# $$V^*(s_0) = \max E[\Sigma^{\infty}_{t=0}\gamma^t r_t (s_t,\pi(s_t),s_{t+1} | s_0] $$
#
# Then
#
# \begin{align*} V^*(s_0) & = \underset{\pi,a_0}{\max} E[r_0(s_0,a_0,s_1) + \Sigma^{\infty}_{t=1} \gamma^t r_t (s_t,\pi(s_t),s_(t+1)|s_0] \\
# & = \underset{a_0}{\max} E_{s_1} [r_0(s_0,a_0,s_1)+ \gamma \underset{\pi}{\max} \Sigma^{\infty}_{t=1} E_{s_2,s_3...} \{ \Sigma^{\infty}_{t=1} \gamma^{t-1} r_t(s_t,\pi(s_{t+1}),s_{t+1} \}  | s_0] \\
# & = \underset{a_0}{\max} E_{s_1}[r_0(s_0,a_0,s_1)+\gamma V^* (s_1)] 
# \end{align*}
#
# Then:
# $$ V^*(s) = \max E_{s'}[r(s,a,s')+\gamma V^*(s')|s] $$
#
# Since $V^*(s) = \underset{a'}{\max}Q^*(s,a')$
#
# \begin{align*}
# \Rightarrow V^*(s) & = \underset{a}{\max}E_{s'}[r(s,a,s')+\gamma \underset{a'}{\max} Q^*(s',a')|s,a] \\
# \Rightarrow Q^*(s,a) & = E_{s'}[r(s,a,s')+\gamma \underset{a'}{\max} Q^*(s',a')|s,a]
# \end{align*}
# $$
# \boxed{Q^*(s,a) = E_{s'\sim\pi^*(.|s,a)}[r(s,a)+\gamma \underset{a'}{\max} Q^*(s',a')]}$$
#
#
#
# **(3)** 
#
# Given the loss function $\mathcal{L}(\theta)= [Q^*(s,a,\theta) - Q(s,a,\theta)]^2$
#
# By substituting $Q^*(s,a,\theta)$, we get:
#
# $$\mathcal{L}(\theta) = [E_{s'\sim \pi^*(.|s,a,s')}[r(s,a,s')+\gamma\max_{a'}Q^{*}(s',a,\theta)] - Q(s,a,\theta)]^2$$
# $$ \Rightarrow \boxed{\mathcal{L}(\theta)=E_{s' \sim \pi^*(.|s,a)}\Vert r+\gamma\max\max_{a'}Q(s',a',\theta)-Q(s,a,\theta)\Vert^{2}} $$

# ***
# The DQN-learning algorithm relies on these derivations to train the parameters $\theta$ of a Deep Neural Network:
#
# 1. At the state $s_t$, select the action $a_t$ with best reward using $Q_t$ and store the results;
#
# 2. Obtain the new state $s_{t+1}$ from the environment $p$;
#
# 3. Store $(s_t,a_t,s_{t+1})$;
#
# 4. Obtain $Q_{t+1}$ by minimizing  $\mathcal{L}$ from a recovered batch from the previously stored results.
#
# ***
# __Question 6__ Implement the class ```Memory``` that stores moves (in a replay buffer) via ```remember``` and provides a ```random_access``` to these. Specify a maximum memory size to avoid side effects. You can for example use a ```list()``` and set by default ```max_memory=100```.

class Memory(object):
    def __init__(self, max_memory=100):
        self.max_memory = max_memory
        self.memory = list()

    def remember(self, m):
        self.memory.append(m)
        pass

    def random_access(self):
        r= np.random.randint(0, len(self.memory))
        return self.memory[r]


# ***
# The pipeline we will use for training is given below:

def train(agent,env,epoch,prefix=''):
    # Number of won games
    score = 0
    loss = 0

    for e in range(1,epoch+1):
        # At each epoch, we restart to a fresh game and get the initial state
        state = env.reset()
        # This assumes that the games will terminate
        game_over = False

        win = 0
        lose = 0

        while not game_over:
            # The agent performs an action
            action = agent.act(state)

            # Apply an action to the environment, get the next state, the reward
            # and if the games end
            prev_state = state
            state, reward, game_over = env.act(action)

            # Update the counters
            if reward > 0:
                win = win + reward
            if reward < 0:
                lose = lose -reward

            # Apply the reinforcement strategy
            loss = agent.reinforce(prev_state, state,  action, reward, game_over)

        # Save as a mp4
        if e % 10 == 0:
            env.draw(prefix+str(e))

        # Update stats
        score += win-lose

        print("Epoch {:03d}/{:03d} | Loss {:.4f} | Win/lose count {}/{} ({})"
              .format(e, epoch, loss, win, lose, win-lose))
        agent.save(name_weights=prefix+'model.h5',name_model=prefix+'model.json')


# ***
# __Question 7__ Implement the DQN training algorithm using a cascade of fully connected layers. You can use different learning rate, batch size or memory size parameters. In particular, the loss might oscillate while the player will start to win the games. You have to find a good criterium.

# +
class DQN(Agent):
    def __init__(self, grid_size,  epsilon = 0.1, memory_size=100, batch_size = 16,n_state=2):
        super(DQN, self).__init__(epsilon = epsilon)

        # Discount for Q learning
        self.discount = 0.99
        
        self.grid_size = grid_size
        
        # number of state
        self.n_state = n_state

        # Memory
        self.memory = Memory(memory_size)
        
        # Batch size when learning
        self.batch_size = batch_size

    def learned_act(self, s):
        acts_ = self.model.predict(np.reshape(s,(1,5,5,self.n_state)))
        return np.argmax(acts_[0])
        

    def reinforce(self, s_, n_s_, a_, r_, game_over_):
        # Two steps: first memorize the states, second learn from the pool

        self.memory.remember([s_, n_s_, a_, r_, game_over_])
        
        input_states = np.zeros((self.batch_size, 5,5,self.n_state))
        target_q = np.zeros((self.batch_size, 4))
        
        for i in range(self.batch_size):
            current_state, next_state, action, reward, game_over_ = self.memory.random_access() #sample some moves
            input_states[i] = current_state
            
            s_reshape=np.reshape(current_state,(1,5,5,self.n_state))
            next_state=np.reshape(next_state,(1,5,5,self.n_state))
            target_q[i] = self.model.predict(s_reshape)
            
           
            if game_over_:

                target_q[i,a_]= reward
            else:
                q = self.model.predict(next_state)
                target_q[i,a_] = reward + self.discount * np.max(q)
        
                
        ######## FILL IN
        # HINT: Clip the target to avoid exploiding gradients.. -- clipping is a bit tighter
        
        
        target_q = np.clip(target_q, -3, 3)

        l = self.model.train_on_batch(input_states, target_q)


        return l

    def save(self,name_weights='model.h5',name_model='model.json'):
        self.model.save_weights(name_weights, overwrite=True)
        with open(name_model, "w") as outfile:
            json.dump(self.model.to_json(), outfile)
            
    def load(self,name_weights='model.h5',name_model='model.json'):
        with open(name_model, "r") as jfile:
            model = model_from_json(json.load(jfile))
        model.load_weights(name_weights)
        model.compile("sgd", "mse")
        self.model = model

            
class DQN_FC(DQN):
    def __init__(self, *args, lr=0.1,**kwargs):
        super(DQN_FC, self).__init__( *args,**kwargs)
        
        # NN Model
        model = Sequential()
        model.add(Reshape((5*5*self.n_state,), input_shape=(5,5,self.n_state)))
        model.add(Dense(128,input_shape=(50,),init='uniform',activation='relu'))
        model.add(Dense(64,init='uniform',activation='relu'))
        model.add(Dense(32,init='uniform',activation='relu'))
        model.add(Dense(16,activation='relu'))
        model.add(Dense(8,activation='relu'))
        model.add(Dense(4,init='uniform',activation='linear'))
        model.compile(sgd(lr=lr, decay=1e-4, momentum=0.0), "mse")
        self.model = model


# -

env = Environment(grid_size=size, max_time=T, temperature=0.3)
agent = DQN_FC(size, lr=.1, epsilon = 0.2, memory_size=16000, batch_size = 32)
train(agent, env, epochs_train, prefix='fc_train_')

HTML(display_videos('fc_train_30.mp4'))


# ***
# ***
# __Question 8__ Implement the DQN training algorithm using a CNN (for example, 2 convolutional layers and one final fully connected layer).

class DQN_CNN(DQN):
    def __init__(self, *args,lr=0.1,**kwargs):
        super(DQN_CNN, self).__init__(*args,**kwargs)
        
        ###### FILL IN
        model = Sequential()
        model.add(Conv2D(32,kernel_size=(2,2),strides=(1,1),
                         activation='relu',input_shape=(5,5,self.n_state)))
        model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))
        model.add(Conv2D(64, (1, 1), activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Flatten())
        model.add(Dense(4, activation='linear'))
        model.compile(sgd(lr=lr, decay=1e-4, momentum=0.0), "mse")
        print(model.summary())
        self.model = model


env = Environment(grid_size=size, max_time=T, temperature=0.3)
agent = DQN_CNN(size, lr=.1, epsilon = 0.1, memory_size=2000, batch_size = 32)
train(agent,env,epochs_train,prefix='cnn_train_')

HTML(display_videos('cnn_train_30.mp4'))

# ***
# ***
# __Question 9__ Test both algorithms and compare their performances. Which issue(s) do you observe? Observe also different behaviors by changing the temperature.

# ### Testing CNN & FC with temperature set to 0.3

# +
env = Environment(grid_size=size, max_time=T,temperature=0.3)
agent_cnn = DQN_CNN(size, lr=.1, epsilon = 0.1, memory_size=2000, batch_size = 32)
agent_cnn.load(name_weights='cnn_train_model.h5',name_model='cnn_train_model.json')

agent_fc = DQN_FC(size, lr=.1, epsilon = 0.1, memory_size=2000, batch_size = 32)
agent_cnn.load(name_weights='fc_train_model.h5',name_model='fc_train_model.json')
print('Test of the CNN')
test(agent_cnn,env,epochs_test,prefix='cnn_test_03_')
print('Test of the FC')
test(agent_fc,env,epochs_test,prefix='fc_test_03_')
# -

HTML(display_videos('cnn_test_03_30.mp4'))

HTML(display_videos('fc_test_03_30.mp4'))

# ### Testing CNN & FC with temperature set to 0.9

# +
env = Environment(grid_size=size, max_time=T,temperature=0.3)
agent_cnn = DQN_CNN(size, lr=.1, epsilon = 0.1, memory_size=2000, batch_size = 32)
agent_cnn.load(name_weights='cnn_train_model.h5',name_model='cnn_train_model.json')

agent_fc = DQN_FC(size, lr=.1, epsilon = 0.1, memory_size=2000, batch_size = 32)
agent_cnn.load(name_weights='fc_train_model.h5',name_model='fc_train_model.json')
print('Test of the CNN')
test(agent_cnn,env,epochs_test,prefix='cnn_test_09_')
print('Test of the FC')
test(agent_fc,env,epochs_test,prefix='fc_test_09_')
# -

HTML(display_videos('cnn_test_09_30.mp4'))

HTML(display_videos('fc_test_09_30.mp4'))


# The CNN model seems to perform better and to be more stable than the FC's.
#
# When increasing the temperature, results improve for both models. Something expected and that makes sense since we are increasing the percentage of cheeses and bonus on the board
#
# Yet, in both models, the agent tends not exploring the island anymore.

# ***
#
# The algorithm tends to not explore the map which can be an issue. We propose two ideas in order to encourage exploration:
# 1. Incorporating a decreasing $\epsilon$-greedy exploration. You can use the method ```set_epsilon```
# 2. Append via the environment a new state that describes if a cell has been visited or not
#
# ***
# __Question 10__ Design a new ```train_explore``` function and environment class ```EnvironmentExploring``` to tackle the issue of exploration.
#
#

# +
def train_explore(agent,env,epoch,prefix=''):
    # Number of won games
    score = 0
    loss = 0

    for e in range(1,epoch+1):
        # At each epoch, we restart to a fresh game and get the initial state
        state = env.reset()
        # This assumes that the games will terminate
        game_over = False

        win = 0
        lose = 0

        while not game_over:
            # The agent performs an action
            action = agent.act(state,train=True)

            # Apply an action to the environment, get the next state, the reward
            # and if the games end
            prev_state = state
            state, reward, game_over = env.act(action)

            # Update the counters
            if reward > 0:
                win = win + reward
            if reward < 0:
                lose = lose - reward

            # Apply the reinforcement strategy
            loss = agent.reinforce(prev_state, state,  action, reward, game_over)

        # Save as a mp4
        if e % 10 == 0:
            env.draw(prefix+str(e))

        # Update stats
        score += win-lose

        print("Epoch {:03d}/{:03d} | Loss {:.4f} | Win/lose count {}/{} ({})"
              .format(e, epoch, loss, win, lose, win-lose))
        agent.save(name_weights=prefix+'model.h5',name_model=prefix+'model.json')
        
class EnvironmentExploring(object):
    def __init__(self, grid_size=10, max_time=500, temperature=0.1):
        grid_size = grid_size+4
        self.grid_size = grid_size
        self.max_time = max_time
        self.temperature = temperature

        #board on which one plays
        self.board = np.zeros((grid_size,grid_size))
        self.position = np.zeros((grid_size,grid_size))

        # coordinate of the rat
        self.x = 0
        self.y = 1
#
        self.malus_position = np.zeros((grid_size,grid_size))      
#         
        # self time
        self.t = 0

        self.scale=16

        self.to_draw = np.zeros((max_time+2, grid_size*self.scale, grid_size*self.scale, 3))


    def draw(self,e):
        skvideo.io.vwrite(str(e) + '.mp4', self.to_draw)

    def get_frame(self,t):
        b = np.zeros((self.grid_size,self.grid_size,3))+128
        b[self.board>0,0] = 256
        b[self.board < 0, 2] = 256
        b[self.x,self.y,:]=256
        b[-2:,:,:]=0
        b[:,-2:,:]=0
        b[:2,:,:]=0
        b[:,:2,:]=0
        
        b =  cv2.resize(b, None, fx=self.scale, fy=self.scale, interpolation=cv2.INTER_NEAREST)

        self.to_draw[t,:,:,:]=b


    def act(self, action):
        """This function returns the new state, reward and decides if the
        game ends."""

        self.get_frame(int(self.t))

        self.position = np.zeros((self.grid_size, self.grid_size))

        self.position[0:2,:]= -1
        self.position[:,0:2] = -1
        self.position[-2:, :] = -1
        self.position[:,-2:] = -1

        self.position[self.x, self.y] = 1
        if action == 0:
            if self.x == self.grid_size-3:
                self.x = self.x-1
            else:
                self.x = self.x + 1
        elif action == 1:
            if self.x == 2:
                self.x = self.x+1
            else:
                self.x = self.x-1
        elif action == 2:
            if self.y == self.grid_size - 3:
                self.y = self.y - 1
            else:
                self.y = self.y + 1
        elif action == 3:
            if self.y == 2:
                self.y = self.y + 1
            else:
                self.y = self.y - 1
        else:
            RuntimeError('Error: action not recognized')

        self.t = self.t + 1
#       #reward = self.board[self.x, self.y]
        reward = 0
        if train:
            reward = -self.malus_position[self.x, self.y]
        self.malus_position[self.x, self.y] = 0.1 
        reward = reward + self.board[self.x, self.y]
#        
        self.board[self.x, self.y] = 0
        game_over = self.t > self.max_time
        #state = np.concatenate((self.board.reshape(self.grid_size, self.grid_size,1),
        #                self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)
#        
        # 3 "feature" states instead of 2
        state = np.concatenate((self.malus_position.reshape(self.grid_size, self.grid_size,1), #OK ajouté dans init
                                self.board.reshape(self.grid_size, self.grid_size,1), #OK deja dans init
                        self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)#OK deja init
#        
        state = state[self.x-2:self.x+3,self.y-2:self.y+3,:]

        return state, reward, game_over

    def reset(self):
        """This function resets the game and returns the initial state"""

        self.x = np.random.randint(3, self.grid_size-3, size=1)[0]
        self.y = np.random.randint(3, self.grid_size-3, size=1)[0]


        bonus = 0.5*np.random.binomial(1,self.temperature,size=self.grid_size**2)
        bonus = bonus.reshape(self.grid_size,self.grid_size)

        malus = -1.0*np.random.binomial(1,self.temperature,size=self.grid_size**2)
        malus = malus.reshape(self.grid_size, self.grid_size)

        self.to_draw = np.zeros((self.max_time+2, self.grid_size*self.scale, self.grid_size*self.scale, 3))


        malus[bonus>0]=0

        self.board = bonus + malus

        self.position = np.zeros((self.grid_size, self.grid_size))
        self.position[0:2,:]= -1
        self.position[:,0:2] = -1
        self.position[-2:, :] = -1
        self.position[:,-2:] = -1
        self.board[self.x,self.y] = 0
        self.t = 0

        #state = np.concatenate((
        #                      self.board.reshape(self.grid_size, self.grid_size,1),
        #                self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)
        # 3 "feature" states instead of 2
#
        state = np.concatenate((self.malus_position.reshape(self.grid_size, self.grid_size,1), #OK ajouté dans init
                                        self.board.reshape(self.grid_size, self.grid_size,1), #OK deja dans init
                                self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)#OK deja init
#
        state = state[self.x - 2:self.x + 3, self.y - 2:self.y + 3, :]
        return state
    
## use those samples of code:
#In train explore: OK
#state, reward, game_over = env.act(action, train=True)
#
## In Environment exploring: #OK 
# You will have to change n_state to 3 because you will use one more layer!
#reward = 0
#if train:
#    reward = -self.malus_position[self.x, self.y]
#self.malus_position[self.x, self.y] = 0.1 
#
#reward = reward + self.board[self.x, self.y]
# 3 "feature" states instead of 2
#state = np.concatenate((self.malus_position.reshape(self.grid_size, self.grid_size,1), #OK ajouté dans init
#                                self.board.reshape(self.grid_size, self.grid_size,1), #OK deja dans init
#                        self.position.reshape(self.grid_size, self.grid_size,1)),axis=2)#OK deja init
# -

# Training
env = EnvironmentExploring(grid_size=size, max_time=T, temperature=0.3)
agent = DQN_CNN(size, lr=.1, epsilon = 0.1, memory_size=2000, batch_size = 32,n_state=3)
train_explore(agent, env, epochs_train, prefix='cnn_train_explore_')

HTML(display_videos('cnn_train_explore_30.mp4'))

# Evaluation
test(agent,env,epochs_test,prefix='cnn_test_explore_')

HTML(display_videos('cnn_test_explore_30.mp4'))

# ***
# ***
# __BONUS question__ Use the expert DQN from the previous question to generate some winning games. Train a model that mimicks its behavior. Compare the performances.

#

# ***
