# To support both python 2 and python 3
from __future__ import division, print_function, unicode_literals

# Common imports
import numpy as np
import os

# to make this notebook's output stable across runs
np.random.seed(42)

# To plot pretty figures
%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt
plt.rcParams['axes.labelsize'] = 14
plt.rcParams['xtick.labelsize'] = 12
plt.rcParams['ytick.labelsize'] = 12

# Where to save the figures
PROJECT_ROOT_DIR = "."
CHAPTER_ID = "classification"

def save_fig(fig_id, tight_layout=True):
    path = os.path.join(PROJECT_ROOT_DIR, "images", CHAPTER_ID, fig_id + ".png")
    print("Saving figure", fig_id)
    if tight_layout:
        plt.tight_layout()
    plt.savefig(path, format='png', dpi=300)

# 导入必要的包
import paddle

import paddle.dataset.imdb as imdb

import paddle.fluid as fluid

import numpy as np

import os

!mkdir -p /home/aistudio/.cache/paddle/dataset/imdb/

!cp /home/aistudio/data/data69/aclImdb_v1.tar.gz /home/aistudio/.cache/paddle/dataset/imdb/

# 获取数据字典

print("加载数据字典中...")

word_dict = imdb.word_dict()#这个数据集是一个常用的数据集，已经被paddle封装到底层代码里面了

# 获取数据字典长度

dict_dim = len(word_dict)

print('完成')

# 获取训练和预测数据

print("加载训练数据中...")

train_reader = paddle.batch(paddle.reader.shuffle(imdb.train(word_dict),

                                                  512),

                            batch_size=128)

print("加载测试数据中...")

test_reader = paddle.batch(imdb.test(word_dict), 

                           batch_size=128)

print('完成')

# 定义长短期记忆网络

def lstm_net(ipt, input_dim):

    # 以数据的IDs作为输入

    emb = fluid.layers.embedding(input=ipt, size=[input_dim, 128], is_sparse=True)

    # 第一个全连接层

    fc1 = fluid.layers.fc(input=emb, size=128)

    # 进行一个长短期记忆操作

    lstm1, _ = fluid.layers.dynamic_lstm(input=fc1, #返回：隐藏状态（hidden state），LSTM的神经元状态

                                         size=128) #size=4*hidden_size

    # 第一个最大序列池操作

    fc2 = fluid.layers.sequence_pool(input=fc1, pool_type='max')

    # 第二个最大序列池操作

    lstm2 = fluid.layers.sequence_pool(input=lstm1, pool_type='max')

    # 以softmax作为全连接的输出层，大小为2,也就是正负面

    out = fluid.layers.fc(input=[fc2, lstm2], size=2, act='softmax')#二分类，1x2的概率分布

    return out

# 定义输入数据， lod_level不为0指定输入数据为序列数据

words = fluid.layers.data(name='words', shape=[1], dtype='int64', lod_level=1)

label = fluid.layers.data(name='label', shape=[1], dtype='int64')#label：正向或者负向

# 获取长短期记忆网络

model = lstm_net(words, dict_dim)

# 获取损失函数和准确率

cost = fluid.layers.cross_entropy(input=model, label=label)

avg_cost = fluid.layers.mean(cost)

acc = fluid.layers.accuracy(input=model, label=label)

# 获取预测程序

test_program = fluid.default_main_program().clone(for_test=True)

# 定义优化方法

optimizer = fluid.optimizer.AdagradOptimizer(learning_rate=0.002)

opt = optimizer.minimize(avg_cost)

# 定义使用CPU还是GPU，使用CPU时use_cuda = False,使用GPU时use_cuda = True
use_cuda = True
place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
exe = fluid.Executor(place)
# 进行参数初始化
exe.run(fluid.default_startup_program())

# 定义输入数据的维度

# 定义数据数据的维度，数据的顺序是一条句子数据对应一个标签

feeder = fluid.DataFeeder(place=place, feed_list=[words, label])

# 开始训练

