如何设计一个 DRL 调度器？ · TPDS '21

Hailiang Zhao | 2022/08/22

Categories: paper Tags: TPDS scheduling

面向深度学习训练任务（以下简称 DL 任务）的调度器，除了通用的调度平台（如 Kubernetes 和 YARN），一些研究工作试图通过建立解析模型并设计启发式策略来实现任务调度。然而，这些解析模型通常是和具体的任务负载类型强关联的，一旦更换待目标任务，这些调度器就很难取得理想的效果。

本文将分析一篇去年发表在 TPDS 上的论文 DL2: A Deep Learning-driven Scheduler for Deep Learning Clusters。在这篇论文中，作者提出了一个黑盒调度器 DL2。DL2 基于策略迭代深度强化学习算法（policy-based DRL）而设计 1，它将 集群当前的资源利用状态 作为输入，将 资源分配方案 作为动作输出，并以 当前任务们的训练进度 作为反馈指导自身作出更加合理的决策。 DL2 在被投入线上进行增量强化学习之前，会先在线下基于 “历史任务数据（job traces）和对应的调度决策” 进行有监督学习，这样可以保证 DL2 至少不逊于一个正在投入使用的 scheduler。

1 研究动机

在集群中，DL 训练任务通常以分布式的方式进行。一种典型的架构是 Parameter Server (PS)-Worker 架构 2。首先，DL 训练任务的提交者需要列出当前任务需要多少个 PS、多少个 worker，以及每个 PS（和 worker）需要多少数量的资源（GPU、CPU、MEM）。提交之后，scheduler 将会根据调度策略适时地将其从队列中取出，并按照其需求说明为其分配资源。然而，对于一个 DL 任务，提交者应当为其指定多少个 PS 和 worker？一般来说，PS（和 worker）越多越好，但是这会占用更多的资源。作者发现，DL 任务的训练效率并不随着 PS（和 worker）的数量呈线性增长——而是一种 次线性 的关系，即存在 边际效应递减 的现象。图 1 反映出了这一现象。

图 1 在不同数量的 worker 下（对于每一个数据点，PS 个数等于 worker 个数），和只有单个 worker、单个 PS 的情形相比，所取得的训练加速。可以发现，训练加速并不随着 PS（和 worker）的数量呈线性增长。

此外，作者还意识到，对于一个 DL 任务，如何设计合适的 PS-worker 比例是很困难的。对不同的 DL 任务而言，使得其训练速度最大的 PS-worker 比例各不相同。图 2 反映出了这一现象。

图 2 对于 Seq2Seq 和 VGG-16 这两个模型，不同的 PS-worker 比例下的训练速度。

上述分析表明，对于一个 DL 任务而言，如何决定其最优的 PS 和 worker 的个数从而 在训练速度和边际效应之间作出折衷 是很困难的。我们应当充分利用集群当前 GPU 的实际占用情况动态地为 DL 任务调整资源分配。DL2 正是这样做的，它会根据集群的资源利用状态不断调整 DL 任务的 PS 和 worker 的个数，从而使得集群资源利用率最大化。

2 强化学习知识梳理

我会在这一章节补充一些强化学习的背景知识 3。这部分内容会比较冗长，但是这对于理解强化学习算法而言是必须的。

2.1 马尔可夫决策过程

在强化学习中，有两个相互交互的对象：智能体（agent）和环境（environment）。

图 3 智能体与环境的交互。

智能体的 策略（policy） 就是智能体如何根据环境状态 $s$ 来决定下一步的动作 $a$，通常可以分为 确定性策略（Deterministic Policy） 和 随机性策略（Stochastic Policy） 两组。确定性策略是从状态空间到动作空间的映射函数 $\pi: \mathcal{S} \to \mathcal{A}$。随机性策略表示在给定环境状态时，智能体选择某个动作的概率分布：

$$ \begin{equation} \begin{aligned} \pi (a | s) := p (a | s), \qquad \sum_{a \in \mathcal{A}} \pi (a | s) = 1, \forall s \in \mathcal{S}. \end{aligned} \end{equation} $$ 通常情况下，强化学习一般使用随机性策略。随机性策略在学习时可以通过引入一定随机性更好地探索环境，从而使得策略更加地多样。比如，在围棋中，确定性策略总是会在同一个位置上下棋，这就很容易被对手预测。

智能体从感知到的初始环境 $s_0$ 开始，然后决定做一个相应的动作 $a_0$，环境相应地发生改变到新的状态 $s_1$，并反馈给智能体一个即时奖励 $r_1$，然后智能体又根据状态 $s_1$ 做一个动作 $a_1$，环境相应改变为 $s_2$，并反馈奖励 $r_2$。这样的交互可以一直进行下去。我们因此可以得到一条智能体与环境的交互轨迹（trajectory）：

$$ \begin{equation} \begin{aligned} s_0, a_0, s_1, r_1, a_1, ..., s_{t-1}, r_{t-1}, a_{t-1}, s_t, r_t, ..., \end{aligned} \end{equation} $$ 其中 $r_t := r(s_{t-1}, a_{t-1}, s_t)$ 是第 $t$ 时刻的即时奖励。

智能体与环境的交互的过程可以看作是一个 马尔可夫决策过程。

马尔可夫过程（Markov Process） 是具有马尔可夫性的随机变量序列 $s_0, s_1, ..., s_t \in \mathcal{S}$，其下一个时刻的状态 $s_{t+1}$ 只取决于当前状态 $s_t$，即 $$ \begin{equation} \begin{aligned} p (s_t | s_t, ..., s_0) = p (s_t | s_{t-1}), \end{aligned} \end{equation} $$ 其中 $p (s_{t+1} | s_t)$ 被称为状态转移概率，显然 $\sum_{s_{t+1} \in \mathcal{S}} p (s_{t+1} | s_t) = 1$。

马尔可夫决策过程（Markov Decision Process，MDP） 在马尔可夫过程中加入一个额外的变量：动作 $a$，即下一个时刻的状态 $s_{t+1}$ 和当前时刻的状态 $s_t$ 以及 动作 $a_t$ 相关，

