基于强化学习的DQN智能体自动生成XSS-先知社区

基于强化学习的DQN智能体自动生成XSS

Q learning 和 DQN

Q learning

DQN实际上就是Q learning+network

那先来看看Q learning，公式如下：

图片加载失败

Q learning和DQN的区别在于，Q learning的Q值是用Q表格来储存的，DQN使用神经网络来储存的

Q learning的工作流程大概：

1.初始化：Q值通常开始随机被初始化，然后在训练的过程中更新

2.探索与利用：在每个时间步智能体都要选择一个动作。这里使用epsilon-greedy策略来完成，该方法会在随机选择动作和选择但前最高Q值的动作之间权衡

3.学习更新：一旦智能体选择了一个行动，环境返回了结果，智能体会根据结果，基于贝尔曼公式和时序差分来更新Q值

那我们来看看Q learning简单的代码实现

可以看到

在实现选择动作着这个函数这里，可以看到epsilon-greedy策略，如果随机数小于epsilon，那就选择基于当前状态下Q值最大的那个动作，否则就是随机选择一个动作并返回

这段代码就是公式

图片加载失败

的具体实现过程，让我着重关注一下q_target和q_predict，q_target 是下一个状态的最大Q值，q_predict是当前状态的Q值，这个公式的目的就是使当前状态尽可能的去你和下一状态，计算下一状态和当前状态Q值的差，再乘以折扣率Gamma（即是这个误差存在一定损失），再乘上学习率alpha，这样就可以逐步的去拟合下一状态的Q值，讲到这里是不是有神经网络梯度下降那味儿了

DQN

在传统的Q-learning中，我们用一个表（Q-table）来存储每个状态-动作对的Q值。然而，当状态和动作的数量非常大时，用表格存储的方式就会变得不现实，因为需要的存储空间和计算资源会非常巨大

那就顺势提出了使用神经网络来充当Q值函数，通过这种方式，我们就可以在连续的状态空间和大规模的动作空间中工作

提到DQN就不得不提提他的两个关键技术：

1.经验回放**（Experience Replay）**：为了打破数据之间的关联性和提高学习效率，DQN会将智能体的经验（状态、动作、奖励、新状态、新动作）储存起来，之后从中随机抽样进行学习

目标网络（Target Network）

：DQN使用了两个神经网络，一个是在线网络，用于选择动作；一个是目标网络，用于计算TD目标（Temporal-Difference Target），这两个网络的结构是完全一样的，只是参数不同，在学习过程中，每个一段时间，会用在线网络的参数去更新目标网络

怎么理解这个target network呢？我这里引用两个师傅的例子

A.把在线网络做一只猫。把监督数据 Q Target 看做是一只老鼠，现在可以把训练的过程看做猫捉老鼠的过程（不断减少之间的距离，类比于在线网络拟合 Q Target 的过程）。现在问题是猫和老鼠都在移动，这样猫想要捉住老鼠是比较困难的

那么让老鼠在一段时间间隔内不动（固定住），而这期间，猫是可以动的，这样就比较容易抓住老鼠了。在 DQN 中也是这样解决的，有两套一样的网络，分别是在线网络和 Q Target 网络。要做的就是固定住 Q target 网络，那如何固定呢？比如可以让在线网路训练10次，然后把在线网络更新后的参数 w 赋给 Q target 网络。然后再让在线网路训练10次，如此往复下去，试想如果不固定 Q Target 网络，两个网络都在不停地变化，这样拟合是很困难的，如果让 Q Target 网络参数一段时间固定不变，那么拟合过程就会容易很多

B.同样的道理，把在线网络去拟合target网络这个过程比作是打靶，如果靶子一直动来动去，那肯定加大了打中的难度，那我们使用target网络把靶子固定起来，那打中的概率是不是就会大很多了呢

介绍一下DQN的整体工作流程：

其实就是在线网络和目标网络的相互配合

在线网络训练

：在线网络和环境交互，在线网络执行了一个动作，环境会返回（状态、动作、奖励、新状态、新动作）然后使用这些数据来更新网络参数，我们希望在线网络的预测值接近于目标值，我们可以使用梯度下降算法来最小化在线网络预测的Q值和目标网络的目标值之间的差距（通常使用平方损失函数）。