for pass_id in range(10):

    # 进行训练

    train_cost = 0

    for batch_id, data in enumerate(train_reader()):              #遍历train_reader迭代器

        train_cost = exe.run(program=fluid.default_main_program(),#运行主程序

                             feed=feeder.feed(data),              #喂入一个batch的数据

                             fetch_list=[avg_cost])               #fetch均方误差


        if batch_id % 40 == 0:                 #每40次batch打印一次训练、进行一次测试

            print('Pass:%d, Batch:%d, Cost:%0.5f' % (pass_id, batch_id, train_cost[0]))

    # 进行测试

    test_costs = []   #测试的损失值

    test_accs = []    #测试的准确率

    for batch_id, data in enumerate(test_reader()):

        test_cost, test_acc = exe.run(program=test_program,

                                            feed=feeder.feed(data),

                                             fetch_list=[avg_cost, acc])

        test_costs.append(test_cost[0])

        test_accs.append(test_acc[0])

    # 计算平均预测损失在和准确率

    test_cost = (sum(test_costs) / len(test_costs))

    test_acc = (sum(test_accs) / len(test_accs))

    print('Test:%d, Cost:%0.5f, ACC:%0.5f' % (pass_id, test_cost, test_acc))

#保存模型

model_save_dir = "/home/aistudio/work/emotionclassify.inference.model"

# 如果保存路径不存在就创建

if not os.path.exists(model_save_dir):

    os.makedirs(model_save_dir)

print ('save models to %s' % (model_save_dir))

fluid.io.save_inference_model(model_save_dir, #保存推理model的路径

                                  ['words'],      #推理（inference）需要 feed 的数据

                                  [model],         #保存推理（inference）结果的 Variables

                                  exe)            #exe 保存 inference model

# 定义预测数据

reviews_str = ['read the book forget the movie', 'this is a great movie', 'this is very bad']#第一句是中性，第二句是正向，第三局是负向

# 把每个句子拆成一个个单词

reviews = [c.split() for c in reviews_str]

# 获取结束符号的标签

UNK = word_dict['<unk>']

# 获取每句话对应的标签

lod = []

for c in reviews:

    # 需要把单词进行字符串编码转换

    lod.append([word_dict.get(words.encode('utf-8'), UNK) for words in c])

# 获取每句话的单词数量

base_shape = [[len(c) for c in lod]]

# 生成预测数据

tensor_words = fluid.create_lod_tensor(lod, base_shape, place)

infer_exe = fluid.Executor(place)    #创建推测用的executor

inference_scope = fluid.core.Scope() #Scope指定作用域

with fluid.scope_guard(inference_scope):#修改全局/默认作用域（scope）, 运行时中的所有变量都将分配给新的scope。

    #从指定目录中加载 推理model(inference model)

    [inference_program,                                            #推理的program

     feed_target_names,                                            #str列表，包含需要在推理program中提供数据的变量名称

     fetch_targets] = fluid.io.load_inference_model(model_save_dir,#fetch_targets: 推断结果，model_save_dir:模型训练路径 

                                                        infer_exe) #infer_exe: 运行 inference model的 executor

    results = infer_exe.run(inference_program,                                 #运行预测程序

                            feed={feed_target_names[0]: tensor_words},#喂入要预测的x值

                            fetch_list=fetch_targets)                           #得到推测结果 

    # 打印每句话的正负面概率

    for i, r in enumerate(results[0]):

        print("\'%s\'的预测结果为：正面概率为：%0.5f，负面概率为：%0.5f" % (reviews_str[i], r[0], r[1]))

import parl
from parl import layers
import paddle.fluid as fluid
import copy
import numpy as np
import os
import gym
from parl.utils import logger