$$ \begin{equation} \begin{aligned} p (s_t | s_t, a_t, ..., s_0, a_0) = p (s_t | s_{t-1}, a_{t-1}), \end{aligned} \end{equation} $$

其中 $p (s_{t+1} | s_t, a_t)$ 被称为状态转移概率。

给定策略 $\pi (a | s)$，马尔可夫决策过程的一个轨迹

$$ \begin{equation} \begin{aligned} \tau := s_0, a_0, s_1, r_1, a_1, ..., s_{T-1}, r_{T-1}, a_{T-1}, s_T, r_T \end{aligned} \end{equation} $$ 的概率为 $$ \begin{equation} \begin{aligned} p (\tau) = p (s_0) \prod_{t=0}^{T-1} \pi (a_t | s_t) p (s_{t+1} | s_t, a_t). \end{aligned} \end{equation} $$

2.2 强化学习的目标函数

给定策略 $\pi (a | s)$，智能体和环境一次交互过程的轨迹 $\tau$ 所收到的累积奖励为总回报（Return）：

$$ \begin{equation} \begin{aligned} G (\tau) := \sum_{t=0}^{T-1} r_{t+1} = \sum_{t=0}^{T-1} r(s_t, a_t, s_{t+1}). \end{aligned} \end{equation} $$

假设环境中有一个或多个特殊的终止状态（Terminal State），当到达终止状态时，一个智能体和环境的交互过程就结束了。这一轮交互的过程称为一个回合（Episode）或试验（Trial）。一般的强化学习任务（比如下棋、游戏）都属于这种回合式的任务。 如果环境中没有终止状态（比如终身学习的机器人），即 $T \to \infty$，称为 持续性强化学习任务，其总回报也可能是无穷大 4。为了解决这个问题，我们可以引入一个折扣率来降低远期回报的权重。折扣回报（Discounted Return）定义为：

$$ \begin{equation} \begin{aligned} G (\tau) := \sum_{t=0}^{T-1} \gamma^t r_{t+1}, \end{aligned} \end{equation} $$ 其中 $\gamma_t \in [0,1]$ 是折扣率。当 $\gamma$ 接近于 $0$ 时，智能体更在意短期回报；而当 $\gamma$ 接近于 $1$ 时，长期回报变得更重要。

强化学习的目标是学习到一个策略 $\pi_\theta (a | s)$，从而最大化期望回报（Expected Return）$\mathcal{J} (\theta)$，即希望智能体执行一系列的动作来获得尽可能多的平均回报：

$$ \begin{equation} \begin{aligned} \mathcal{J} (\theta) := \mathbb{E}_{\tau \sim p_\theta (\tau)} \big[G (\tau) \big] = \mathbb{E}_{\tau \sim p_\theta (\tau)} \bigg[ \sum_{t=0}^{T-1} \gamma^t r_{t+1} \bigg]. \end{aligned} \tag{e.q. 1} \end{equation} $$

持续性强化学习的优化目标也可以定义为 MDP 到达平稳分布时 “即时奖励” 的期望。

2.3 值函数

为了评估一个策略 $\pi$ 的期望回报，我们定义两个值函数：状态值函数（state value function） 和 状态 - 动作值函数（action-state value function）。

一个策略 $\pi$ 的期望回报可以分解为

$$ \begin{equation} \begin{aligned} \mathbb{E}_{\tau \sim p (\tau)} \big[G (\tau) \big] = \mathbb{E}_{s \sim p(s_0)} \Bigg[ \mathbb{E}_{\tau \sim p (\tau)} \Big[ \sum_{t=0}^{T-1} \gamma^t r_{t+1} | \tau_{s_0} = s \Big] \Bigg] = \mathbb{E}_{s \sim p(s_0)} \big[V^\pi (s)\big], \end{aligned} \end{equation} $$ 其中 $\tau_{s_0}$ 表示轨迹 $\tau$ 的初始状态。

我们将 $V^\pi (s) := \mathbb{E}_{\tau \sim p (\tau)} \Big[ \sum_{t=0}^{T-1} \gamma^t r_{t+1} | \tau_{s_0} = s \Big]$ 称为状态值函数，它表示从状态 $s$ 开始，执行策略 $\pi$ 得到的期望总回报。

根据马尔可夫性，将 $V^\pi (s)$ 展开，我们可以得到著名的贝尔曼方程（a.k.a. 动态规划方程）：

$$ \begin{equation} \begin{aligned} V^\pi (s) = \mathbb{E}_{a \sim \pi (a | s)} \mathbb{E}_{s'\sim p (s' | s, a)} \Big[ r(s, a, s') + \gamma V^\pi (s') \Big]. \end{aligned} \end{equation} $$

这个公式表明当前状态的值函数可以通过下个状态的值函数来计算。如果给定策略 $\pi (a | s)$，状态转移概率 $p (s'| s, a)$ 和奖励 $r(s, a, s')$，我们就可以通过迭代的方式来计算 $V^\pi (s)$。由于存在折扣率，迭代一定步数后，每个状态的值函数就会固定不变。

我们将状态 - 动作值函数 $Q^\pi (s, a)$ 定义为：在初始状态为 $s$ 并进行动作 $a$ 的条件下，执行策略 $\pi$ 得到的期望总回报，即

$$ \begin{equation} \begin{aligned} Q^\pi (s, a) := \mathbb{E}_{s'\sim p (s' | s, a)} \Big[ r(s, a, s') + \gamma V^\pi (s') \Big]. \end{aligned} \tag{e.q. 2} \end{equation} $$

可以发现，状态值函数 $V^\pi (s)$ 是 $Q$ 函数 $Q^\pi (s, a)$ 关于动作 $a$ 的期望：

$$ \begin{equation} \begin{aligned} V^\pi (s) := \mathbb{E}_{a \sim \pi (a | s)} \Big[ Q^\pi (s, a) \Big]. \end{aligned} \tag{e.q. 3} \end{equation} $$

结合 $(\textrm{e.q. 2})$ 和 $(\textrm{e.q. 3})$，我们可以得到 $Q$ 函数的贝尔曼方程：

