本文解读 TreeQN 和 ATreeC,这两种创新的、基于树结构的深度强化学习模型。它们通过动态构建和学习抽象状态空间中的规划树,有效提升了智能体的决策能力。实验证明,TreeQN 和 ATreeC 在复杂任务和经典游戏中都展现出超越现有方法的卓越性能。
文章来自牛津大学的 Gregory Farquhar 等人,作者四人中现在有三人在 DeepMind 。
1th author: Gregory Farquhar - Google Scholar
paper: [1710.11417] TreeQN and ATreeC: Differentiable Tree-Structured Models for Deep Reinforcement Learning
code: oxwhirl/treeqn
1. 背景问题与动机
深度强化学习(Deep RL)近年来在许多领域取得了显著成功,尤其是结合深度学习来处理高维感知输入。其中,无模型方法(model-free)如 DQN, A2C 等,直接学习价值函数或策略,简单高效。然而,它们通常需要大量的试错经验,且缺乏对环境动力学的显式理解,难以进行前向预测和规划。
另一方面,基于模型的方法(model-based)试图学习环境的动力学模型(environment dynamics),然后利用学到的模型进行规划。理论上,这能更高效地利用经验,甚至在模型准确的情况下实现完美的规划。但实际中,在高维、复杂的环境中准确学习一个全局的、完全忠实的动力学模型极具挑战性,且模型误差会随着规划深度迅速累积,导致“幻觉”(hallucination)问题,严重影响规划效果。
那么,能否将两者结合,既享受无模型的端到端优化带来的便利,又利用模型进行局部的、短视角的规划,以提升样本效率和决策质量?
为了解决这些问题,Farquhar等人在2018年提出了 TreeQN 和 ATreeC ,思路是: 在价值函数(例如如 Figure 1 的抽象 DQN 结构)处理状态编码后增加一个可微分的树形结构,其用内部规划的方式替换掉原本输出 Q 值的全连接层。
2. TreeQN 的计算图
TreeQN 将传统的树状规划过程转化为一个可微分的神经网络结构。不同于仅仅使用一个预训练或独立训练的模型进行规划,TreeQN 将模型 components (转移函数、奖励函数、价值函数) 作为神经网络的一部分进行联合训练。在学习到的抽象状态空间中构建一个前向查找树,并通过可微分的方式回传梯度。
TreeQN 的整体结构可以视为 DQN 的一个扩展,它在状态编码
z
t
z_t
zt 和最终的 Q 值输出之间插入了一个递归的树状网络。如论文 Figure 2 所示,这个树状网络包含了多个层级的节点,每个节点代表一个预测的未来抽象状态。
2.1 TreeQN 的组成部分
-
状态编码器 (Encoder Function):
是 TreeQN 的起始点,负责将原始的状态观测 s t s_t st (通常是高维数据,如图像)编码成一个低维的抽象状态表示 z 0 ∣ t z_{0|t} z0∣t。
在论文的实现中,对于视觉观测,编码器通常由一系列卷积层构成 ,最后接一个全连接层。其数学表达可以概括为:
z 0 ∣ t = encode ( s t ) z_{0|t} = \text{encode}(s_t) z0∣t=encode(st)
这里的 encode ( ⋅ ) \text{encode}(\cdot) encode(⋅) 代表了整个编码网络的复杂非线性变换。 z 0 ∣ t z_{0|t} z0∣t 表示在时间 t t t 进行了 0 0 0 步内部规划后的状态(规划树的根节点)。 -
转移函数 (Transition Function):
预测在抽象状态下执行某个动作后,下一个抽象状态的表示。与传统的环境模型不同,这里的转移是在学习到的潜在空间中进行的。
对于树中的任意节点,其状态表示为 z l ∣ t z_{l|t} zl∣t(表示在时间 t t t 的真实状态基础上,经过了 l l l 步的内部规划转移, l l l 也表示节点的深度)。执行动作 a i ∈ A a_i \in \mathcal{A} ai∈A 后,预测的下一个状态表示为 z l + 1 ∣ t a i z_{l+1|t}^{a_i} zl+1∣tai :
z l + 1 ∣ t env = z l ∣ t + tanh ( W env z l ∣ t + b env ) z l + 1 ∣ t a i = z l + 1 ∣ t env + tanh ( W a i z l + 1 ∣ t env ) z_{l+1|t}^{\text{env}} = z_{l|t} + \tanh(W^{\text{env}} z_{l|t} + b^{\text{env}}) \\ z_{l+1|t}^{a_i} = z_{l+1|t}^{\text{env}} + \tanh(W^{a_i} z_{l+1|t}^{\text{env}}) zl+1∣tenv=zl∣t+tanh(Wenvzl∣t+benv)zl+1∣tai=zl+1∣tenv+tanh(Waizl+1∣tenv)
其中 W env , W a i ∈ R k × k W^{\text{env}}, W^{a_i} \in \mathbb{R}^{k \times k} Wenv,Wai∈Rk×k 和 b env ∈ R k b^{\text{env}} \in \mathbb{R}^k benv∈Rk 是可学习的参数, k k k 是抽象状态空间的维度。注意,对于同一个动作 a i a_i ai,其转移矩阵 W a i W^{a_i} Wai 在整个规划树的所有层级中都是共享的 。这体现了马尔可夫过程的 stationary property 归纳偏置。
为了防止状态表示的范数在递归应用中爆炸或消失,论文还在每次应用转移函数前对状态向量进行 L2 范数归一化 :
z l ∣ t : = z l ∣ t ∣ ∣ z l ∣ t ∣ ∣ 2 z_{l|t} := {z_{l|t} \over ||z_{l|t}||_2} zl∣t:=∣∣zl∣t∣∣2zl∣t
这个转移函数被递归地应用于树中的每个节点,直到达到预设的深度 d d d。 -
奖励函数 (Reward Function):
与转移函数并行,TreeQN 也学习一个奖励函数来预测在抽象状态下执行某个动作所获得的即时奖励。对于状态 z l ∣ t z_{l|t} zl∣t 和动作 a i a_i ai,预测的奖励记为 r ^ l ∣ t a i \hat{r}_{l|t}^{a_i} r^l∣tai。
论文中使用了一个带有 ReLU 激活的两层 MLP 来预测奖励 :
r ^ ( z l ∣ t ) = W 2 r ReLU ( W 1 r z l ∣ t + b 1 r ) + b 2 r \hat{r}(z_{l|t}) = W_2^r \text{ReLU}(W_1^r z_{l|t} + b_1^r) + b_2^r r^(zl∣t)=W2rReLU(W1rzl∣t+b1r)+b2r
其中 W 1 r ∈ R m × k , W 2 r ∈ R ∣ A ∣ × m W_1^r \in \mathbb{R}^{m \times k}, W_2^r \in \mathbb{R}^{|\mathcal{A}| \times m} W1r∈Rm×k,W2r∈R∣A∣×m 是可学习参数, m m m 是隐藏层维度, ∣ A ∣ |\mathcal{A}| ∣A∣ 是动作空间的数量。预测的奖励 r ^ l ∣ t a i \hat{r}_{l|t}^{a_i} r^l∣tai 是向量 r ^ ( z l ∣ t ) \hat{r}(z_{l|t}) r^(zl∣t) 的第 i i i 个元素。 -
价值函数 (Value Function):
该函数用于估计规划树叶子节点(即达到最大规划深度 d d d 时的抽象状态)的价值。论文中使用了一个简单的线性函数 :
V ( z ) = w ⊤ z + b V(z) = w^\top z + b V(z)=w⊤z+b
其中 w ∈ R k w \in \mathbb{R}^k w∈Rk 和 b ∈ R b\in\mathbb{R} b∈R 是可学习参数。这个价值函数在所有叶子节点上也是参数共享的。 -
回传函数 (Backup Function):
是 TreeQN 中连接预测奖励和价值估计,并计算最终 Q 值的关键。它通过树结构反向传播价值信息。论文中使用了一个基于 softmax 的函数来进行回传 :
b ( x ) = ∑ i x i softmax ( x ) i \text b(\boldsymbol x) = \sum_i x_i \text{softmax}(\boldsymbol x)_i b(x)=i∑xisoftmax(x)i
其中 x \boldsymbol x x 是一个向量。这个函数可以递归地应用于树的每个层级。
2.2 计算 Q 值
利用上面的组件,TreeQN 中深度为 l l l 的节点执行动作 a i a_i ai 的 Q 值估计 Q l ( z l ∣ t , a i ) Ql(z_{l∣t},a_i) Ql(zl∣t,ai) 可以递归计算:
-
叶子节点 ( l = d l=d l=d):
Q d ( z d ∣ t , a i ) = r ^ d ∣ t a i + γ V ( z d + 1 ∣ t a i ) Q^d(z_{d|t},a_i)=\hat{r}_{d|t}^{a_i}+\gamma V(z_{d+1|t}^{a_i}) Qd(zd∣t,ai)=r^d∣tai+γV(zd+1∣tai)
这里的 z d + 1 ∣ t a i z_{d+1∣t}^{a_i} zd+1∣tai 是通过转移函数从 z d ∣ t z_{d∣t} zd∣t 和 a i a_i ai 预测得到的抽象状态。 -
非叶子节点 ( l < d l<d l<d):
Q l ( z l ∣ t , a i ) = r ^ l ∣ t a i + γ [ ( 1 − λ ) V ( z l ∣ t a i ) + λ b ( Q l + 1 ( z l + 1 ∣ t a i , a j ) ) ] Q^l(z_{l|t},a_i)=\hat{r}_{l|t}^{a_i}+\gamma\left [ (1-\lambda)V(z_{l|t}^{a_i}) +\lambda\text b\left (Q^{l+1}(z_{l+1|t}^{a_i},a_j)\right )\right ] Ql(zl∣t,ai)=r^l∣tai+γ[(1−λ)V(zl∣tai)+λb(Ql+1(zl+1∣tai,aj))]
注意这里的 Q l + 1 ( z l + 1 ∣ t a i , a j ) Q^{l+1}(z_{l+1|t}^{a_{i}},a_{j}) Ql+1(zl+1∣tai,aj) 是一个向量,包含了从状态 z l + 1 ∣ t a i z_{l+1∣t}^{a_i} zl+1∣tai 执行所有可能动作 a j a_j aj 得到的下一层 Q 值估计。0 < λ < 1 0<\lambda <1 0<λ<1 是一个超参数,作者的思路是: 混合中间状态的价值 V ( z l ∣ t a i ) V(z_{l|t}^{a_i}) V(zl∣tai) 和 b \text b b 回传的价值,而非仅使用回传,这样可以使中间状态的价值更有实际意义。
如果 λ = 1 \lambda=1 λ=1 且回传函数 b \text b b 是标准的最大值函数 max \max max 时,退化为标准的 Bellman 回传 :
Q l ( z l ∣ t , a i ) = r ^ l ∣ t a i + γ max a ′ Q l + 1 ( z l + 1 ∣ t a i , a ′ ) Q^l(z_{l|t}, a_i) = \hat{r}_{l|t}^{a_i} + \gamma \max_{a'} Q^{l+1}(z_{l+1|t}^{a_i}, a') Ql(zl∣t,ai)=r^l∣tai+γa′maxQl+1(zl+1∣tai,a′)
使用 softmax 的回传函数而非 hard max 的好处在于,它允许梯度流经所有分支,而不仅仅是最大 Q 值对应的最优分支 。这使得模型能够从所有可能的规划路径中学习,潜在的提高训练的稳定性并减少对最优路径预测准确性的过度依赖。
通过将上述所有模块组合起来,TreeQN 构建了一个从原始状态 s t s_t st 到最终每个动作的 Q 值估计 Q ( s t , a ) Q(s_t, a) Q(st,a) 的可微分计算图。在训练过程中,梯度可以直接从最终的 Q 值损失回传到编码器、转移函数、奖励函数和价值函数的参数。这种端到端的训练方式确保了学习到的模型 components 能够有效地服务于在线规划,以提高 Q 值估计的准确性。
3. 训练与 Grounding
TreeQN 架构的完全可微分的特性意味着可以将 TreeQN 直接“插入”到现有的离散动作空间的深度强化学习算法中,作为价值函数网络(例如 DQN)或策略网络(例如 A2C)的替代品,并进行端到端的训练。这种训练确保了学习到的模型 components 是为了提升最终的 RL 性能而优化的,避免了独立训练模型可能导致的失配问题。
3.1 损失函数
论文主要在两种流行的算法框架下评估了 TreeQN:n-step Q-learning 和 A2C。
在 n-step Q-learning 框架下,TreeQN 取代了标准 DQN 的最后全连接层,直接输出每个离散动作的 Q 值 Q ( s t , a ) Q(s_t, a) Q(st,a)。训练的目标是最小化预测 Q 值与目标值之间的误差。n-step Q-learning 的目标值结合了实际观测到的奖励和未来状态的 Q 值估计。具体而言,对于从时间步 t t t 开始的 n-step 转移,损失函数通常包括一个关于 Q 值的均方误差项。使用 TreeQN 后,这个损失函数变为:
L
nstep-Q
=
∑
envs
∑
j
=
1
n
(
Q
(
s
t
+
n
−
j
,
a
t
+
n
−
j
,
θ
)
−
(
∑
k
=
1
j
[
γ
j
−
k
r
t
+
n
−
k
]
+
γ
j
max
a
′
Q
(
s
t
+
n
,
a
′
,
θ
−
)
)
⏟
n-step TD target
)
2
\mathcal{L}_{\text{nstep-Q}}=\sum_{\text{envs}}\sum_{j=1}^n \left( Q\left(\mathbf{s}_{t+n-j},a_{t+n-j},\theta\right)- \underbrace{\left( \sum_{k=1}^j\left[\gamma^{j-k}r_{t+n-k}\right]+\gamma^j\max_{a^{\prime}}Q\left(\mathbf{s}_{t+n},a^{\prime},\theta^-\right)\right) } _{\text{n-step TD target}} \right)^2
Lnstep-Q=envs∑j=1∑n
Q(st+n−j,at+n−j,θ)−n-step TD target
(k=1∑j[γj−krt+n−k]+γja′maxQ(st+n,a′,θ−))
2
解释:
-
∑ envs \sum_{\text{envs}} ∑envs: 这表示对所有的并行环境进行求和。在异步强化学习中,通常会同时运行多个环境,论文中每次搜索 (可以并行) 都可以看作一个环境。
-
∑ j = 1 n \sum_{j=1}^n ∑j=1n: 这表示对从 j = 1 j=1 j=1 到 n n n 的步数进行求和。这是 n-step 的核心思想。不只计算一步的 TD 误差,而是考虑从当前状态到未来 n 步的状态的TD误差。这里的 j j j 表示从当前状态 s t + n − j \mathbf{s}_{t+n-j} st+n−j 开始,往前看 j j j 步。
-
∑ k = 1 j [ γ j − k r t + n − k ] \sum_{k=1}^j\left[\gamma^{j-k}r_{t+n-k}\right] ∑k=1j[γj−krt+n−k]: 这是从状态 s t + n − j \mathbf{s}_{t+n-j} st+n−j 开始,未来 j j j 步的折扣累积回报。
- r t + n − k r_{t+n-k} rt+n−k 是在时刻 t + n − k t+n-k t+n−k 获得的奖励。
- γ \gamma γ 是折扣因子,用于衡量未来奖励的重要性。
- γ j − k \gamma^{j-k} γj−k 是对应的折扣系数,距离当前状态越远的奖励,折扣越大。
- 这个求和项计算的是从状态 t + n − j t+n-j t+n−j 开始,到状态 t + n − 1 t+n-1 t+n−1 为止,这 j − 1 j-1 j−1 步获得的折扣奖励。
-
γ j max a ′ Q ( s t + n , a ′ , θ − ) \gamma^j\max_{a^{\prime}}Q\left(\mathbf{s}_{t+n},a^{\prime},\theta^-\right) γjmaxa′Q(st+n,a′,θ−): 这是在未来状态 s t + n \mathbf{s}_{t+n} st+n 的目标 Q 值,用目标网络 θ − \theta^- θ− 计算。
- s t + n \mathbf{s}_{t+n} st+n 是从状态 s t + n − j \mathbf{s}_{t+n-j} st+n−j 经过 j j j 步后的状态。
- max a ′ Q ( s t + n , a ′ , θ − ) \max_{a^{\prime}}Q\left(\mathbf{s}_{t+n},a^{\prime},\theta^-\right) maxa′Q(st+n,a′,θ−) 表示在状态 s t + n \mathbf{s}_{t+n} st+n 下,使用目标网络 θ − \theta^- θ− 评估所有可能行动的 Q 值,并取最大值。这代表了在未来状态 s t + n \mathbf{s}_{t+n} st+n 下能获得的最佳期望回报。
- γ j \gamma^j γj 是对应的折扣因子,因为这是从状态 s t + n − j \mathbf{s}_{t+n-j} st+n−j 经过 j j j 步到达的状态 s t + n \mathbf{s}_{t+n} st+n 的价值。
-
Q ( s t + n − j , a t + n − j , θ ) Q\left(\mathbf{s}_{t+n-j},a_{t+n-j},\theta\right) Q(st+n−j,at+n−j,θ): 这是在状态 s t + n − j \mathbf{s}_{t+n-j} st+n−j 下,执行动作 a t + n − j a_{t+n-j} at+n−j 的当前 Q 值,用当前网络 θ \theta θ 计算。
整个损失函数的核心思想是:
对于每个环境和每个 j j j (从 1 到 n),我们计算一个 TD 误差。这个误差是 n 步实际累积回报 + n 步后状态的折扣最大 Q 值(这是我们的目标值,由目标网络计算)与 当前状态的 Q 值(由当前网络计算)之间的差。然后对所有环境和所有 j j j 的这个误差的平方进行求和,作为要最小化的损失。
理解 n-step Q-learning。由标准 Q-learning (使用 One-step TD) 推广而来。 回顾 Q-learning:
E [ U t ∣ S t = s t , A t = a t ⋆ ] = max a Q ⋆ ( s t , a ) \mathbb E[U_t|S_t=s_t,A_t=a^{\star}_t]=\max_aQ^\star(s_t,a) E[Ut∣St=st,At=at⋆]=amaxQ⋆(st,a)
关键是知道 t t t 时刻的回报 U t U_t Ut 可以写作:
U t = R t + γ ⋅ U t + 1 = R t + γ ⋅ ( R t + 1 + γ ⋅ U t + 2 ) = R t + γ ⋅ R t + 1 + γ 2 ⋅ U t + 2 = R t + γ ⋅ R t + 1 + γ 2 ⋅ R t + 2 + ⋯ + γ n ⋅ U t + n = ∑ i = 0 n − 1 γ i ⋅ R t + i + γ n ⋅ U t + n \begin{align*} U_t&=R_t+\gamma \cdot U_{t+1} \\ &= R_t+\gamma \cdot (R_{t+1}+\gamma\cdot U_{t+2}) \\ &= R_t + \gamma \cdot R_{t+1} + \gamma^2 \cdot U_{t+2} \\ &= R_t + \gamma \cdot R_{t+1} + \gamma^2 \cdot R_{t+2} + \cdots + \gamma^n \cdot U_{t+n} \\ &= \sum_{i=0}^{n-1}{\gamma^i\cdot R_{t+i}} + \gamma^n\cdot U_{t+n} \end{align*} Ut=Rt+γ⋅Ut+1=Rt+γ⋅(Rt+1+γ⋅Ut+2)=Rt+γ⋅Rt+1+γ2⋅Ut+2=Rt+γ⋅Rt+1+γ2⋅Rt+2+⋯+γn⋅Ut+n=i=0∑n−1γi⋅Rt+i+γn⋅Ut+n
3.2 辅助损失
尽管端到端训练能够使模型 components 为 RL 任务服务,但学习到的抽象状态转移和奖励预测可能与真实环境并不完全对应,缺乏“接地”(Grounding)。为了鼓励 TreeQN 的内部规划过程更接近真实环境的动力学和奖励,论文引入了辅助损失(Auxiliary Losses)。这些辅助损失在训练 RL 主目标的同时,额外惩罚模型预测与真实观测之间的差异。
使用辅助损失,TreeQN 可以被视为一种 model-based 方法。而从另一个角度看,如果完全不考虑 Grounding,则可以被视为一种 model-free 的方法。
论文研究了两种主要的辅助损失:
-
奖励 Grounding (Reward Grounding):
这个辅助任旨在训练 TreeQN 的奖励函数 r ^ \hat{r} r^,使其预测值接近真实环境提供的即时奖励。对于规划树中在时间步 t t t 经过 l l l 步内部转移、对应于真实动作序列 { a t , a t + 1 , … , a t + l − 1 } \{a_t, a_{t+1}, \dots, a_{t+l-1}\} {at,at+1,…,at+l−1} 的节点,其预测奖励 r ^ l ∣ t a t + l − 1 \hat{r}_{l|t}^{a_{t+l-1}} r^l∣tat+l−1 被强制接近真实的奖励 r t + l − 1 r_{t+l-1} rt+l−1。
考虑到 n-step Q-learning 的训练批量包含多个时间步的转移,奖励 Grounding 损失被定义为:
L r = η r ∑ envs ∑ j = 1 n ∑ l = 1 min ( d , n − j + 1 ) ( r ^ l ∣ t + j a t + j : t + j + l − 1 − r t + j + l − 1 ) 2 \mathcal{L}_{r} = \eta_r \sum_{\text{envs}} \sum_{j=1}^n \sum_{l=1}^{\min(d, n-j+1)} (\hat{r}_{l|t+j}^{a_{t+j:t+j+l-1}} - r_{t+j+l-1})^2 Lr=ηrenvs∑j=1∑nl=1∑min(d,n−j+1)(r^l∣t+jat+j:t+j+l−1−rt+j+l−1)2
其中 η r \eta_r ηr 是控制奖励 Grounding 强度的超参数。外层的两个求和项遍历了并行环境和 n-step 转移序列中的每个时间步 t + j t+j t+j。内层的求和遍历了规划树的深度 l l l,但只考虑那些其对应真实奖励 r t + j + l − 1 r_{t+j+l-1} rt+j+l−1 已经被观测到的节点( min \min min 函数,使 l l l 不超过当前 n-step 序列的剩余步数)。 r ^ l ∣ t + j a t + j : t + j + l − 1 \hat{r}_{l|t+j}^{a_{t+j:t+j+l-1}} r^l∣t+jat+j:t+j+l−1 表示在时间步 t + j t+j t+j 对应的状态下,沿着真实采取的动作序列 { a t + j , … , a t + j + l − 1 } \{a_{t+j}, \dots, a_{t+j+l-1}\} {at+j,…,at+j+l−1} 在规划树中预测的第 l l l 步的奖励。 -
状态 Grounding (State Grounding):
这个辅助任务旨在训练 TreeQN 的转移函数,使其预测的抽象状态接近真实状态的编码。对于规划树中在时间步 t t t 经过 l l l 步内部转移、对应于真实动作序列 { a t , a t + 1 , … , a t + l − 1 } \{a_t, a_{t+1}, \dots, a_{t+l-1}\} {at,at+1,…,at+l−1} 的节点,其预测状态表示 z l ∣ t a t : t + l − 1 z_{l|t}^{a_{t:t+l-1}} zl∣tat:t+l−1 被强制接近真实环境在时间步 t + l t+l t+l 的状态编码 z 0 ∣ t + l z_{0|t+l} z0∣t+l。
状态 Grounding 损失被定义为:
L s = η s ∑ envs ∑ j = 1 n ∑ l = 1 d ( z l ∣ t + j a t + j : t + j + l − 1 − z 0 ∣ t + j + l ) 2 \mathcal{L}_{s} = \eta_s \sum_{\text{envs}} \sum_{j=1}^n \sum_{l=1}^d (z_{l|t+j}^{a_{t+j:t+j+l-1}} - z_{0|t+j+l})^2 Ls=ηsenvs∑j=1∑nl=1∑d(zl∣t+jat+j:t+j+l−1−z0∣t+j+l)2
结构与奖励 Grounding 类似。 z l ∣ t + j a t + j : t + j + l − 1 z_{l|t+j}^{a_{t+j:t+j+l-1}} zl∣t+jat+j:t+j+l−1 表示在时间步 t + j t+j t+j 对应的状态下,沿着真实采取的动作序列在规划树中预测的第 l l l 步的抽象状态表示。 z 0 ∣ t + j + l z_{0|t+j+l} z0∣t+j+l 是真实环境在时间步 t + j + l t+j+l t+j+l 的状态经过编码器得到的表示。
总的训练损失可以是主要的 RL 损失与辅助损失的加权和。例如,在 n-step Q-learning 中,总损失为:
L
total
=
L
nstep-Q
+
L
r
+
L
s
\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{nstep-Q}} + \mathcal{L}_{r} + \mathcal{L}_{s}
Ltotal=Lnstep-Q+Lr+Ls
通过同时优化主 RL 目标和这些辅助目标,TreeQN 被鼓励学习到一个不仅能预测准确 Q 值,而且其内部规划过程(转移和奖励预测)也能在一定程度上反映真实环境的模型。论文通过实验发现,奖励 Grounding 对性能有积极影响,而简单的潜在状态 Grounding(L2 Loss)并没有带来提升,这表明在抽象空间中进行状态 Grounding 仍是一个有待深入研究的问题。
4. ATreeC 架构
TreeQN 成功地将可微分树状规划集成到基于价值的方法(如 DQN)中。自然而然地,这种思想也可以应用于基于策略的方法,尤其是 Actor-Critic 框架。ATreeC (Actor-TreeC) 正是 TreeQN 在 Actor-Critic 领域的扩展。它的目标是学习一个直接输出动作概率的策略,同时利用树状结构赋予策略一定的“远见”或规划能力。
在 Actor-Critic 框架下,通常包含两个主要部分:一个 Actor(策略网络)负责选择动作,一个 Critic(价值网络)负责评估当前状态或状态-动作对的价值,并用这个价值信号来指导 Actor 的学习。ATreeC 将 TreeQN 的树状结构用于构建策略网络(Actor)。
4.1 ATreeC 的组成部分
-
状态编码器 (Encoder Function):
与 TreeQN 相同,ATreeC 也使用一个状态编码器将原始状态 s t s_t st 转换为抽象状态表示 z t z_t zt。这个编码器通常与 Critic 的编码器共享参数。
z t = encode ( s t ) z_t = \text{encode}(s_t) zt=encode(st) -
策略网络 (Actor):
ATreeC 的策略网络是基于 TreeQN 的树状结构构建的。唯一不同是在输出 Q 向量后添加了一个 softmax 层,将其转化为一个离散动作的概率分布 π ( a ∣ s t ) \pi(a|s_t) π(a∣st)。
Q ( s t , a ) = TreeQN ( z t ) π ( a ∣ s t ) = softmax ( Q ( s t , a ) ) Q(s_t, a) = \text{TreeQN}(z_t)\\ \pi(a|s_t) = \text{softmax}(Q(s_t, a)) Q(st,a)=TreeQN(zt)π(a∣st)=softmax(Q(st,a)) -
价值网络 (Critic):
ATreeC 的 Critic 负责估计状态价值 V ( s ) V(s) V(s)。在论文的实现中,Critic 网络结构相对简单 (也可以用 TreeQN ),它共享 Actor 的状态编码器,然后接一个全连接层来预测一个标量状态价值。
V cr ( s t ) = w cr ⊤ z t + b cr V_{\text{cr}}(s_t) = w_{\text{cr}}^\top z_t + b_{\text{cr}} Vcr(st)=wcr⊤zt+bcr
其中 w cr w_{\text{cr}} wcr 和 b cr b_{\text{cr}} bcr 是 Critic 独有的可学习参数。Critic 的训练目标是最小化预测价值与实际回报(或其估计)之间的误差。
4.2 损失函数
ATreeC 的训练通常在 Actor-Critic 算法框架下进行,例如异步优势 Actor-Critic (A3C) 或其同步变种 A2C。论文中使用了 A2C 进行训练。A2C 的训练目标是同时优化策略(最大化期望回报)和价值函数(最小化价值预测误差)。
在 A2C 框架下,Actor 的参数更新方向由优势函数(Advantage Function)决定,优势函数衡量了在当前状态下采取某个动作相对于平均水平的“好坏”。优势函数通常定义为
A
(
s
t
,
a
t
)
=
Q
(
s
t
,
a
t
)
−
V
(
s
t
)
A(s_t, a_t) = Q(s_t, a_t) - V(s_t)
A(st,at)=Q(st,at)−V(st)。在 A2C 中,使用 n-step return 来估计 Q 值,并使用 Critic 的价值估计作为基线:
A
t
=
∑
k
=
0
n
−
1
γ
k
r
t
+
k
+
γ
n
V
cr
(
s
t
+
n
)
−
V
cr
(
s
t
)
A_t = \sum_{k=0}^{n-1} \gamma^k r_{t+k} + \gamma^n V_{\text{cr}}(s_{t+n}) - V_{\text{cr}}(s_t)
At=k=0∑n−1γkrt+k+γnVcr(st+n)−Vcr(st)
Actor 的策略梯度损失函数包含策略的对数概率项乘以优势函数,以及一个鼓励探索的熵正则项:
L
actor
=
−
E
[
log
π
(
a
t
∣
s
t
;
θ
actor
)
A
t
+
β
H
(
π
(
⋅
∣
s
t
;
θ
actor
)
)
]
\mathcal{L}_{\text{actor}} = - \mathbb{E} \left[ \log \pi(a_t|s_t; \theta_{\text{actor}}) A_t + \beta H(\pi(\cdot|s_t; \theta_{\text{actor}})) \right]
Lactor=−E[logπ(at∣st;θactor)At+βH(π(⋅∣st;θactor))]
其中
θ
actor
\theta_{\text{actor}}
θactor 是 Actor 网络的参数,
H
H
H 是策略的熵,
β
\beta
β 是熵正则化系数。
Critic 的价值函数通过最小化预测价值与 n-step return 之间的均方误差进行训练:
L
critic
=
E
[
(
V
cr
(
s
t
;
θ
critic
)
−
(
∑
k
=
0
n
−
1
γ
k
r
t
+
k
+
γ
n
V
cr
(
s
t
+
n
)
)
)
2
]
\mathcal{L}_{\text{critic}} = \mathbb{E} \left[ (V_{\text{cr}}(s_t; \theta_{\text{critic}}) - (\sum_{k=0}^{n-1} \gamma^k r_{t+k} + \gamma^n V_{\text{cr}}(s_{t+n})))^2 \right]
Lcritic=E[(Vcr(st;θcritic)−(k=0∑n−1γkrt+k+γnVcr(st+n)))2]
与 TreeQN 类似,ATreeC 也可以引入奖励 Grounding 和状态 Grounding 的辅助损失来帮助学习到的模型 components 更贴近真实环境。辅助损失的数学形式与第三节中描述的类似,但现在它们是添加到 Actor 和 Critic 的总损失中进行联合优化。
ATreeC 将 TreeQN 的可微分规划能力赋予了策略网络,使得策略在选择动作时能够利用内部的预测树进行多步的考量。实验结果表明,ATreeC 在某些任务上相对于基线 A2C 取得了更好的性能,尤其是在那些需要一定前瞻性才能做出最优决策的环境中。Actor-Critic 框架下的训练也可能带来不同的优势,例如策略梯度方法在处理随机性方面可能更具优势。
5. 总结与未来方向
论文探索了一种将深度 model-free 强化学习与在线规划深度融合的新范式。通过 TreeQN 和 ATreeC 这两个具有可微分树状结构的神经网络模型,他们有效地解决了在复杂环境中学习精确模型并将其用于规划的挑战。
核心贡献在于构建了一个端到端可微分的规划树,使得模型 Components (状态转移、奖励预测、价值函数) 可以直接通过强化学习的主目标进行优化。这种优化方式确保了学习到的模型对于其在规划树中的实际用途是有效的。TreeQN 作为价值函数网络,通过预测和回溯未来状态的 Q 值来指导决策;而 ATreeC 则将相同的树状结构应用于策略网络,使得策略能够基于多步规划的結果来输出动作概率。
通过在 Box Pushing 和 Atari 游戏上的实验,论文验证了 TreeQN 和 ATreeC 的有效性。
在 Box Pushing 环境中,TreeQN 和 ATreeC 均显著优于其对应的基线(DQN 和 A2C)。这表明在动力学相对容易学习的环境中,TreeQN 和 ATreeC 结构所引入的归纳偏置和规划能力能够有效地加速学习并获得更好的策略。不同深度的树在 TreeQN 中表现出差异,更深的树带来了性能提升,而在 ATreeC 中这种差异不明显,这可能与环境复杂度或 Actor-Critic 训练的特性有关。
在更具挑战性的 Atari 环境中,TreeQN 在多数游戏上超越了 DQN 和 VPN。ATreeC 也始终与 A2C 持平或优于 A2C。实验结果表明,TreeQN 和 ATreeC 的优势在那些需要一定前瞻性思考才能做出最优决策的游戏中尤为突出,例如 Alien 和 Seaquest。与简单增加模型容量的基线(DQN-Deep 和 DQN-Wide)相比,TreeQN 的结构和辅助损失带来的提升更为显著。然而,论文也指出,在一些探索性要求较高的游戏(如 Krull 和 Frostbite)中,所有方法的表现都不尽理想,这提示了 TreeQN 和 ATreeC 在探索机制上的潜在不足。
关于模型 Grounding 的辅助损失,实验表明奖励 Grounding 对性能有积极影响,这符合直觉,因为奖励是 RL 任务中与目标直接相关的信号。然而,简单的潜在状态 Grounding 并没有带来性能提升,这可能是因为潜在空间的误差最小化与最终的 RL 目标存在不一致,或者需要更复杂的机制来确保潜在状态的语义对应性。
TreeQN 和 ATreeC 的主要优势在于:
- 端到端可微分规划: 将规划过程融入神经网络,通过联合训练优化所有 Components,避免了传统 Model-Based RL 中模型训练与策略训练的脱节。
- 结构化的归纳偏置: 树状结构和参数共享编码了马尔可夫过程的先验知识,有助于学习更有效的状态表示和价值函数。
- 抽象状态空间规划: 在学习到的低维空间中进行规划,降低了建模难度,提高了效率。
潜在的局限性包括:
- 规划深度限制: 受限于计算资源,规划树的深度不能无限增加,这可能限制了其解决需要长远规划任务的能力。
- 抽象模型的可解释性与准确性: 学习到的抽象状态转移和奖励预测可能与真实环境存在偏差,尤其是缺乏强 Grounding 时,内部规划过程的意义可能不明确。
- 探索挑战: 在某些需要充分探索的环境中,TreeQN 和 ATreeC 并未展现出显著优势,可能需要结合专门的探索机制。
论文的结论部分也为未来的研究指明了方向:
- 更深层树的优化: 如何更有效地训练和优化具有更深规划树的模型,以处理更复杂的长期依赖问题。
- 可解释的规划: 探索新的辅助目标或模型结构,鼓励学习到的模型 components 能够产生更具可解释性的内部规划过程,同时不牺牲性能。
- 智能探索机制的融合: 将 TreeQN 和 ATreeC 与先进的探索策略相结合,提升其在探索性环境中的表现。
- 连续动作空间的扩展: 当前模型仅针对离散动作空间,将其扩展到连续动作空间是另一个重要的研究方向。
总而言之,TreeQN 和 ATreeC 提供了一种将 model-free 学习与在线规划有机结合的富有前景的方法。它们的可微分树状结构为深度强化学习模型注入了新的活力,展示了在复杂环境中实现更具“前瞻性”智能体的潜力。