强化学习 Reinforcement Learning（六）——好奇心驱动的强化学习

最新推荐文章于 2024-10-06 10:44:39 发布

ZiSeoi

最新推荐文章于 2024-10-06 10:44:39 发布

阅读量4.3k

点赞数 12

分类专栏：强化学习整活小项目

本文链接：https://blog.csdn.net/wzduang/article/details/112533271

版权

强化学习同时被 2 个专栏收录

9 篇文章 7 订阅

订阅专栏

整活小项目

7 篇文章 0 订阅

订阅专栏

强化学习 Reinforcement Learning（六）——好奇心驱动的强化学习

写在前面

本文将讲述关于好奇心机制的主要内容，向读者展示什么是好奇心机制，最后复现 ICM ，并与我们熟悉的 DQN 相结合

噢噢，差点忘了，本文的代码全部基于 paddle2.0rc 版本

第一部分：好奇心机制与稀疏奖励

标准的强化学习算法在对智能体反馈稀疏的环境中表现不佳，更关键的是，这种环境在现实世界中非常常见。

举例来讲，假如您正在尝试学习如何在一个迷宫般的大超市里找到您想要买的东西。您找来找去，但就是找不到。

如果每一步都得不到“胡萝卜”，连“大棒”也没有，无法得到反馈，您可能每走一步都会怀疑自己走的方向是否正确。

在这种情况下，要如何避免原地打转呢？这时候就可以利用好奇心啦。好奇心会驱使你走进看似陌生的产品区，最终寻找到想要买的东西。

在逛超市时，您会尝试预测“我现在在肉类区，所以我觉得拐角的区域是水产区”。

如果预测错了，您会感到意外“哟，没想到”，并因此得到好奇心的满足奖励。这让您以后更愿意探索新的地点，只为看看预期与现实是否一致。

2017 年，OpenAI 和 UC 伯克利发表了 Large-Scale Study of Curiosity-Driven Learning

本文最大的亮点是直接剔除了死亡（terminaton of a game, death）所提供的信息，即当 agent 死亡，直接让 agent 恢复到出生点（变成了infinite horizon episode），而不终结episode或者给予任何extrinsic惩罚。

同时由于 dynamic-based curiosity 的 agent 很熟悉出生点的情况，会失去兴趣（获得很少 intrinsic reward ）从而间接驱使 agent 不要死亡。

如果环境中的 transition 是完全随机的，那么即使有一个完美的 dynamics model，期望的回报也会是 transition 的熵，agent 会寻找熵最大的 transition（因为不确定性越大，熵越大，agent 越感兴趣）。

意思就是如果 agent 正在看电视，有一个遥控器可以更换频道。由于频道切换会带来巨大的奖励，因为每次更换都是不可预测和令人意外的。

agent 就会被吸引在电视机前而停止对环境的探索，感觉自己 high 到不行（狗头）

即使在所有可观看的频道都循环了一次后，随机的频道选择也能确保每次更换都令人意外，因为智能体会预测更换频道后电视会播放什么节目，而这个预测很可能出错，从而导致意外的出现。

重要的是，即使智能体已经看过了每个频道播出的每个节目，这种更换依然是不可预测的。

因此，具有基于意外之好奇心的智能体最终会一直留在电视前，而不去寻找奖励非常丰厚的物品，这种情况与拖延很相似。那么，如何定义好奇心才能避免这种行为？

第二部分：构建特征空间

由于好奇心机制如果以原始环境信息作为输入具有各种各样的问题，所以首先我们需要一个良好的特征空间

在深入描述模型之前，我们先必须问自己，根据我们当前的状态和行动，智能体该如何预测下一个状态？

我们知道可以将好奇心定义为给定当前状态 $s_t$ 和行动 $a_t$ 的预测新状态 $s_{t + 1}$ 与真实新状态之间的误差。

但是，请记住，大多数情况下，我们的状态是4帧图像的堆栈。这意味着我们需要找到一种方法来预测下一个叠帧（Frame Stacks），这有两个难点：

首先，很难直接预测像素，想象如果你以图像信息作为输入，那么你的预测也就需要生成完整的图像
其次，这其实是一种错误的做法。想象一下，如果你需要在微风中研究树叶的运动。我们很难对微风进行建模，预测每片叶子的像素位置就更困难了。