图片加载失败

Q值的更新公式为：

DQN工作的整体流程：

1.初始化：初始化在线网络和目标网络，创建一个经验回放储存区

2.探索与利用：在每个时间步智能体都要选择一个动作。这里使用epsilon-greedy策略来完成，该方法会在随机选择动作和选择但前最高Q值的动作之间权衡

3.交互与储存：智能体与环境进行交互，产生的（状态、动作、奖励、新状态、新动作）储存在经验回放区中

4.学习：从经验回放储存区中随机抽取一些样本来训练在线网络，通过最小化网络预测的Q值和这个目标值之间的差距来更新网络的参数

5.更新网络：每个一定的时间会将在线网络的参数直接拷贝给目标网络，是目标网络的参数保持相对稳定，使学习过程更相对稳定

6.迭代：重复2～5步骤

最后贴一下我项目的部分代码吧

这段代码就是基于DQN写的，其他的类由于篇幅原因，就不贴出来了

features特征提取

先来看看项目的整体流程：

主要有DQNAgent和WAF_env组成

除开这两个重要的，还有接受DQNAgent命令执行具体免杀操作的XSS_Manipulator模块，正则检测的WAF模块，还有一个Features特征提取模块，下图是整个流程图：

图片加载失败

那接下来讲讲features模块

老规矩，先上代码

这段代码的作用类比方块走迷宫中的迷宫环境，使用

np.bincount()

统计列表中

每个 ASCII 值出现的次数

，生成一个

长度为 256（ASCII 码范围）

的数组

，表示每个字符的频率，这个列表即为我们的xss样本，这样做的目的说白了就是为了，在之后的训练中提供给智能体一个类似于刚刚所说的迷宫环境，要训练的就是智能体针对不同的样本（从字符串中字符出现的频次出现出来）给出不同的免杀方法

这里我们采用的维度是1+256维(即257维)，后面的256维是字符的频率分布，前面拼接的1维是样本的长度

接着还是老生常谈的问题：归一化

归一化

这里为什么要归一化呢？因为：

1.梯度爆炸或消失：如果输入的数值跨度太大，神经网络的梯度更新可能不稳定，影响学习效果

更新方向偏差

：模型可能倾向于优先优化数值较大的特征，而忽略数值较小的特征

这里用公式解释一下：

神经网络的参数更新公式：

图片加载失败

其中：

●

w是神经网络的权重

●

alphaα 是学习率（learning rate）

●

L 是损失函数

在反向传播（Backpropagation）时，权重的梯度计算：

图片加载失败

其中：

●

由损失函数决定

●

由输入数据的特征值决定

如果特征x的数值范围过大（例如 1000），那么梯度计算时：

图片加载失败

梯度就会变得很大，导致：

●

该特征的权重变化幅度大，模型更关注它

●

数值小的特征（例如 0.01）的梯度较小，更新幅度小，可能被忽略

在代码实现的过程中遇到一个问题：虽然 h.astype(self.dtype) 强制 h 变成了 float32，但是 h.sum() 是 int64 类型，在 NumPy 中，当 float32 除以 int64，结果会被提升为 float64

所以h.sum也要使用astype

基于正则的WAF

老规矩，先上代码

说白了就是用re做一个正则匹配

re.IGNORECASE 是 Python 标准库 re（正则表达式模块）中的一个标志，用于在正则表达式匹配时忽略大小写。它允许正则表达式在匹配字符串时，不区分字母的大小写

这里的waf我没给完，可以自己添加

Xss_Manipulator免杀模块

也就是当智能体做出某个免杀操作是，xss_manipulator接受智能体的命令对xss的payload做出免杀操作

这里要讲一下getattr 函数，可以从对象中根据字符串名称获取属性或方法。如果属性是一个可调用的方法，可以直接调用它

我在这个类实现的就是，传递动作字符串，然后使用getattr来调用对应的方法动作列表

ENV

这是强化学习最重要的部分

而这里最重要的就是利用gym提供的框架，那就来学习一下

如何写一个gym.Env的子类

以平衡车为例简单了解一下

看起来很简单吧

那我们照葫芦画瓢，写一下本项目的env

DQN_Agent