LEARN_FREQ = 5 # 训练频率，不需要每一个step都learn，攒一些新增经验后再learn，提高效率
MEMORY_SIZE = 20000    # replay memory的大小，越大越占用内存
MEMORY_WARMUP_SIZE = 200  # replay_memory 里需要预存一些经验数据，再从里面sample一个batch的经验让agent去learn
BATCH_SIZE = 32   # 每次给agent learn的数据数量，从replay memory随机里sample一批数据出来
GAMMA = 0.99 # reward 的衰减因子，一般取 0.9 到 0.999 不等
LEARNING_RATE = 0.005 # 学习率

#Model用来定义前向(Forward)网络，用户可以自由的定制自己的网络结构。
class Model(parl.Model):
    def __init__(self, act_dim):
        hid1_size = 128
        hid2_size = 128
        # 3层全连接网络
        self.fc1 = layers.fc(size=hid1_size, act='relu')
        self.fc2 = layers.fc(size=hid2_size, act='relu')
        self.fc3 = layers.fc(size=act_dim, act=None)

    def value(self, obs):
        # 定义网络
        # 输入state，输出所有action对应的Q，[Q(s,a1), Q(s,a2), Q(s,a3)...]
        
        h1 = self.fc1(obs)
        h2 = self.fc2(h1)
        Q = self.fc3(h2)
        return Q

# from parl.algorithms import DQN # 也可以直接从parl库中导入DQN算法

class DQN(parl.Algorithm):
    def __init__(self, model, act_dim=None, gamma=None, lr=None):
        """ DQN algorithm
        
        Args:
            model (parl.Model): 定义Q函数的前向网络结构
            act_dim (int): action空间的维度，即有几个action
            gamma (float): reward的衰减因子
            lr (float): learning rate 学习率.
        """
        self.model = model # 我们用来获取 current Q 的模型
        self.target_model = copy.deepcopy(model) # 创建一个target Q模型，创建的策略是直接从model复制给target

        assert isinstance(act_dim, int)
        assert isinstance(gamma, float)
        assert isinstance(lr, float)
        self.act_dim = act_dim # 把这些参数变成class properties
        self.gamma = gamma
        self.lr = lr

    def predict(self, obs): # 使用 current Q network 获取所有action的 Q values
        """ 使用self.model的value网络来获取 [Q(s,a1),Q(s,a2),...]
        """
        return self.model.value(obs)

    def learn(self, obs, action, reward, next_obs, terminal):
        """ 使用DQN算法更新self.model的value网络
        """
        # 从target_model中获取 max Q' 的值，用于计算target_Q
        next_pred_value = self.target_model.value(next_obs) # 获取 target Q network 的所有action的 Q values
        best_v = layers.reduce_max(next_pred_value, dim=1) # 获取最大的Q值
        best_v.stop_gradient = True  # 阻止梯度传递
        terminal = layers.cast(terminal, dtype='float32') # 把terminal （是否终止）换为一个float32类型的数组，如果终止里面存储1，如果不终止里面存储0
        target = reward + (1.0 - terminal) * self.gamma * best_v # 这里如果终止， 1-terminal 对应的元素为0，就不需要取best_v，不然还是要取best_v

        pred_value = self.model.value(obs)  # 获取 current Q network 的所有action的 Q values

        # 接着我们需要获取action对应的Q，这里使用了一个one-hot encoding来做乘法运算，相当于选中了Q values中action对应的那个值

        # 将action转one-hot向量，比如：3 => [0,0,0,1,0]
        action_onehot = layers.one_hot(action, self.act_dim)
        action_onehot = layers.cast(action_onehot, dtype='float32')
        # 下面一行是逐元素相乘，拿到action对应的 Q(s,a)
        # 比如：pred_value = [[2.3, 5.7, 1.2, 3.9, 1.4]], action_onehot = [[0,0,0,1,0]]
        #  ==> pred_action_value = [[3.9]]
        pred_action_value = layers.reduce_sum(
            layers.elementwise_mul(action_onehot, pred_value), dim=1)

        # 计算 Q(s,a) 与 target_Q的MSE均方差，得到loss
        cost = layers.square_error_cost(pred_action_value, target)
        cost = layers.reduce_mean(cost) # Loss 对于每一个样本都是一个数字，为了优化我们求平均数
        optimizer = fluid.optimizer.Adam(learning_rate=self.lr)  # 使用Adam优化器，Adam是一种优化算法
        optimizer.minimize(cost)
        return cost

    def sync_target(self):
        """ 把 self.model 的模型参数值同步到 self.target_model
        """
        self.model.sync_weights_to(self.target_model) # 这个函数主要是为了更新 Target Q，因为每一段时间我们就需要使用 Current Q Network 更新一次Target Q Network

