用强化学习来玩Atari游戏(基于Tensorflow的深度Q学习模型)
在之前的博客用Tensorflow Agents实现强化学习DQN_gzroy的博客-CSDN博客中，我用TF-Agents实现了一个深度Q学习模型，并且对小车上山这个环境进行了训练。那么更进一步，我们可以看看能否用深度Q学习解决一些更复杂的问题，例如我们是否能训练一个模型，可以玩Atari的游戏，并取得比人类更好的成绩。
在2015年，DeepMind发表的论文《Human-level control through deep reinforcement learning》,提出了用带经验回放和目标网络的深度Q学习算法，直接从游戏界面中学习。本文将基于论文中提到的算法，来实现对Atari游戏进行学习。
首先加载Atari的环境，这里选择的是Pong这个游戏。环境名称中的deterministic表示每次调用Step的时候会得到4帧后的观测，返回的是这4帧的总奖励值。V4表示执行某个动作后，下一个动作没有限定要和上个动作相同。如果是V0则表示有25%的机率下一个动作和上一个动作相同。
from tf_agents.environments import suite_gym
from tf_agents.environments import tf_py_environment
from tf_agents.networks import q_network
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.trajectories import trajectory
from tf_agents.trajectories import Trajectory
from tf_agents.trajectories import time_step, TimeStep
from tf_agents.specs import tensor_spec
from tqdm import trange
from tf_agents.policies.q_policy import QPolicy
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.initializers import GlorotUniform
from tf_agents.networks.q_network import QNetwork
from tf_agents.agents.dqn import dqn_agent
from tf_agents.utils import common
from tf_agents.policies import random_tf_policy
from tf_agents.drivers import dynamic_step_driver
import random
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
%matplotlib notebook
import os
env_name = 'PongDeterministic-v4'
train_py_env = suite_gym.load(env_name, max_episode_steps=0)
eval_py_env = suite_gym.load(env_name, max_episode_steps=0)
train_env = tf_py_environment.TFPyEnvironment(train_py_env)
eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)加载环境后，我们可以用随机策略试玩一下，看看这个游戏是怎样的
import imageio
import base64
import IPython
def embed_mp4(filename):
"""Embeds an mp4 file in the notebook."""
video = open(filename,'rb').read()
b64 = base64.b64encode(video)
tag = '''
'''.format(b64.decode())
return IPython.display.HTML(tag)
def create_policy_eval_video(policy, filename, num_episodes=1, fps=30):
filename = filename + ".mp4"
with imageio.get_writer(filename, fps=fps) as video:
for _ in range(num_episodes):
time_step = eval_env.reset()
video.append_data(eval_py_env.render())
images = []
while not time_step.is_last():
time_step = eval_env.step(random.randint(action_min,action_max))
video.append_data(eval_py_env.render())
return embed_mp4(filename)
create_policy_eval_video(agent.policy, "trained-agent-pong_random")视频如下：
trained-agent-pong_random
为了能理解游戏界面，需要设计一个卷积神经网络来提取图形的特征。这里需要把游戏最近的4帧灰度画面作为输入，因为这样网络才能理解物体运动的信息。整个神经网络的结构如下：
以下代码定义了一个DQNAgent，这里需要修改原有环境的Timestep的spec，因为现在网络的输入是需要最近的4帧画面，所以输入的维度是110*84*4
input_shape = (110, 84, 4)
target_network_update_freq = 4000
learning_rate = 0.001
gamma = 0.99
time_step_spec = train_env.time_step_spec()
new_observation_spec = tensor_spec.BoundedTensorSpec(input_shape, tf.uint8, 0, 255, 'observation')
new_time_step_spec = time_step.TimeStep(
time_step_spec.step_type,
time_step_spec.reward,
time_step_spec.discount,
new_observation_spec)
q_net = q_network.QNetwork(
new_observation_spec,
train_env.action_spec(),
preprocessing_layers=tf.keras.layers.Rescaling(scale=1./127.5, offset=-1),
conv_layer_params=[(32,8,4), (64,4,2), (64,3,1)],
fc_layer_params=(512,))
global_step = tf.compat.v1.train.get_or_create_global_step()
agent = dqn_agent.DqnAgent(
new_time_step_spec,
train_env.action_spec(),
q_network=q_net,
target_update_tau=1./target_network_update_freq,
target_update_period=1,
gamma=gamma,
epsilon_greedy=0.,
td_errors_loss_fn=common.element_wise_squared_loss,
optimizer=tf.keras.optimizers.RMSprop(learning_rate, 0.95, momentum=0.95, epsilon=0.01),
train_step_counter=global_step)
agent.initialize()定义一个replay buffer用于存储和回放历史经验，这里注意的是在dataset里面需要设置num_steps=2，因为对于深度Q学习，每次需要用两条相邻的经验来进行学习。
replay_memory_size = 100000
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
agent.collect_data_spec,
batch_size=train_env.batch_size,
max_length=replay_memory_size)
dataset = replay_buffer.as_dataset(
num_parallel_calls=3,
sample_batch_size=batch_size,
num_steps=2).prefetch(3)
iterator = iter(dataset)定义一个函数来评估模型的性能，以及用图表描绘模型的训练过程：
def compute_avg_return(environment, policy, num_episodes=10):
total_return = 0.0
for _ in range(num_episodes):
time_step = environment.reset()
episode_return = 0.0
step = 0
images = []
while not time_step.is_last():
step += 1
images, observation = get_observation(images, time_step.observation)
step_type = tf.squeeze(time_step.step_type)
time_step = TimeStep(
time_step.step_type,
time_step.reward,
time_step.discount,
tf.expand_dims(observation, axis=0))
if random.random()<=0.05 or step<=random_initial_steps:
action = random.randint(action_min, action_max)
else:
action = policy.action(time_step).action
time_step = environment.step(action)
episode_return += time_step.reward
total_return += episode_return
avg_return = total_return / num_episodes
return avg_return.numpy()[0]
class Chart:
def __init__(self):
self.fig, self.ax = plt.subplots(figsize = (8, 6))
def plot(self, data, x_name, y_name, hue_name):
self.ax.clear()
sns.lineplot(data=data, x=data[x_name], y=data[y_name], hue=data[hue_name], ax=self.ax)
self.fig.canvas.draw()因为整个模型的训练非常耗时，因此需要定义checkpointer saver来定期保存训练结果以及replay buffer的数据
checkpoint_dir = os.path.join('./', 'checkpoint')
train_checkpointer = common.Checkpointer(
ckpt_dir=checkpoint_dir,
max_to_keep=1,
agent=agent,
policy=agent.policy,
replay_buffer=replay_buffer,
global_step=global_step
)之后我们就可以训练和评估模型了，这里设置了在头10个回合，只搜集经验但不进行训练。10个回合之后，每个回合的头30步采用随机动作，这样可以避免模型记住某些特定的开局。在训练开始epsilon设置为1，然后逐步递减，直到epsilon减少为0.1为止。在训练的时候，需要保存前3帧的图像，并和当前帧一起叠加后，输入到模型中。代码如下：
batch_size = 128
update_freq = 4 # 训练网络的间隔
epsilon = 1. # 初始探索率
min_epsilon = 0.1 # 最终探索率
epsilon_decrease = 9e-7 # 探索减小速度
replay_start_episode = 10 # 开始训练前的经验数
random_initial_steps = 30 # 每个回合开始时随机步数
train_episodes = 400
test_episodes = 2 # 每次验证智能体的回合数
start_episode = 0
end_episode = start_episode+train_episodes
action_min = agent.collect_data_spec.action.minimum
action_max = agent.collect_data_spec.action.maximum
rewards_df = pd.DataFrame(data=None, columns=['episode','reward','type'])
loss_df = pd.DataFrame(data=None, columns=['episode','loss','type'])
chart_reward = Chart()
chart_loss = Chart()
def get_observation(images, observation):
image = tf.squeeze(observation)
image = tf.image.rgb_to_grayscale(image)
image = tf.image.resize(image, [input_shape[0],input_shape[1]])
image = tf.cast(image, tf.uint8)
image = tf.squeeze(image)
if len(images)==0:
images = [image, image, image, image]
images = images[1:]
images.append(image)
observation = tf.stack(images)
observation = tf.transpose(observation, perm=[1,2,0])
return images, observation
for episode in trange(start_episode, end_episode):
time_step = train_env.reset()
step = 0
images = []
episode_reward = 0
total_loss = 0
while not time_step.is_last():
step += 1
images, observation = get_observation(images, time_step.observation)
step_type = tf.squeeze(time_step.step_type)
time_step = TimeStep(time_step.step_type, time_step.reward, time_step.discount, tf.expand_dims(observation, axis=0))
if step<=random_initial_steps:
action = random.randint(action_min, action_max)
action = tf.convert_to_tensor(action, dtype=tf.int64)
else:
epsilon = max(epsilon-epsilon_decrease, min_epsilon)
if random.random()=replay_start_episode and step%update_freq==0:
experience, unused_info = next(iterator)
train_loss = agent.train(experience).loss
total_loss += train_loss
episode_reward += time_step.reward.numpy()[0]
if episode>replay_start_episode:
rewards_df = rewards_df.append({'episode':episode, 'reward':episode_reward, 'type':'train'}, ignore_index=True)
loss_df = loss_df.append({'episode':episode, 'loss':total_loss.numpy()/step, 'type':'train'}, ignore_index=True)
if (episode-replay_start_episode)%test_episodes==0:
avg_return = compute_avg_return(eval_env, agent.policy, test_episodes)
rewards_df = rewards_df.append({'episode':episode, 'reward':avg_return, 'type':'eval'}, ignore_index=True)
chart_reward.plot(rewards_df, 'episode', 'reward', 'type')
chart_loss.plot(loss_df, 'episode', 'loss', 'type')
train_checkpointer.save(global_step) 训练非常花时间，在训练了460回合后，我们可以看一下效果，修改一下之前生成video的代码：
def create_policy_eval_video(policy, filename, num_episodes=1, fps=30):
filename = filename + ".mp4"
with imageio.get_writer(filename, fps=fps) as video:
for _ in range(num_episodes):
time_step = eval_env.reset()
video.append_data(eval_py_env.render())
images = []
while not time_step.is_last():
images, observation = get_observation(images, time_step.observation)
step_type = tf.squeeze(time_step.step_type)
time_step = TimeStep(
time_step.step_type,
time_step.reward,
time_step.discount,
tf.expand_dims(observation, axis=0))
action_step = policy.action(time_step)
time_step = eval_env.step(action_step.action)
video.append_data(eval_py_env.render())
return embed_mp4(filename)
create_policy_eval_video(agent.policy, "trained-agent-pong_1")视频如下：
trained-agent-pong
从视频中可以见到模型已经学习到了一些技巧，可以获得一些分数。
当继续训练到800回合左右时，我们可以看到模型的每个回合的回报已经去到了21，表明已经取得了最好的成绩了：
我们可以把评估结果输出为视频看一下，如下：
trained-agent-pong_862