这样的后果是，由于总是会有一个很大的像素预测误差，即使叶子的运动不是智能体动作的结果，智能体也无法对叶子的运动做出控制，智能体也会一直好奇。

因此，我们需要将原始感官输入（像素阵列）转换为仅包含相关信息的特征空间，而不是在原始感官空间（像素）中进行预测。

可以通过定义以下三点规则来构建一个好的特征空间：

要对可由智能体控制的物体进行建模。
还要对智能体无法控制但可能对其产生影响的事物进行建模。
不要对智能体无法控制且对其没有影响的事物建模（因此不受影响）。

即智能体能够控制所有能对其产生影响的事物，同时不再关心无关的事物。

由此，所需的嵌入空间应该：

空间紧凑（去除不相关部分）。
保留有关观察的充分信息。
稳定——非固定奖励使强化学习体难以学习。

第三部分：内在好奇心模块（ICM）

针对上面提到的问题，次年发表的 Curiosity-driven Exploration by Self-supervised Prediction，提出一种有效的内在好奇心模块

内在好奇心模块是帮助我们产生好奇心奖励的系统。它由三部分神经网络组成。

由于直接预测图象非常困难，我们需要的不是从原始感觉空间（像素）进行预测，而是将输入转换为特征向量，其中仅表示与智能体执行的动作相关的信息。

这里的 Encoder 就是负责提取输入信号中的有用信息，并嵌入特征空间。

那么如何让 Encoder 学习特征空间？

ICM 采用自监督的方法，Inverse Model 根据两个相邻的 state 来推断智能体所选的 action，即给定其当前状态和下一个状态 $s_t$ 和 $s_{t+ 1}$ 预测学习体行为 $â_t$ 。

然后利用 inverse prediction error （ $â_t$ 和 $a_t$ 之间的差距）来训练Encoder。

这样处理的好处是让 Encoder 输出的特征空间限于智能体能够控制、改变的空间里。

回到刚刚 agent 对“看电视上瘾”的问题上，当 Encoder 无法从电视中获得奖励（因为 agent 发现无法通过 action 改变电视上播放的内容），Encoder 就不会在意电视屏幕上的内容了，也就不会提供好奇心奖励。

这里有以下几点需要注意：

Forward Model 的 Prediction error（ $\hat{s}_{t + 1}$ 与 $s_{t + 1}$ 的损失）只用来训练 Forward model ，而不用于训练 Encoder
Inverse Model 的 Inverse prediction error（ $\hat{a}_{t + 1}$ 与 $a_{t + 1}$ 的损失）既用来训练 Inverse model ，也用来 Encoder
同时 Forward Model 的 Prediction error 同时也是 intrinsic reward （好奇心奖励）来训练我们的 agent

举个例子，假设我们在看球赛，

突然因为转播信号不好，画面没了。

当重新看到画面的时候北京已经开始庆祝了。

那么我们能够做出的推测就是北京赢了。

如果画面切回时我们看到江苏在庆祝了，那就是江苏逆风翻盘了。

注意，帮助我们做出推断的并不是看台上的观众或者场下的广告牌，虽然它们也出现在了画面里面。

最重要的信息是球员和解说等等。我们的大脑对这些信息进行了编码才做出了准确的推断。

当然，其实直接看记分牌就行了：）

第三部分 ICM-DQN

这里将 ICM 与 DQN 结合，组一个 ICM-DQN 来试试马里奥兄弟

安装依赖

pip install atari_py
pip install gym
pip install box2d-py
pip install scikit-image
pip install gym-super-mario-bros

导入需要用到的库

import gym, random, pickle, os.path, math, glob
from skimage import transform

from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym_super_mario_bros.actions import RIGHT_ONLY,SIMPLE_MOVEMENT,COMPLEX_MOVEMENT

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import collections
from collections import deque

import cv2 as cv

import paddle
import paddle.nn as nn
from paddle.nn import Conv2D, Linear
import numpy as np
import paddle.nn.functional as F
from visualdl import LogWriter

Frame Stack 定义

本文将输入维度限制在 100*100 ，同时叠合多帧图像，用于获取运动信息