Algorithm定义了具体的算法来更新前向网络(Model)，也就是通过定义损失函数来更新Model，和算法相关的计算都放在algorithm中。
class Agent(parl.Agent):
    def __init__(self,
                 algorithm,
                 obs_dim,
                 act_dim,
                 e_greed=0.1,
                 e_greed_decrement=0):
        assert isinstance(obs_dim, int)
        assert isinstance(act_dim, int)
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        super(Agent, self).__init__(algorithm)

        self.global_step = 0
        self.update_target_steps = 200  # 每隔200个training steps再把model的参数复制到target_model中

        self.e_greed = e_greed  # 有一定概率随机选取动作，探索
        self.e_greed_decrement = e_greed_decrement  # 随着训练逐步收敛，探索的程度慢慢降低

    def build_program(self):
        self.pred_program = fluid.Program()
        self.learn_program = fluid.Program()

        with fluid.program_guard(self.pred_program):  # 搭建计算图用于 预测动作，定义输入输出变量
            obs = layers.data(
                name='obs', shape=[self.obs_dim], dtype='float32')
            self.value = self.alg.predict(obs)

        with fluid.program_guard(self.learn_program):  # 搭建计算图用于 更新Q网络，定义输入输出变量
            obs = layers.data(
                name='obs', shape=[self.obs_dim], dtype='float32')
            action = layers.data(name='act', shape=[1], dtype='int32')
            reward = layers.data(name='reward', shape=[], dtype='float32')
            next_obs = layers.data(
                name='next_obs', shape=[self.obs_dim], dtype='float32')
            terminal = layers.data(name='terminal', shape=[], dtype='bool')
            self.cost = self.alg.learn(obs, action, reward, next_obs, terminal)

    def sample(self, obs):
        sample = np.random.rand()  # 产生0~1之间的小数
        if sample < self.e_greed:
            act = np.random.randint(self.act_dim)  # 探索：每个动作都有概率被选择
        else:
            act = self.predict(obs)  # 选择最优动作
        self.e_greed = max(
            0.01, self.e_greed - self.e_greed_decrement)  # 随着训练逐步收敛，探索的程度慢慢降低
        return act

    def predict(self, obs):  # 选择最优动作
        obs = np.expand_dims(obs, axis=0)
        pred_Q = self.fluid_executor.run(
            self.pred_program,
            feed={'obs': obs.astype('float32')},
            fetch_list=[self.value])[0]
        pred_Q = np.squeeze(pred_Q, axis=0)
        act = np.argmax(pred_Q)  # 选择Q最大的下标，即对应的动作
        return act

    def learn(self, obs, act, reward, next_obs, terminal):
        # 每隔200个training steps同步一次model和target_model的参数
        if self.global_step % self.update_target_steps == 0:
            self.alg.sync_target()
        self.global_step += 1

        act = np.expand_dims(act, -1)
        feed = {
            'obs': obs.astype('float32'),
            'act': act.astype('int32'),
            'reward': reward,
            'next_obs': next_obs.astype('float32'),
            'terminal': terminal
        }
        cost = self.fluid_executor.run(
            self.learn_program, feed=feed, fetch_list=[self.cost])[0]  # 训练一次网络
        return cost

# replay_memory.py
# 经验池，用于存储多条经验，实现经验回放。
import random
import collections
import numpy as np