$$ \begin{equation} \begin{aligned} Q^\pi (s, a) = \mathbb{E}_{s'\sim p(s' | s, a)} \bigg[ r(s, a, s') + \mathbb{E}_{a' \sim \pi (a'| s')} \Big[ Q^\pi (s', a') \Big] \bigg]. \end{aligned} \end{equation} $$

至此，强化学习的基本知识介绍完毕。接下来，我将梳理常用的强化学习算法。

2.4 强化学习算法

不同强化学习算法之间的关系总结如图 4 所示。我们首先来看基于值函数的学习方法。

图 4 不同强化学习算法之间的关系。

2.4.1 基于值函数的学习方法

注意，我们的目标是找到一个策略 $\pi^*$，使得

$$ \begin{equation} \begin{aligned} \forall s \in \mathcal{S}, \pi^* = \operatorname{argmax}_{\pi} V^\pi (s). \end{aligned} \end{equation} $$

直接求解上述优化问题是很困难的，我们通过迭代的方法不断优化策略，直到选出最优策略。 对于一个策略 $\pi (a | s)$，其 $Q$ 函数为 $Q^\pi (s, a)$，我们可以设置一个新的策略 $\pi' (a | s)$，它满足：

$$ \begin{equation} \begin{aligned} \pi' (a | s) = \left\{ \begin{array}{ll} 1 & a = \operatorname{argmax}_{\hat{a}} Q^\pi (s, \hat{a}) \\ 0 & \textrm{otherwise}. \end{array} \right. \end{aligned} \end{equation} $$

这样，当我们执行策略 $\pi' (a | s)$，总会有