def stack_frames(stacked_frames, state, is_new_episode):
    # Preprocess frame
    frame = preprocess_frame(state)
    
    if is_new_episode:
        # Clear our stacked_frames
        stacked_frames = deque([np.zeros((100,100), dtype=np.int) for i in range(4)], maxlen=4)
        
        # Because we're in a new episode, copy the same frame 4x
        stacked_frames.append(frame)
        stacked_frames.append(frame)
        stacked_frames.append(frame)
        stacked_frames.append(frame)
        
        # Stack the frames
        stacked_state = np.stack(stacked_frames, axis=0)
        
    else:
        # Append frame to deque, automatically removes the oldest frame
        stacked_frames.append(frame)

        # Build the stacked state (first dimension specifies different frames)
        stacked_state = np.stack(stacked_frames, axis=0) 
    
    return stacked_state, stacked_frames

def preprocess_frame(frame):
    # Crop the screen (remove the roof because it contains no information)
    frame = cv.cvtColor(frame,cv.COLOR_RGB2GRAY)
    cropped_frame = np.array(cv.resize(frame, (100, 100)))
    # Normalize Pixel Values
    normalized_frame = cropped_frame / 255.0
    # Resize
    preprocessed_frame = transform.resize(normalized_frame, [100,100])   
    return preprocessed_frame

stacked_frames  =  deque([np.zeros((100,100), dtype=np.int) for i in range(4)], maxlen=4) # 定义 frame stack 全局变量

定义模型

定义 DQN 与 ICM 结构体

class DQN(nn.Layer):

    def __init__(self, h, w, in_channels=4, num_actions=5):
        super(DQN, self).__init__()
        self.conv1 = Conv2D(in_channels=in_channels, out_channels=16, kernel_size=8, stride=2)
        self.conv2 = Conv2D(in_channels=16, out_channels=32, kernel_size=4, stride=2)
        self.conv3 = Conv2D(in_channels=32, out_channels=32, kernel_size=3, stride=2)

        def conv2d_size_out(size, kernel_size = 5, stride = 2):
            return (size - (kernel_size - 1) - 1) // stride  + 1 # 计算卷积后的特征图面积
        
        convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w, 8), 4), 3)
        convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h, 8), 4), 3)
        linear_input_size = convw * convh * 32
        self.head = Linear(linear_input_size, num_actions)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = paddle.reshape(x, [x.shape[0], -1])
        return self.head(x)

class ICM(nn.Layer):
    def __init__(self, h, w, in_channels=4, num_actions=5):
        super(ICM, self).__init__()
        self.conv1 = Conv2D(in_channels=in_channels, out_channels=32, kernel_size=5, stride=2)
        self.conv2 = Conv2D(in_channels=32, out_channels=64, kernel_size=3, stride=2)
        self.conv3 = Conv2D(in_channels=64, out_channels=64, kernel_size=3, stride=2)

        def conv2d_size_out(size, kernel_size = 5, stride = 2):
            return (size - (kernel_size - 1) - 1) // stride  + 1 # 计算卷积后的特征图面积
        
        convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w, 5), 3), 3)
        convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h, 5), 3), 3)

        self.fc4 = Linear(convw * convh * 64, 512)
                
        self.pred_module1 = Linear(512 + num_actions, 256)
        self.pred_module2 = Linear(256, 512)
            
        self.invpred_module1 = Linear(512 + 512, 256) # 用于生成特征状态
        self.invpred_module2 = Linear(256, num_actions) # 用于生成（反向）动作

    def Encoder(self, x): # 编码器结构，负责给输入降维，尤其是图像结构，生成器无法生成太高维的状态量
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = paddle.reshape(x, [x.shape[0], -1])
        x = F.relu(self.fc4(x))
        return x
    
    def forward(self, x): # Encoder 的前向结构：
        # 得到特征
        feature = self.Encoder(x)
        return feature

    def Predict(self, feature_x, a_vec):
        # 输入当前的状态特征和动作(以 one-hot 的形式), 生成下一特征状态(编码后的特征)
        pred_s_next = paddle.concat([feature_x, paddle.cast(a_vec, "float32")], axis = -1).detach()
        pred_s_next = self.pred_module1(pred_s_next)
        pred_s_next = F.relu(pred_s_next)
        pred_s_next = self.pred_module2(pred_s_next) # 状态维度 512
        return pred_s_next
    
    def Invpred(self,feature_x, feature_x_next):
        # 相反动作预测: 输入现态和次态, 输出预测动作(one-hot 形式)
        pred_a_vec = paddle.concat([feature_x, feature_x_next], axis = -1)
        pred_a_vec = self.invpred_module1(pred_a_vec)
        pred_a_vec = F.relu(pred_a_vec)
        pred_a_vec = self.invpred_module2(pred_a_vec)
        return F.softmax(pred_a_vec, axis = -1)
    
    def Get_Full(self, x, x_next, a_vec):
        # 得到所有需要的特征(下一特征状态，下一动作)
        feature = self.Encoder(x)
        feature_next = self.Encoder(x_next)

        pred_s_next = self.Predict(feature, a_vec) # 预测下一状态
        pred_a_vec = self.Invpred(feature, feature_next) # （反向）动作预测：该动作

        return pred_s_next, pred_a_vec, feature_next