class ReplayMemory(object):
    def __init__(self, max_size):
        self.buffer = collections.deque(maxlen=max_size)

    # 增加一条经验到经验池中
    def append(self, exp):
        self.buffer.append(exp)

    # 从经验池中选取N条经验出来
    def sample(self, batch_size):
        mini_batch = random.sample(self.buffer, batch_size)
        obs_batch, action_batch, reward_batch, next_obs_batch, done_batch = [], [], [], [], []

        for experience in mini_batch:
            s, a, r, s_p, done = experience
            obs_batch.append(s)
            action_batch.append(a)
            reward_batch.append(r)
            next_obs_batch.append(s_p)
            done_batch.append(done)

        return np.array(obs_batch).astype('float32'), \
            np.array(action_batch).astype('float32'), np.array(reward_batch).astype('float32'),\
            np.array(next_obs_batch).astype('float32'), np.array(done_batch).astype('float32')

    def __len__(self):
        return len(self.buffer)

# 训练一个episode
def run_episode(env, agent, rpm):
    total_reward = 0
    obs = env.reset()
    step = 0
    while True:
        step += 1
        action = agent.sample(obs)  # 采样动作，所有动作都有概率被尝试到
        next_obs, reward, done, _ = env.step(action)
        rpm.append((obs, action, reward, next_obs, done))

        # train model
        if (len(rpm) > MEMORY_WARMUP_SIZE) and (step % LEARN_FREQ == 0):
            (batch_obs, batch_action, batch_reward, batch_next_obs,
             batch_done) = rpm.sample(BATCH_SIZE)
            train_loss = agent.learn(batch_obs, batch_action, batch_reward,
                                     batch_next_obs,
                                     batch_done)  # s,a,r,s',done

        total_reward += reward
        obs = next_obs
        if done:
            break
    return total_reward


# 评估 agent, 跑 5 个episode，总reward求平均
def evaluate(env, agent, render=False):
    eval_reward = []
    for i in range(5):
        obs = env.reset()
        episode_reward = 0
        while True:
            action = agent.predict(obs)  # 预测动作，只选最优动作
            obs, reward, done, _ = env.step(action)
            episode_reward += reward
            if render:
                env.render()
            if done:
                break
        eval_reward.append(episode_reward)
    return np.mean(eval_reward)

# 创建环境
env = gym.make('MountainCar-v0')
action_dim = env.action_space.n  # MountainCar-v0: 3
obs_shape = env.observation_space.shape  # MountainCar-v0: (2,)

# 创建经验池
rpm = ReplayMemory(MEMORY_SIZE)  # DQN的经验回放池



# 根据parl框架构建agent
######################################################################
#
#请参考课堂Demo，嵌套Model, DQN, Agent构建 agent
#
######################################################################
model = Model(act_dim=action_dim)
algorithm = DQN(model, act_dim=action_dim, gamma=GAMMA, lr=LEARNING_RATE)
agent = Agent(
    algorithm,
    obs_dim=obs_shape[0],
    act_dim=action_dim,
    e_greed=0.1,  # 有一定概率随机选取动作，探索
    e_greed_decrement=1e-6)  # 随着训练逐步收敛，探索的程度慢慢降低
    
# 加载模型
# save_path = './dqn_model.ckpt'
# agent.restore(save_path)

# 先往经验池里存一些数据，避免最开始训练的时候样本丰富度不够
while len(rpm) < MEMORY_WARMUP_SIZE:
    run_episode(env, agent, rpm)

max_episode = 2000

# 开始训练
episode = 0
while episode < max_episode:  # 训练max_episode个回合，test部分不计算入episode数量
    # train part
    for i in range(0, 50):
        total_reward = run_episode(env, agent, rpm)
        episode += 1

    # test part
    eval_reward = evaluate(env, agent, render=False)  # render=True 查看显示效果
    logger.info('episode:{}    e_greed:{}   test_reward:{}'.format(
        episode, agent.e_greed, eval_reward))

# 训练结束，保存模型
save_path = './dqn_model.ckpt'
agent.save(save_path)