tensorflow2.0 实现 DQN

DQN 简介

• 表格型方法存储的状态数量有限，当面对围棋或机器人控制这类有数不清的状态的环境时，表格型方法在存储和查找效率上都受局限，DQN的提出解决了这一局限，使用神经网络来近似替代Q表格。
• 本质上DQN还是一个Q-learning算法，更新方式一致。为了更好的探索环境，同样的也采用ε-greedy方法训练。
• 在Q-learning的基础上，DQN提出了两个技巧使得Q网络的更新迭代更稳定。
• 经验回放 Experience
  Replay：主要解决样本关联性和利用效率的问题。使用一个经验池存储多条经验s,a,r,s’，再从中随机抽取一批数据送去训练。
• 固定Q目标 Fixed-Q-Target：主要解决算法训练不稳定的问题。复制一个和原来Q网络结构一样的Target
  Q网络，用于计算Q目标值。

DQN 实践

• 使用DQN解决CartPole问题，移动小车使得车上的摆杆倒立起来

搭建Model、Algorithm、Agent架构

1. Model
   Model用来定义前向(Forward)网络，用户可以自由的定制自己的网络结构。

这里搭建两个相同结构的模型，其中 model 进行训练，target_model 不进行训练，model 训练到一定程度传递权重到 target_model

import tensorflow as tf
from tensorflow.keras import layers,models

class Model:
   def __init__(self,obs_n,act_dim):
       self.act_dim = act_dim
       self.obs_n = obs_n
       self._build_model()
       
   def _build_model(self):
       hid1_size = 128
       hid2_size = 128
       # ------------------ build evaluate_net ------------------
       model = models.Sequential()
       model.add(layers.Input(shape=(self.obs_n)))
       model.add(layers.Dense(hid1_size,activation='relu',name='l1'))
       model.add(layers.Dense(hid2_size,activation='relu',name='l2'))
       model.add(layers.Dense(self.act_dim,name='l3'))
       model.summary()
       self.model = model
       # ------------------ build target_model ------------------
       target_model = models.Sequential()
       target_model.add(layers.Input(shape=(self.obs_n)))
       target_model.add(layers.Dense(hid2_size,activation='relu',name='l1'))
       target_model.add(layers.Dense(hid2_size,activation='relu',name='l2'))
       target_model.add(layers.Dense(self.act_dim,name='l3'))
       target_model.summary()
       self.target_model = target_model

2. Algorithm
   Algorithm 定义了具体的算法来更新前向网络(Model)，也就是通过定义损失函数来更新Model，和算法相关的计算都放在algorithm中。

在计算 loss 函数时，y = Q(s,a)=r+gamma*max Q’(s’,a’)
在 DQN 中是通过网络对于各个 Q 值进行估计的，即：
Q(s,a;θ)≈Q’(s,a)
根据 Bellman 方程得到的 Q 函数的估计以及通过网络估计的 Q 值函数就存在一个差异。因此，可以在网络中引入一个损失函数，即
L(θ) = 1/n [(y-Q(s,a;θ))^2]

Q(s,a;θ) 在 tensorflow 中可以这样计算

# obs,action 为存储在经验池中的数据
pred_value = self.model.predict(obs)	
enum_action = list(enumerate(action))
pred_action_value = tf.gather_nd(pred_value,indices=enum_action)

import tensorflow as tf