经验池定义

由于我们使用的是 DQN，需要经验回放，故在此定义经验池

class ReplayMemory(object):
    def __init__(self, memory_size=1000):
        self.buffer = []
        self.memory_size = max_size
        self.next_idx = 0

    def push(self, state, action, reward, next_state, done):
        data = (state, action, reward, next_state, done)
        if len(self.buffer) <= self.memory_size: # 经验池未满
            self.buffer.append(data)
        else: # 经验池已满
            self.buffer[self.next_idx] = data
        self.next_idx = (self.next_idx + 1) % self.memory_size

    def sample(self, batch_size):
        states, actions, rewards, next_states, dones = [], [], [], [], []
        for i in range(batch_size):
            idx = random.randint(0, self.size() - 1)
            data = self.buffer[idx]
            state, action, reward, next_state, done= data
            states.append(state)
            actions.append(action)
            rewards.append(reward)
            next_states.append(next_state)
            dones.append(done)
            
            
        return np.concatenate(states), actions, rewards, np.concatenate(next_states), dones
    
    def size(self):
        return len(self.buffer)

定义 ICM-DQN agent

结构体 agent 里，定义的方法有计算 Q 值表、包含随机探索的动作采样与、模型训练（计算 td 损失）

class ICM_DQN_Agent(object): # HWC格式输入
    def __init__(self, obs_dim = [], in_channels = 1, action_space = [], max_size = 10000, epsilon  = 1, lr = 1e-4, 
                 forward_scale = 0.8, inverse_scale = 0.2, Qloss_scale = 0.1, intrinsic_scale= 1, use_extrinsic = True):
        self.epsilon = epsilon
        self.action_space = action_space
        # ICM 的参数
        self.forward_scale = forward_scale # 前向预测模型损失函数的比例为 0.8        
        self.inverse_scale = inverse_scale # 反向预测模型损失函数的比例为 0.2
        self.Qloss_scale = Qloss_scale # Q 值损失函数比例为1
        self.intrinsic_scale = intrinsic_scale # 内在奖励的比例为 1
        self.use_extrinsic = use_extrinsic # 是否使用外在奖励, 如果为 False , 模型只接受来自 ICM 的损失
        
        self.memory_buffer = ReplayMemory(max_size) # 经验池
        self.DQN = DQN(obs_dim[1], obs_dim[2], in_channels = in_channels, num_actions = action_space.n)
        self.DQN_target = DQN(obs_dim[1], obs_dim[2], in_channels = in_channels, num_actions = action_space.n)
        self.DQN_target.load_dict(self.DQN.state_dict())
        self.ICM = ICM(obs_dim[1], obs_dim[2], in_channels = in_channels, num_actions = action_space.n)
        self.clip = paddle.nn.ClipGradByNorm(clip_norm=1.0)
        self.optimizer = paddle.optimizer.Adam(learning_rate=lr, parameters=list(self.DQN.parameters())+list(self.ICM.parameters()), grad_clip=self.clip)
        self.DQN.train()
        self.DQN_target.train()
    
    def observe(self, frame):
        # 帧到 tensor 
        frame = np.expand_dims(frame, axis=0).astype('float32')
        state = paddle.to_tensor(frame, dtype='float32')
        return state

    def value(self, state):
        q_values = self.DQN(state)
        return q_values
    
    def act(self, state, epsilon = None):
        """
        包含随即探索的动作采样
        """
        if epsilon is None: epsilon = self.epsilon

        q_values = self.value(state).detach().numpy()
        if random.random() < epsilon:
            aciton = random.randrange(self.action_space.n)
        else:
            aciton = q_values.argmax(1)[0]
        return aciton
    
    def compute_td_loss(self, states, actions, rewards, next_states, is_done, gamma=0.99):
        """ 计算 td 损失 """
        actions = paddle.to_tensor(actions, dtype='int64')    # shape: [batch_size]
        rewards = paddle.to_tensor(rewards, dtype='float32')  # shape: [batch_size]
        is_done = paddle.to_tensor(is_done, dtype='bool')  # shape: [batch_size]

        # 得到现态的 Q 值表
        predicted_qvalues = self.DQN(states)
        
        # 得到 ICM 的结果
        a_vec = F.one_hot(actions, num_classes = self.action_space.n) # 将动作转化为 one-hot 向量
        pred_s_next, pred_a_vec, feature_next = self.ICM.Get_Full(states, next_states, a_vec)

        # 计算前向模型和反向动作预测的损失
        forward_loss = F.mse_loss(pred_s_next, feature_next.detach(), reduction='none')
        inverse_pred_loss = F.cross_entropy(pred_a_vec, actions.detach(), reduction='none')
        
        # 计算奖励
        intrinsic_rewards = self.intrinsic_scale * forward_loss.mean(-1)
        total_rewards = intrinsic_rewards.numpy()
        total_rewards = paddle.to_tensor(total_rewards)
        if self.use_extrinsic:
            total_rewards += rewards
            
        # 选择动作对应的 Q 值
        predicted_qvalues_for_actions = paddle.sum(
            paddle.multiply(a_vec, predicted_qvalues), axis=1)

        # 计算现态所有动作的 Q 值
        predicted_next_qvalues = self.DQN_target(next_states)

        # 计算次态所有动作的 Q 值
        next_state_values =  predicted_next_qvalues.max(-1)[0]
        
        # Q 学习核心公式
        target_qvalues_for_actions = total_rewards + gamma *next_state_values 

        # 对于最后的状态，可以进行一些简化Q(s,a) = r(s,a), 因为 s' 不存在
        target_qvalues_for_actions = paddle.where(
            is_done, total_rewards, target_qvalues_for_actions)
        
        # Lasso 损失
        Q_loss = F.smooth_l1_loss(predicted_qvalues_for_actions, target_qvalues_for_actions.detach())
        loss = self.Qloss_scale*Q_loss + self.forward_scale*forward_loss.mean() + self.inverse_scale* inverse_pred_loss.mean()

        return loss, Q_loss.numpy().astype("float32"), forward_loss.mean().numpy().astype("float32"), \
                inverse_pred_loss.mean().numpy().astype("float32"), intrinsic_rewards.mean().numpy().astype("float32")
    
    def sample_from_buffer(self, batch_size):
        # 从经验池中选取数据
        states, actions, rewards, next_states, dones = [], [], [], [], []
        for i in range(batch_size):
            idx = random.randint(0, self.memory_buffer.size() - 1)
            data = self.memory_buffer.buffer[idx]
            frame, action, reward, next_frame, done= data
            states.append(self.observe(frame))
            actions.append(action)
            rewards.append(reward)
            next_states.append(self.observe(next_frame))
            dones.append(done)
        return paddle.concat(states), actions, rewards, paddle.concat(next_states), dones

    def learn_from_experience(self, batch_size):
        if self.memory_buffer.size() > batch_size:
            states, actions, rewards, next_states, dones = self.sample_from_buffer(batch_size)
            td_loss, Q_loss, forward_loss, inverse_pred_loss,intrinsic_rewards = self.compute_td_loss(states, actions, rewards, next_states, dones)
            self.optimizer.clear_grad()
            td_loss.backward()
            self.optimizer.step()
            return(td_loss.numpy().astype("float32"),Q_loss, forward_loss, inverse_pred_loss,intrinsic_rewards)
        else:
            return(0,0,0,0,0)

