重复造轮子没啥意义，但是手打一遍还是有收获。

import tensorflow as tf
import random
from collections import namedtuple
import math
import numpy as np

def huber_loss(y_true, y_pred, clip_delta=1.0):
    '''
    Huber_loss:回归loss函数，专用于DQN算法
    :param y_true: 真实集
    :param y_pred: 预测集
    :param clip_delta: 超参数，决定平方结果随误差变化的速度
    :return: loss值
    '''
    error = y_true - y_pred
    cond = tf.keras.backend.abs(error) < clip_delta

    squared_loss = 0.5 * tf.keras.backend.square(error)
    linear_loss = clip_delta * (tf.keras.backend.abs(error) - 0.5 * clip_delta)

    return tf.where(cond, squared_loss, linear_loss)

def huber_loss_mean(y_true, y_pred, clip_delta=1.0):
  return tf.keras.backend.mean(huber_loss(y_true, y_pred, clip_delta))


class BasicPool(object):
    def __init__(self, capacity):
        '''
        基于python内置环境，开辟数据池空间。理论上更快。
        :param capacity: 数据池容量
        '''
        self.capacity = capacity
        self.memory = []
        self.position = 0
        self.Transition = namedtuple('Transition',('state', 'action', 'next_state', 'reward'))

    def push(self, *args):
        '''
        存入任意参数
        '''
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = self.Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

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


class BasicNet(tf.keras.Model):
    '''
    简单的卷积+FC，双网络的baseline
    '''
    def __init__(self, outputs):
        super(BasicNet, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(filters=16, kernel_size=(5,5), strides=(2,2), padding="same")  # 二维卷积，卷积核大小为5
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.conv2 = tf.keras.layers.Conv2D(filters=32, kernel_size=(5,5), strides=(2,2), padding="same")
        self.bn2 = tf.keras.layers.BatchNormalization()
        self.conv3 = tf.keras.layers.Conv2D(filters=32, kernel_size=(5,5), strides=(2,2), padding="same")
        self.bn3 = tf.keras.layers.BatchNormalization()
        self.Flatten = tf.keras.layers.Flatten()
        self.FC2 = tf.keras.layers.Dense(units=outputs, activation=tf.keras.activations.softmax)

    def call(self, inputs, training=None, mask=None):
        x = tf.nn.relu(self.bn1(self.conv1(inputs)))
        x = tf.nn.relu(self.bn2(self.conv2(x)))
        x = tf.nn.relu(self.bn3(self.conv3(x)))
        x = self.Flatten(x)
        return self.FC2(x)

class ArmAgent:
    def __init__(self, n_actions, screen_height, screen_weight,
                 batch_size=128, reward_decay=0.99, e_start=0.9,
                 e_end=0.05, e_decay=200, replace_target_iter=10,
                 memory_size=8000):
        self.BATCH_SIZE = batch_size
        self.GAMMA = reward_decay
        self.EPS_START = e_start
        self.EPS_END = e_end
        self.EPS_DECAY = e_decay
        self.TARGET_UPDATE = replace_target_iter
        self.NAction = n_actions
        self.H = screen_height
        self.W = screen_weight
        self.Optimizer = tf.keras.optimizers.RMSprop(learning_rate=0.001)
        self.build_net()
        self.Memory = BasicPool(memory_size)
        self.steps = 0

    def build_net(self):
        self.eval_net = BasicNet(self.NAction)
        self.target_net = BasicNet(self.NAction)
        for layers in self.eval_net.layers:
            layers.trainable = True
        self.eval_net.build(input_shape=(None, self.H, self.W, 3))
        self.target_net.build(input_shape=(None, self.H, self.W, 3))
        self.eval_net.summary()
        self.target_net.summary()

    def learn(self):
        print("learning")
        if len(self.Memory) < self.BATCH_SIZE:
            return

        with tf.GradientTape() as tape:
            transitions = self.Memory.sample(self.BATCH_SIZE)  # 一个batch进行计算
            batch = self.Memory.Transition(*zip(*transitions))
            non_final_mask = tf.constant(tuple(map(lambda s: s is not None, batch.next_state)), dtype=tf.bool)
            # bool形式的数组作为索引
            non_final_next_states = tf.concat(tuple([s for s in batch.next_state
                                                     if s is not None]), axis=0)
            state_batch = tf.concat(batch.state, axis=0)
            action_batch = tf.concat(batch.action, axis=0)
            reward_batch = tf.concat(batch.reward, axis=0)
            # 计算loss
            for s in self.eval_net.trainable_variables:
                tape.watch(s)
            first_index = tf.constant(np.linspace(0, self.BATCH_SIZE - 1, self.BATCH_SIZE, dtype=np.int32),
                                      dtype=tf.int32)
            index_action = tf.stack((first_index, np.squeeze(action_batch)), axis=1)
            state_action_values = tf.gather_nd(self.eval_net(state_batch, training=True), index_action)
            next_state_values = tf.constant([], dtype=tf.float32)
            copy_array = tf.reduce_max(self.target_net(non_final_next_states), axis=1)
            index_copy = 0
            for a in range(non_final_mask.shape[0]):
                if bool(non_final_mask[a]):
                    next_state_values = tf.concat((next_state_values, tf.expand_dims(copy_array[index_copy], axis=0)), axis=0)
                    index_copy += 1
                else:
                    next_state_values = tf.concat((next_state_values, tf.constant([0.])), axis=0)
            # 计算预期的Q值
            expected_state_action_values = tf.multiply(next_state_values, self.GAMMA) + reward_batch
            loss = huber_loss_mean(expected_state_action_values, state_action_values)

        print("loss:{}".format(loss))
        gradients = tape.gradient(loss, self.eval_net.trainable_variables)
        self.Optimizer.apply_gradients(grads_and_vars=zip(gradients, self.eval_net.trainable_variables))

    def choose_action(self, state):
        sample = random.random()
        eps_threshold = self.EPS_END + (self.EPS_START - self.EPS_END) * math.exp(-1. * self.steps / self.EPS_DECAY)  # eps阈值
        self.steps += 1
        if sample > eps_threshold:
            return np.argmax(self.eval_net(state))
        else:
            return np.array([[random.randrange(self.NAction)]], dtype=np.int64)

    def store_data(self, state, action, next_state, reward):
        if type(action) != np.ndarray:
            action = np.array([[action]])
        self.Memory.push(state, action, next_state, np.array(reward, dtype=np.float32))

    def update_target(self):
        modelWeights = self.eval_net.trainable_weights
        targetModelWeights = self.target_net.trainable_weights

        for i in range(len(targetModelWeights)):
            targetModelWeights[i].assign(modelWeights[i])