repaly经验重放那里前面已经讲了很多了，没什么好说的了

讲讲train_dqn函数吧

整个过程训练了100轮次，每个轮次先初始化加入500个记忆

训练时，从这个经验数据中随机抽取batch_size个大小进行训练

本代码的模型很简单：

值得注意的是我这里在梯度下降时使用的是mse来计算损失函数

图片加载失败

损失函数

学到这里突然意识到损失函数使用正确与否的重要性，记录一下

分为两大类吧：

1.回归任务

2.分类任务（二分类和多分类）

回归损失函数

适用于

回归任务

（预测连续值）

损失函数	代码	说明
均方误差（MSE, Mean Squared Error）	'mse' 或 tf.keras.losses.MeanSquaredError()	计算预测值与真实值的平方误差，适用于大多数回归任务
均方对数误差（MSLE, Mean Squared Logarithmic Error）	'msle' 或 tf.keras.losses.MeanSquaredLogarithmicError()	适用于数据分布跨度较大的情况（比如 0.01 和 1000），减少大值对损失的影响
均绝对误差（MAE, Mean Absolute Error）	'mae' 或 tf.keras.losses.MeanAbsoluteError()	计算预测值与真实值的绝对误差，适用于稳健回归任务，避免平方项带来的影响
均绝对百分比误差（MAPE, Mean Absolute Percentage Error）	'mape' 或 tf.keras.losses.MeanAbsolutePercentageError()	适用于对相对误差敏感的回归任务
Huber Loss	'huber' 或 tf.keras.losses.Huber(delta=1.0)	结合 MSE 和 MAE 的优势，对异常值更具鲁棒性，适用于 DQN 强化学习

分类损失函数

适用于

分类任务

（如二分类、多分类）。

损失函数	代码	适用场景
二元交叉熵（Binary Crossentropy）	'binary_crossentropy' 或 tf.keras.losses.BinaryCrossentropy()	适用于二分类任务
分类交叉熵（Categorical Crossentropy）	'categorical_crossentropy' 或 tf.keras.losses.CategoricalCrossentropy()	适用于独热编码（one-hot）多分类任务
稀疏分类交叉熵（Sparse Categorical Crossentropy）	'sparse_categorical_crossentropy' 或 tf.keras.losses.SparseCategoricalCrossentropy()	适用于整数标签编码的多分类任务（不需要 one-hot）

二分类（0/1）二元交叉熵函数→

binary_crossentropy

（sigmoid）

多分类（one-hot）多元交叉熵函数→

categorical_crossentropy

（softmax）

多分类（整数索引）稀疏多元交叉熵函数→

sparse_categorical_crossentropy

（softmax）

回归任务 vs. 二分类任务

均方误差（MSE，Mean Squared Error）：

图片加载失败

回归任务可以根据公式来看，主要是数值逼近问题，所以均方误差用来衡量预测值和真实值的误差再合适不过

二元交叉熵（Binary Crossentropy, BCE）：

公式：

图片加载失败

想必大家在刚接触时都有疑惑，为什么分类任务不用均方误差？

因为

均方误差是线性的

，但在二分类任务中，概率的误差往往是非线性的（接近 0 或 1 的误差比 0.5 附近的误差影响更大）。

还没有理解的话，举一个实际的数学例子就明白了

直观理解 MSE 和 BCE

假设我们在训练一个二分类模型，预测某个邮件是否是垃圾邮件（Spam: 1，Not Spam: 0）：

邮件	真实标签 (y)	预测概率 (y^)	MSE 计算	BCE 计算
A	1	0.9	(1 - 0.9)² = 0.01	-[1 * log(0.9) + 0 * log(0.1)] ≈ 0.105
B	0	0.1	(0 - 0.1)² = 0.01	-[0 * log(0.1) + 1 * log(0.9)] ≈ 0.105
C	1	0.5	(1 - 0.5)² = 0.25	-[1 * log(0.5) + 0 * log(0.5)] ≈ 0.693

可以看到如果使用二元交叉熵的话，越接近0或1，值越小

但均方误差因为是线性的，所以 0.9 预测 1 ， 0.5 预测 1 的损失这两个的损失变化不大，这对于分类任务来说是致命的