模型训练

这里把外在的奖励完全拿掉了，看一下仅凭好奇心我们的 agent 能拿多少分

# 训练
env = gym_super_mario_bros.make('SuperMarioBros-v2')
env = JoypadSpace(env, SIMPLE_MOVEMENT)


gamma = 0.99
epsilon_max = 1
epsilon_min = 0.01
eps_decay = 50000
frames = int(1e8)
learning_rate = 1e-4
max_size = 3*int(1e5)
update_tar_interval = 1000
batch_size = 128
print_interval = 1000
log_interval = 1000
learning_start = 5000 # 10000
win_reward = 18     # Pong-v4
win_break = True

# ICM 的参数
forward_scale = 0.8 # 前向预测模型损失函数的比例, 默认为 0.8
inverse_scale = 0.2 # 反向预测模型损失函数的比例为, 默认为 0.2
Qloss_scale = 1 # Q 值损失函数比例为1
intrinsic_scale = 100 # 内在奖励的比例为 100
use_extrinsic = False # 是否使用外在奖励, 如果为 False , 模型只接受来自 ICM 的损失

action_space = env.action_space
action_dim = env.action_space.n
# state_dim = env.observation_space.shape
state_dim = [4, 100, 100]
state_channel = state_dim[0]

agent = ICM_DQN_Agent(obs_dim = state_dim, in_channels = state_channel, action_space= action_space, lr = learning_rate,
                    forward_scale = forward_scale, inverse_scale =inverse_scale, Qloss_scale = Qloss_scale, intrinsic_scale= intrinsic_scale,
                     use_extrinsic = use_extrinsic)