$$ \begin{equation} \begin{aligned} \forall s: V^{\pi'} (s) \geq V^\pi (s). \end{aligned} \end{equation} $$

由此，我们便将策略 $\pi$ 更新为了 $\pi'$。

如果知道马尔可夫决策过程的状态转移概率 $p (s'| s, a)$ 和奖励 $r (s, a, s')$，我们直接可以通过贝尔曼方程（动态规划）来迭代计算其值函数。这种模型已知的强化学习算法也称为 基于模型的强化学习（Model-Based Reinforcement Learning） 算法，这里的模型就是指马尔可夫决策过程。

第一个要介绍的基于动态规划的算法是 策略迭代算法（Policy Iteration），如算法 1 所示。

算法 1 策略迭代算法。

策略迭代中的 策略评估 和 策略改进 交替轮流进行，其中策略评估也是通过一个内部迭代来进行计算，其计算量比较大。事实上，我们不需要每次都精确地计算出值函数，也就是说内部迭代不需要执行到完全收敛。值迭代方法（Value Iteration） 将策略评估和策略改进两个过程合并，来直接计算出最优策略。值迭代算法的步骤如算法 2 所示。

算法 2 值迭代算法。

策略迭代算法是根据贝尔曼方程来更新值函数，并根据当前的值函数来改进策略。而值迭代算法是直接使用贝尔曼最优方程来更新值函数，收敛时的值函数就是最优的值函数，其对应的策略也就是最优的策略。值迭代和策略迭代都需要经过非常多的迭代次数才能完全收敛。在实际应用中，可以不必等到完全收敛。这样，当状态和动作数量有限时，经过有限次迭代就能收敛到近似最优策略。

以上两个方法都是基于模型的强化学习算法。它们都要求模型已知，即 需要知道马尔可夫决策过程的状态转移概率 $p(s'|s, a)$ 和奖励函数 $r(s,a,s')$，这个要求很难满足。此外，这两个算法都要求枚举所有的状态和动作，对于很多实际问题而言，这是很不现实的。

当我们无法获取状态转移概率 $p(s'|s, a)$ 和奖励函数 $r(s,a,s')$ 的时候，我们可以通过让智能体和环境交互来收集一些轨迹样本，然后再根据这些样本来求解马尔可夫决策过程最优策略。这种模型未知，基于采样的学习算法也称为模型无关的强化学习（Model-Free Reinforcement Learning）算法。 此时我们有

$$ \begin{equation} \begin{aligned} \forall s, a: Q^\pi (s, a) \approx \hat{Q}^\pi_N (s, a) := \frac{1}{N} \sum_{n=1}^N G \Big(\tau^{(n)}_{s_0 = s, a_0 = a}\Big), \end{aligned} \tag{e.q. 4} \end{equation} $$

其中 $\tau^{(n)}_{s_0 = s, a_0 = a}$ 表示第 $n$ 个轨迹，其起始状态和对应的动作分别为 $s$ 和 $a$。这一类基于采样的方法被称为 蒙特卡罗方法。

在蒙特卡罗方法中，如果采用确定性策略 $\pi$，每次试验得到的轨迹是一样的，只能计算出 $Q^\pi (s, \pi (s))$，而无法计算其它动作 $a'$ 的 $Q$ 函数，因此也无法进一步改进策略。这仅仅是对当前策略的利用（exploitation），而缺失了对环境的探索（exploration），即试验的轨迹尽可能覆盖所有的状态和动作，以找到更好的策略。对此，对于任意目标策略 $\pi$，我们引入对应的 $\epsilon$-greedy 策略如下：

$$ \begin{equation} \begin{aligned} \pi^\epsilon (s) = \left\{ \begin{array}{ll} \pi (s) & \textrm{with prob. } 1 - \epsilon \\ \textrm{randomly select from } \mathcal{A} & \textrm{with prob. } \epsilon. \end{array} \right. \end{aligned} \end{equation} $$

有了 $\pi^\epsilon (s)$，我们可以将蒙特卡洛算法进一步划分为 同策略方法 和 异策略方法：

蒙特卡罗采样方法一般需要拿到完整的轨迹，才能对策略进行评估并更新模型，因此效率也比较低。时序差分学习（Temporal-Difference Learning） 结合了动态规划和蒙特卡罗方法，比仅仅使用蒙特卡罗采样方法的效率要高很多。时序差分学习是模拟一段轨迹，每行动一步 (或者几步)，就利用贝尔曼方程来评估行动前状态的价值。当时序差分学习中每次更新的动作数为最大步数时，就等价于蒙特卡罗方法。

根据 $(\textrm{e.q. 4})$ 的定义，可以得到

$$ \begin{equation} \begin{split} \hat{Q}^\pi_N (s, a) &= \frac{1}{N} \sum_{n=1}^N G \Big(\tau^{(n)}_{s_0 = s, a_0 = a}\Big) \\ &= \frac{1}{N} \Bigg( G \Big(\tau^{(N)}_{s_0 = s, a_0 = a}\Big) + \sum_{n=1}^{N-1} G \Big(\tau^{(n)}_{s_0 = s, a_0 = a}\Big) \Bigg) \\ &= \hat{Q}^\pi_{N-1} (s, a) + \frac{1}{N} \Bigg( G \Big(\tau^{(N)}_{s_0 = s, a_0 = a}\Big) - \hat{Q}_{N-1}^\pi (s, a) \Bigg). \end{split} \end{equation} $$

这意味着，值函数 $\hat{Q}^\pi_N (s, a)$ 在第 $N$ 试验后的平均等于第 $N − 1$ 试验后的平均加上一个增量。更一般性地，我们将权重系数 $\frac{1}{N}$ 改为一个比较小的正数 $\alpha$。这样每次采样一个新的轨迹 $\tau_{s_0 = s, a_0 = a}$，就可以更新 $\hat{Q}^\pi_N (s, a)$：

$$ \begin{equation} \begin{split} \hat{Q}^\pi_N (s, a) \leftarrow \hat{Q}^\pi_N (s, a) + \alpha \bigg( G (\tau_{s_0 = s, a_0 = a}) - \hat{Q}^\pi_N (s, a) \bigg), \end{split} \tag{e.q. 5} \end{equation} $$

其中的增量 $\delta := G (\tau_{s_0 = s, a_0 = a}) - \hat{Q}^\pi_N (s, a)$ 叫做 蒙特卡罗误差，表示当前轨迹的真实回报 $G (\tau_{s_0 = s, a_0 = a})$ 与期望回报 $\hat{Q}^\pi_N (s, a)$ 之间的差距。然而，如果不采样出一条完整的轨迹，我们就无法知道 $G (\tau_{s_0 = s, a_0 = a})$ 的数值。为了提高效率，我们可以借助动态规划的方法来计算 $G (\tau_{s_0 = s, a_0 = a})$，而不需要得到完整的轨迹。具体地，不妨设当前的状态和动作分别为 $s$ 和 $a$，我们采样下一步的状态 $s'$ 和动作和 $a'$，并得到奖励 $r(s, a, s')$，然后我们利用贝尔曼方程来近似估计 $G (\tau_{s_0 = s, a_0 = a})$：

$$ \begin{equation} \begin{split} G (\tau_{s_0 = s, a_0 = a, s_1 = s', a_1 = a'}) &= r(s, a, s') + \gamma G (\tau_{s_0 = s', a_0 = a'}) \\ &\approx r(s, a, s') + \gamma \hat{Q}^\pi (s', a'), \end{split} \tag{e.q. 6} \end{equation} $$

其中 $\hat{Q}^\pi (s', a')$ 是对 $Q^\pi (s', a')$ 的近似估计。结合 $(\textrm{e.q. 5})$ 和 $(\textrm{e.q. 6})$，我们可以得到近似 $Q$ 函数的更新公式：

$$ \begin{equation} \begin{split} \hat{Q}^\pi (s, a) \leftarrow \hat{Q}^\pi (s, a) + \alpha \bigg( r(s, a, s') + \gamma \hat{Q}^\pi (s', a') - \hat{Q}^\pi (s, a) \bigg). \end{split} \tag{e.q. 7} \end{equation} $$

基于上述迭代公式的方法被称为 SARSA 算法（State Action Reward State Action，SARSA），其步骤总结如算法 3 所示。

算法 3 SARSA 算法。

SARSA 采样和优化的策略都是 $\pi^\epsilon$，因此是一种同策略算法。为了提高计算效率，我们不需要对环境中所有的 $s, a$ 组合进行穷举，并计算值函数。只需要将当前的探索 $s', a'$ 作为下一次估计的起始状态和动作。时序差分学习是强化学习的主要学习方法，其关键步骤就是在每次迭代中优化 $Q$ 函数来减少现实 $r + \gamma \hat{Q}^\pi (s', a')$ 和预期 $Q(s, a)$ 的差距。 这和动物学习的机制十分相像。在大脑神经元中，多巴胺的释放机制和时序差分学习十分吻合。多巴胺的释放, 来自实际奖励和预期奖励的差异，而不是奖励本身。

时序差分学习和蒙特卡罗方法的主要不同为：蒙特卡罗需要完整一个路径完成才能知道其总回报，也不依赖马尔可夫性质；而时序差分学习只需要一步，其总回报需要依赖马尔可夫性质来进行近似估计。

$Q$ 学习（$Q$-Learning）算法是一种异策略的时序差分学习算法。在 $Q$-Learning 中，$Q$ 函数的估计方法为

$$ \begin{equation} \begin{split} \hat{Q}^\pi (s, a) \leftarrow \hat{Q}^\pi (s, a) + \alpha \bigg( r + \gamma \max_{a'} \hat{Q}^\pi (s', a') - \hat{Q}^\pi (s, a) \bigg). \end{split} \end{equation} $$

这相当于让 $Q(s, a)$ 直接去估计最优状态值函数 $Q^*(s, a)$。$Q$ 学习 的步骤如算法 4 所示。与 SARSA 算法的不同，$Q$ 学习算法不通过 $\pi^\epsilon$ 来选下一步的动作 $a'$，而是直接选最优的 $Q$ 函数，因此更新后的 $Q$ 函数是关于策略 $\pi$ 的，而不是策略 $\pi^\epsilon$ 的。

算法 4 \(Q\)-Learning 算法。

为了在连续的状态和动作空间中计算值函数 $Q^\pi (s, a)$，我们可以用一个函数 $Q_\phi (\vec{s}, \vec{a})$ 来表示近似计算，称为值函数近似（Value Function Approximation）：

$$ \begin{equation} \begin{aligned} Q_\phi (\vec{s}, \vec{a}) \approx Q^\pi (s, a), \end{aligned} \end{equation} $$ 其中 $\vec{s}, \vec{a}$ 分别是状态 $s$ 和动作 $a$ 的向量表示。

函数 $Q_\phi (\vec{s}, \vec{a})$ 通常是一个参数为 $\phi$ 的函数，比如神经网络，输出为一个实数，称为 Q 网络（Q-network）。 我们需要学习参数 $\phi$ 来使得函数 $Q_\phi (\vec{s}, \vec{a})$ 可以逼近值函数 $Q^\pi (s, a)$。如果采用蒙特卡罗方法，就直接让 $Q_\phi (\vec{s}, \vec{a})$ 去逼近平均的总回报 $Q^\pi (s, a)$; 如果采用时序差分方法，就让 $Q_\phi (\vec{s}, \vec{a})$ 去逼近 $\mathbb{E}_{\vec{s}', \vec{a}'} \big[ r + \gamma Q_\phi (\vec{s}', \vec{a}') \big]$。以 $Q$ 学习为例，采用随机梯度下降，目标函数为

$$ \begin{equation} \begin{aligned} \mathcal{L} (s, a, s'| \phi) := \Bigg( \underbrace{r + \gamma Q_\phi (\vec{s}', \vec{a}')}_{\approx \textrm{ reality}} - \underbrace{Q_\phi (\vec{s}, \vec{a})}_{\textrm{approximation}} \bigg)^2, \end{aligned} \end{equation} $$

其中 $\vec{s}', \vec{a}'$ 分别是状态 $s'$ 和动作 $a'$ 的向量表示。

然而，这个目标函数存在两个问题：一是目标不稳定，参数学习的目标依赖于参数本身；二是样本之间有很强的相关性。为了解决这两个问题，研究人员提出了一种深度 Q 网络 (Deep Q-Networks，DQN)。深度 Q 网络采取两个措施：

训练时，随机从经验池中抽取样本来代替当前的样本进行训练。这样，也可以就打破了和相邻训练样本的相似性，避免模型陷入局部最优。经验回放在一定程度上类似于监督学习——先收集样本，然后在这些样本上进行训练。深度 $Q$ 网络的学习过程如算法 5 所示。

算法 5 深度 Q 网络。

2.4.2 基于策略函数的学习方法

强化学习的目标是学习到一个策略 $\pi_\theta (a | s)$ 来最大化期望回报。一种直接的方法是在策略空间直接搜索来得到最佳策略，称为策略搜索（Policy Search）。策略搜索本质是一个优化问题，可以分为基于梯度的优化和无梯度优化。策略搜索和基于值函数的方法相比，策略搜索可以不需要值函数，直接优化策略。参数化的策略能够处理连续状态和动作，可以直接学出随机性策略。

策略梯度（Policy Gradient）是一种基于梯度的强化学习方法。假设 $\pi_\theta (a | s)$ 是一个关于 $\theta$ 的连续可微函数，我们可以用梯度上升的方法来优化参数 $\theta$ 使得目标函数 $\mathcal{J} (\theta)$ 最大。根据 $(\textrm{e.q. 1})$，我们可以推导目标函数 $\mathcal{J} (\theta)$ 关于参数 $\theta$ 的导数如下：

$$ \begin{equation} \begin{split} \frac{\partial \mathcal{J}(\theta)}{\partial \theta} &= \frac{\partial}{\partial \theta} \int_{\tau} p_\theta (\tau) G (\tau) d \tau \\ &= \int_{\tau} p_\theta (\tau) \bigg( \frac{1}{p_\theta (\tau)} \frac{\partial}{\partial \theta} p_\theta (\tau) \bigg) G(\tau) d\tau \\ &= \int_{\tau} p_\theta (\tau) \bigg( \frac{\partial}{\partial \theta} \log p_\theta (\tau) \bigg) G(\tau) d\tau \\ &= \mathbb{E}_{\tau \sim p_\theta (\tau)} \bigg[ \frac{\partial}{\partial \theta} \log p_\theta (\tau) G (\tau) \bigg]. \end{split} \end{equation} $$

可以发现，参数 $\theta$ 优化的方向是 使得总回报 $G(\tau)$ 越大的轨迹 $\tau$ 的概率 $p_\theta (\tau)$ 也越大。进一步地，有：

$$ \begin{equation} \begin{split} \frac{\partial}{\partial \theta} \log p_\theta (\tau) &= \frac{\partial}{\partial \theta} \log \Bigg( p(s_0) \prod_{t=0}^{T-1} \pi_\theta (a_t | s_t) p (s_{t+1} | s_t, a_t) \Bigg) \\ &= \frac{\partial}{\partial \theta} \Bigg( \log p (s_0) + \sum_{t=0}^{T-1} \log \pi_\theta (a_t | s_t) + \sum_{t=0}^{T-1} \log p (s_{t+1} | s_t, a_t) \Bigg)\\ &= \sum_{t=0}^{T-1} \frac{\partial}{\partial \theta} \log \pi_\theta (a_t | s_t). \end{split} \end{equation} $$

可以发现，$\frac{\partial}{\partial \theta} \log p_\theta (\tau)$ 和状态转移概率无关，只和策略函数相关。进一步可得：

$$ \begin{equation} \begin{split} \frac{\partial \mathcal{J}(\theta)}{\partial \theta} &= \mathbb{E}_{\tau \sim p_\theta (\tau)} \bigg[ \frac{\partial}{\partial \theta} \log p_\theta (\tau) G (\tau) \bigg] \\ &= \mathbb{E}_{\tau \sim p_\theta (\tau)} \bigg[ \sum_{t=0}^{T-1} \Big( \frac{\partial}{\partial\theta} \log \pi_\theta (a_t | s_t) \gamma^t G (\tau_{t: T}) \Big) \bigg], \end{split} \tag{e.q. 8} \end{equation} $$

其中 $G (\tau_{t: T}) = \sum_{t'=t}^{T-1} \gamma^{t' - t} r_{t' + 1}$。

基于以上分析，我接下来介绍 REINFORCE 算法。我们可以通过随机游走来采样轨迹，从而去除 $(\textrm{e.q. 8})$ 中的期望运算。即

$$ \begin{equation} \begin{split} \frac{\partial \mathcal{J}(\theta)}{\partial \theta} \approx \frac{1}{N} \sum_{n=1}^N \Bigg( \sum_{t=0}^{T-1} \frac{\partial}{\partial\theta} \log \pi_\theta (a_t^{(n)} | s_t^{(n)}) \gamma^t G (\tau_{t: T}^{(n)}) \Bigg) \end{split} \tag{e.q. 9} \end{equation} $$

我们可以每采样一条轨迹，就利用 $(\textrm{e.q. 9})$ 计算每个时刻的梯度并更新参数，这就是 REINFORCE 算法。其步骤如算法 6 所示。

算法 6 REINFORCE 算法。

REINFORCE 算法的一个主要缺点是不同路径之间的方差很大，导致训练不稳定，这是在高维空间中使用蒙特卡罗方法的的通病。一种减少方差的通用方法是引入一个控制变量。假设要估计函数 $f$ 的期望，为了减少 $f$ 的方差，我们引入一个已知期望的函数 $g$，令

$$ \begin{equation} \begin{split} \hat{f} := f - \alpha \Big(g - \mathbb{E}[g] \Big), \end{split} \end{equation} $$

则可得 $\mathbb{E}[\hat{f}] = \mathbb{E}[f]$。这意味着，我们可以用 $\hat{f}$ 的期望来估计函数 $f$ 的期望，同时利用函数 $g$ 来减小 $\hat{f}$ 的方差。因为

$$ \begin{equation} \begin{split} \operatorname{Var}(\hat{f}) := \operatorname{Var} (f) - 2\alpha \operatorname{Cov}(f, g) + \alpha^2 \operatorname{Var}(g), \end{split} \end{equation} $$

若想要最小化 $\operatorname{Var}(\hat{f})$，即 $\frac{\partial \operatorname{Var}(\hat{f})}{\partial \alpha} = 0$，可得 $\alpha = \frac{\operatorname{Cov}(f, g)}{\operatorname{Var}(g)}$。因此

$$ \begin{equation} \begin{split} \operatorname{Var}(\hat{f}) = \bigg( 1- \frac{\operatorname{Cov}(f, g)^2}{\operatorname{Var}(f) \operatorname{Var}(g)} \bigg) \operatorname{Var}(f) = \Big( 1- \operatorname{corr}(f, g)^2 \Big) \operatorname{Var}(f), \end{split} \end{equation} $$ 其中 $\operatorname{corr}(f, g)$ 是 $f$ 和 $g$ 的相关系数。这意味着，$f$ 和 $g$ 的相关系数越高，$\operatorname{Var}(\hat{f})$ 就越小。

回顾 $(\textrm{e.q. 8})$，在每个时刻 $t$，其策略梯度为

$$ \begin{equation} \begin{split} \frac{\partial \mathcal{J}_t(\theta)}{\partial \theta} = \mathbb{E}_{s_t} \Bigg[ \mathbb{E}_{a_t} \bigg[ \gamma^t G (\tau_{t: T}) \frac{\partial}{\partial\theta} \log \pi_\theta (a_t | s_t) \bigg] \Bigg], \end{split} \end{equation} $$

为了减小策略梯度的方差，我们引入一个和 $a_t$ 无关的基准函数 $b(s_t)$，

$$ \begin{equation} \begin{split} \frac{\partial \hat{\mathcal{J}}_t(\theta)}{\partial \theta} = \mathbb{E}_{s_t} \Bigg[ \mathbb{E}_{a_t} \bigg[ \gamma^t \Big( G (\tau_{t: T}) - b(s_t) \Big) \frac{\partial}{\partial\theta} \log \pi_\theta (a_t | s_t) \bigg] \Bigg]. \end{split} \tag{e.q. 10} \end{equation} $$

因为 $b(s_t)$ 和 $a_t$ 无关，所以有

$$ \begin{equation} \begin{split} \mathbb{E}_{a_t} \bigg[b(s_t) \frac{\partial}{\partial \theta} \log \pi_\theta (a_t | s_t) \bigg] &= \int_{a_t} \Bigg( b(s_t) \frac{\partial}{\partial \theta} \log \pi_\theta (a_t | s_t) \Bigg) \pi_\theta (a_t | s_t) d a_t \\ &= \int_{a_t} b(s_t) \frac{\partial}{\partial \theta} \pi_\theta (a_t | s_t) d a_t \\ &= \frac{\partial}{\partial \theta} b(s_t) \int_{a_t} \pi_\theta (a_t | s_t) d a_t \\ &= \frac{\partial}{\partial \theta} b(s_t) = 0. \end{split} \end{equation} $$

这意味着 $\frac{\partial \hat{\mathcal{J}}_t(\theta)}{\partial \theta} = \frac{\partial \mathcal{J}_t(\theta)}{\partial \theta}$。也就是说，$(\textrm{e.q. 10})$ 不会改变梯度优化的方向。现在的问题是，我们应该选取哪个函数作为基准函数 $b(s_t)$ 呢？ 为了可以有效地减小方差，$b(s_t)$ 和 $G(\tau_{t:T})$ 越相关越好，一个很自然的选择是令 $b(s_t)$ 为值函数 $V^{\pi_\theta}(s_t)$。但是由于值函数未知，我们可以用一个可学习的函数 $V_{\phi} (s_t)$ 来近似值函数，目标函数为

$$ \begin{equation} \begin{split} \mathcal{L} (\phi | s_t, \pi_\theta) := \Big( V^{\pi_\theta}(s_t) - V_{\phi} (s_t) \Big)^2, \end{split} \end{equation} $$

其中 $V^{\pi_\theta}(s_t) = \mathbb{E} [G(\tau_{t:T})]$ 也用蒙特卡罗方法进行估计。上述目标函数关于 $\theta$ 的梯度为

$$ \begin{equation} \begin{split} \frac{\partial \mathcal{L} (\phi | s_t, \pi_\theta)}{\partial \theta} := - \Big( G(\tau_{t:T}) - V^{\pi_\theta}(s_t)\Big) \frac{\partial V^{\pi_\theta}(s_t)}{\partial \theta}. \end{split} \end{equation} $$

相应地，$(\textrm{e.q. 9})$ 变成

$$ \begin{equation} \begin{split} \frac{\partial \hat{\mathcal{J}}_t(\theta)}{\partial \theta} = \mathbb{E}_{s_t} \Bigg[ \mathbb{E}_{a_t} \bigg[ \gamma^t \Big( G (\tau_{t: T}) - V^{\pi_\theta}(s_t) \Big) \frac{\partial}{\partial\theta} \log \pi_\theta (a_t | s_t) \bigg] \Bigg]. \end{split} \end{equation} $$

综上所述，我们可以交替更新近似值函数的参数 $\phi$ 和策略自身的参数 $\theta$，从而在缩小方差的前提下优化策略。这就是 带基准线的 REINFORCE 算法。其步骤总结如算法 7 所示。

算法 7 带基准线的 REINFORCE 算法。

在 REINFORCE 算法中，每次需要根据一个策略采集一条完整的轨迹，并计算这条轨迹上的回报。这种采样方式的方差比较大，学习效率也比较低。我们可以借鉴时序差分学习的思想，使用动态规划方法来提高采样的效率，即从状态 $s$ 开始的总回报可以通过当前动作的即时奖励 $r(s, a, s')$ 和下一个状态 $s'$ 的值函数来近似估计。演员 - 评论员算法（Actor-Critic Algorithm） 是一种结合策略梯度和时序差分学习的强化学习方法。其中演员（Actor）是指策略函数 $\pi_\theta (a | s)$，即学习一个策略来得到尽量高的回报，评论员（Critic）是指值函数 $V_\phi (s)$，是对当前策略的值函数进行估计，即评估 actor 的好坏。借助于值函数，Actor-Critic 算法可以单步就更新参数，不需要等到回合结束才进行更新。和带基准线的 REINFORCE 算法一样，在 Actor-Critic 算法中的策略函数和值函数都是待学习的函数，需要在训练过程中同时学习。

对于从时刻 $t$ 开始的回报 $G(\tau_{t:T})$，我们如下近似计算：

$$ \begin{equation} \begin{split} G(\tau_{t:T}) \approx \hat{G} (\tau_{t:T}) = r_{t+1} + \gamma V_\phi (s_{t+1}). \end{split} \end{equation} $$

在每一步中，一方面，我们更新值函数的参数使其接近于估计的真实回报：

$$ \begin{equation} \begin{split} \min_\phi \bigg(\hat{G} (\tau_{t:T}) - V_\phi (s_t) \bigg)^2, \end{split} \end{equation} $$

另一方面，我们将值函数 $V_\phi (s_t)$ 作为基函数来更新参数 $\theta$，减少策略梯度的方差：

$$ \begin{equation} \begin{split} \theta \leftarrow \theta + \alpha \gamma^t \Big(\hat{G}(\tau_{t:T}) - V_\phi (s_t) \Big) \frac{\partial}{\partial \theta} \log \pi_\theta (a_t | s_t). \end{split} \end{equation} $$

在每步更新中，演员根据当前的环境状态 $s$ 和策略 $\pi_\theta (a | s)$ 去执行动作 $a$，环境状态变为 $s'$，并得到即时奖励 $r$。评论员根据环境给出的真实奖励和之前标准下的打分（也就是 $r + \gamma V_\phi (s')$），来调整自己的打分标准，使得自己的评分更接近环境的真实回报。演员则跟据评论员的打分，调整自己的策略 $\pi_\theta$，争取下次做得更好。开始训练时，演员随机表演，评论员随机打分。通过不断的学习，评论员的评分越来越准，演员的动作越来越好 5。Actor-Critic 算法的步骤参见算法 8.

算法 8 Actor-Critic 算法。

虽然在带基准线的 REINFORCE 算法也同时学习策略函数和值函数，但是它并不是一种 Actor-Critic 算法。因为其中值函数只是用作基线函数以减少方差，并不用来估计回报（即评论员的角色）。

2.5 总结

至此，我们完整地总结了强化学习的理论以及 8 个最常见的算法。读者可查阅下表对集中典型的算法进行分析和比较。

表 1 四种强化学习算法的比较。

3 DL2：设计与架构

DL2 的目标是最小化 DL 任务的平均完成时间。设计一个基于强化学习的算法，首先需要识别出问题中的状态空间，然后合理地设计动作空间和奖励函数。动作空间不应在维度和规模上过大，奖励函数需要使策略可以朝着期望的方向得到优化。

DL 集群包含多个节点，每个节点配备有一个或多个 GPU。每当用户提交一个 DL 任务的时候，他需要指明该任务每个 PS（和 worker）所需要的资源需求，以及本次训练需要进行多少轮（epoch）。DL2 周期性地执行资源分配的决策，它允许资源的动态分配和调整——即单个 DL 任务在不同的时间片可能有不同数量的 PS（和 worker）在运行。对于同步训练（所有 worker 都完成本地参数更新之后，PS 才会开始 averaging），为了保证单个 DL 任务的总 batch size 不变，当 worker 的个数发生变更时，输入给每个 worker 的 mini-batch 大小会被调整；对于异步训练，输入给每个 worker 的 mini-batch 大小则保持不变。

DL2 采用的是基于策略梯度的方法（算法 6、7、8），其架构如图 5 所示。它包含三个组成部分：

图 5 DL2 的架构。

3.1 策略网络的设计

如图 5 所示，DL2 的策略网络（policy network）和值网络（value network，也就是近似值函数，approximation value function）均是全联接神经网络。

状态设计。 在每个调度周期内，输入到 policy network 的状态是一个矩阵 $s = (x, \vec{d}, \vec{e}, \vec{r}, \vec{w}, \vec{u})$，如图 6 所示，它包含如下组成部分：

图 6 输入状态。

对于上述的每个向量，任务按照抵达时间进行排序，从而保证每个调度周期内 $s$ 的唯一性。

动作设计。 DL2 采用随机性策略，policy network 的输出是一个随机性策略 $\pi: \pi (a | s; \theta) \to [0, 1]$。作者对于动作空间的设计是很精巧的，这和通俗的设计方案相比，极大地减小了动作空间的维度和大小。具体地，policy network 的动作空间大小只有 $3J + 1$，对于每个任务 $j$，针对其的动作有：

或者什么也不做（a void action）。我们将针对一个 DL 任务的动作输出称之为一个 inference，DL2 的运行过程是这样的：在单个调度周期内，对于当前状态 $s$，基于策略 $\pi$ 产生一个动作，然后更新状态为 $s'$，继续执行 inference。如此反复执行，直到一个 void 动作被输出。对于 policy network 而言，当它输出 void 动作时，这意味着 DL2 任务继续给任务们增加资源不会带来有效的训练速度提升。

奖励函数的设计。 在每个调度周期 $t$ 内，作者将即时奖励 $r_t$ 设计为

$$ \begin{equation} \begin{split} r_t := \sum_{j \in J} \frac{t_j}{E_j}, \end{split} \tag{e.q. 11} \end{equation} $$

其中 $t_i$ 是任务 $j$ 在当前调度周期内完成的 epoch 个数，$E_j$ 是作业在提交时注明的总 epoch 个数。显然，单个调度周期内完成的 epoch 越多，奖励就越大。

算法。 DL2 采用基于策略梯度的强化学习算法。具体地，DL2 采用的是 Actor-Critic 算法（算法 8）。 和标准的强化学习算法不同的是，DL2 在单个调度周期内，执行多次 inference，因此会得到多个四元组 $(s, a, r, s')$。只有当所有的 inference 都完成之后，才会更新 policy network 的梯度。和算法 8 中的步骤略有不同的是，作者采用 experience reply 的方式来更新 policy network 和 value network。在每个调度周期内采样得到的四元组 $(s, a, r, s')$ 将会被放入一个经验池中，在更新 policy network 和 value network 时，会随机从经验池中选出一组四元组（a mini-batch of data）用于梯度下降。这意味着，于之前的某些调度周期内产生的四元组也是可能被选中的。这使得网络在环境探索上具有更多的可能性。

任务感知的环境探索。 为了保证 DL2 在前期可以充分探索环境，作者引入了熵 $H(\cdot)$ 作为正则化项，即梯度更新的公式变为

$$ \begin{equation} \begin{split} \theta \leftarrow \theta + \alpha \gamma^t \Big(\hat{G}(\tau_{t:T}) - V_\phi (s_t) \Big) \frac{\partial}{\partial \theta} \log \pi_\theta (a_t | s_t) + \beta H (\pi (\cdot | s; \theta)). \end{split} \end{equation} $$

除了添加正则化项，作者也引入了 $\epsilon$-greedy 策略来探索环境。在做每一次 inference 的时候，如果当前的状态 $s$ 是我们已经识别了的 pool state，就启用 $\epsilon$-greedy 策略：按照 $1-\epsilon$ 的概率，执行 policy network 的输出动作；按照 $\epsilon$ 的概率，直接执行一个特定的动作（这个动作来自之前的探索和分析）。

3.2 PS 和 Worker 的弹性伸缩

现有的深度学习框架并不支持 PS 和 worker 的动态调整。一种简单的实现是：基于 checkpoint 技术，将当前 DL 任务终止，并将其参数全部保存到文件中（checkpoint image），然后按照新的 PS 和 worker 的要求从 checkpoint image 中重新加载模型参数启动。这样做虽然简单，但是时间开销较大。

对此，DL2 是这样操作的。DL2 添加了一个协调器模块（coordinator），对于每一个 DL 任务，DL2 会为其启动一个 coordinator，负责该 DL 任务的 PS 和 worker 的动态伸缩。接下来，如图 7 所示，我以为某个 DL 任务添加一个 PS 为例阐述这个流程。

图 7 添加一个 PS 的流程。

  1. 注册（新的 PS）。 当添加一个新的 PS 时，这个 PS 需要将自身注册到 coordinator 中，然后这个 PS 将会获得 coordinator 为其分配的 ID、需要负责 averaging 的模型参数、以及需要建立连接的 PS 和 worker 的 ID。
  2. 模型参数分配（coordinator）。 每当 coordinator 收到一个 PS 的注册请求时，它首先更新自身所维护的 PS 和 worker 的列表，然后基于 best-fit 算法计算将要分配给该 PS 的模型参数。这个算法使得参数在这些 PS 上的移动开销尽可能小。Coordinator 还实现了一个版本的计数器（version counter）来保证模型参数迁移时的准确性和一致性。
  3. 模型参数迁移（其他 PS）。 对于每个 PS，当其收到来自 coordinator 的时钟信号时，就将自身的部分参数迁移至新的 PS。
  4. Worker 更新连接关系（worker）。 对于每个 worker，当其收到来自 coordinator 的时钟信号时，就暂停当前操作，等待 PS 的参数迁移完成。然后 worker 更新自身和这些 PS 的连接关系。

PS 的删除、worker 的添加 / 删除也遵循类似的步骤，此处不再展开。

4 总结

本文分析了面向 DL 任务的、基于 DRL 的调度器 DL2。虽然 DL2 采用的是通用的 Actor-Critic 算法，但是也设计了诸多 trick 来保证调度器的性能，例如熵正则化项、基于 $\epsilon$-greedy 的强制动作、使用经验池等。我为在未来几天补充在 Kubernetes 使用 DL2 的一些感受和实现细节。

本文关于 DL2 论文部分的图表全部来自此文，关于强化学习理论基础的图全部来自于教材 《神经网络与深度学习》。

转载申请

本作品采用 知识共享署名 4.0 国际许可协议 进行许可，转载时请注明原文链接。您必须给出适当的署名，并标明是否对本文作了修改。

  1. 没错！要用魔法打败魔法。 ↩︎

  2. 我在前文 面向深度学习的互适应调度 · OSDI'21 介绍过 PS-worker 架构，读者可以前往阅读此文。 ↩︎

  3. 这一部分的内容主要来自邱锡鹏编写的教材 《神经网络与深度学习》，此书电子版公开下载。 ↩︎

  4. 强化学习的目标是最大化累计奖赏。如果累计奖赏是无穷大，这个优化问题就不是良定义的了（没有意义）。 ↩︎

  5. 这和 GAN 的思想如出一辙。 ↩︎