class DQN:
    def __init__(self,model,gamma=0.9,learnging_rate=0.01):
        self.model = model.model
        self.target_model = model.target_model
        self.gamma = gamma
        self.lr = learnging_rate
        # --------------------------训练模型--------------------------- # 
        self.model.optimizer = tf.optimizers.Adam(learning_rate=self.lr)
        self.model.loss_func = tf.losses.MeanSquaredError()
        # self.model.train_loss = tf.metrics.Mean(name="train_loss")
        # ------------------------------------------------------------ #
        self.global_step = 0
        self.update_target_steps = 200  # 每隔200个training steps再把model的参数复制到target_model中


    def predict(self, obs):
        """ 使用self.model的value网络来获取 [Q(s,a1),Q(s,a2),...]
        """
        return self.model.predict(obs)

    def _train_step(self,action,features,labels):
        """ 训练步骤
        """
        with tf.GradientTape() as tape:
            # 计算 Q(s,a) 与 target_Q的均方差，得到loss
            predictions = self.model(features,training=True)
            enum_action = list(enumerate(action))
            pred_action_value = tf.gather_nd(predictions,indices=enum_action)
            loss = self.model.loss_func(labels,pred_action_value)
        gradients = tape.gradient(loss,self.model.trainable_variables)
        self.model.optimizer.apply_gradients(zip(gradients,self.model.trainable_variables))
        # self.model.train_loss.update_state(loss)
    def _train_model(self,action,features,labels,epochs=1):
        """ 训练模型
        """
        for epoch in tf.range(1,epochs+1):
            self._train_step(action,features,labels)

    def learn(self,obs,action,reward,next_obs,terminal):
        """ 使用DQN算法更新self.model的value网络
        """
        # 每隔200个training steps同步一次model和target_model的参数
        if self.global_step % self.update_target_steps == 0:
            self.replace_target()

        # 从target_model中获取 max Q' 的值，用于计算target_Q
        next_pred_value = self.target_model.predict(next_obs)
        best_v = tf.reduce_max(next_pred_value,axis=1)
        terminal = tf.cast(terminal,dtype=tf.float32)
        target = reward + self.gamma * (1.0 - terminal) * best_v

        # 训练模型
        self._train_model(action,obs,target,epochs=1)
        self.global_step += 1

    def replace_target(self):
        '''预测模型权重更新到target模型权重'''
        self.target_model.get_layer(name='l1').set_weights(self.model.get_layer(name='l1').get_weights())
        self.target_model.get_layer(name='l2').set_weights(self.model.get_layer(name='l2').get_weights())
        self.target_model.get_layer(name='l3').set_weights(self.model.get_layer(name='l3').get_weights())

3. Agent
   Agent 负责算法与环境的交互，在交互过程中把生成的数据提供给Algorithm来更新模型(Model)，数据的预处理流程也一般定义在这里。

import numpy as np
import tensorflow as tf

class Agent:
    def __init__(self,act_dim,algorithm,e_greed=0.1,e_greed_decrement=0):
        self.act_dim = act_dim
        self.algorithm = algorithm
        self.e_greed = e_greed
        self.e_greed_decrement = e_greed_decrement


    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 = tf.expand_dims(obs,axis=0)
        action = self.algorithm.model.predict(obs)
        return np.argmax(action)

4. ReplayMemory
   经验池：用于存储多条经验，实现 经验回放。

import random
import collections
import numpy as np

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

    def append(self,exp):
        self.buffer.append(exp)

    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('int32'), 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)

5. Training && Test（训练&&测试）

import gym
import numpy as np
from model import Model
from algorithm import DQN
from agent import Agent
from replay_memory import ReplayMemory

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一批数据出来
LEARNING_RATE = 0.001  # 学习率
GAMMA = 0.99  # reward 的衰减因子，一般取 0.9 到 0.999 不等

def run_episode(env,algorithm,agent,rpm):
    step = 0
    total_reward = 0
    obs = env.reset()

    while True:
        step += 1
        action = agent.sample(obs)
        next_obs,reward,done,_ = env.step(action)
        rpm.append((obs,action,reward,next_obs,done))

        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)
            algorithm.learn(batch_obs,batch_action,batch_reward,batch_next_obs,batch_done)
            
        obs = next_obs
        total_reward += reward
        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)

6. 创建环境和Agent，创建经验池，启动训练，保存模型

def main():
    env = gym.make(
        'CartPole-v0'
    )
    action_dim = env.action_space.n         # 2
    obs_shape = env.observation_space.shape # (4,)

    rpm = ReplayMemory(MEMORY_SIZE)         # DQN的经验回放池
    model = Model(obs_shape[0],action_dim)
    algorithm = DQN(model,gamma=GAMMA,learnging_rate=LEARNING_RATE)
    agent = Agent(action_dim,algorithm,e_greed=0.1,e_greed_decrement=1e-6)

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

    max_episode = 2000

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

        # 测试
        eval_reward = evaluate(env,agent,render=True)
        print('episode:{}   e_greed:{}   Test reward:{}'.format(episode,agent.e_greed,eval_reward))

    # 训练结束，保存模型
    save_path = './dqn_model.h5'
    model.model.save(save_path)
    env.close()

注意事项

1. target_model 要记得进行传递参数
2. 收敛不了，要检查学习率