frame = env.reset()
frame, stacked_frames = stack_frames(stacked_frames, frame, True)

episode_reward = 0
intrinsic_reward = 0
all_rewards = []
losses = []
episode_num = 0
is_win = False

log_writer = LogWriter(logdir = "./log/ICM_DQN_Pong", comment= "good_makeatari")

# e-greedy 衰减
epsilon_by_frame = lambda frame_idx: epsilon_min + (epsilon_max - epsilon_min) * math.exp(
            -1. * frame_idx / eps_decay)

loss = 0
Q_loss_record = 0
forward_loss_record=0
inverse_pred_loss_record=0
intrinsic_rewards_rec = 0
for i in range(frames):
    epsilon = epsilon_by_frame(i)
    state_tensor = agent.observe(frame)
    action = agent.act(state_tensor, epsilon)
    
    next_frame, reward, done, _ = env.step(action)
    next_frame, stacked_frames = stack_frames(stacked_frames, next_frame, False)
    
    all_rewards.append(reward)
    
    episode_reward += reward
    agent.memory_buffer.push(frame, action, reward, next_frame, done)
    frame = next_frame
    
    if agent.memory_buffer.size() >= learning_start:
        loss, Q_loss_record, forward_loss_record, inverse_pred_loss_record, intrinsic_rewards_rec = agent.learn_from_experience(batch_size)
        losses.append(loss)
        intrinsic_reward += intrinsic_rewards_rec

    if i % print_interval == 0:
        print("frames: %5d, reward: %5f, total_loss: %4f, forward_loss: %4f, inverse_pred_loss: %4f, Q_loss: %4f, intrinsic_rewards: %4f, epsilon: %5f, episode: %4d" % 
              (i, np.mean(all_rewards[-10:]), loss, forward_loss_record, inverse_pred_loss_record, Q_loss_record, intrinsic_rewards_rec, epsilon, episode_num))
        log_writer.add_scalar("Temporal Difference Loss", value=loss, step=i)
        log_writer.add_scalar("Intrinsic Reward", value=intrinsic_reward, step=i)
        log_writer.add_scalar("Extrinsic Reward", value=episode_reward, step=i)
        log_writer.add_scalar("Mean Reward", value=np.mean(all_rewards[-10:]), step=i)
        log_writer.add_scalar("Epsilon Delay", value=epsilon, step=i)

    if i % update_tar_interval == 0:
        agent.DQN_target.load_dict(agent.DQN.state_dict())
    
    if done:
        frame = env.reset()
        frame, stacked_frames = stack_frames(stacked_frames, frame, True)
        all_rewards.append(episode_reward)
        log_writer.add_scalar("Reward", value=episode_reward, step=episode_num)
        episode_reward = 0
        episode_num += 1
        avg_reward = float(np.mean(all_rewards[-100:]))

打印结果（其实可以用 visualDL 指定 ./log 文件夹也是可以的啦）

def plot_training(frame_idx, rewards, losses):
    clear_output(True)
    plt.figure(figsize=(20,5))
    plt.subplot(131)
    plt.title('frame %s. reward: %s' % (frame_idx, np.mean(rewards[-10:])))
    plt.plot(rewards)
    plt.subplot(132)
    plt.title('loss')
    plt.plot(losses)
    plt.show()

plot_training(i, all_rewards, losses)