Representation RL : DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototyp

ICML 2022
paper
code

Intro

Dreamer系列的model-based的方法需要通过解码器，将状态$z_t=[h_t,s_t]$重构为观测状态进而优化世界模型。这导致面对高位图像输入时计算复杂度高。
$${\mathcal{J}}_{\mathrm{DREAMER}}=\sum _{t=1}^T{\mathbb{E}}_q[\underset{{\mathcal{J}}_{\mathrm{o}}^t}{\underbrace{\log p(o_t\mid h_t,s_t)}}+\underset{{\mathcal{J}}_{\mathrm{R}}^t}{\underbrace{\log p(r_t\mid h_t,s_t)}}-\underset{{\mathcal{J}}_{\mathrm{KL}}^t}{\underbrace{D_{\mathrm{KL}}(q(s_t\mid h_t,o_t)\parallel p(s_t\mid h_t))}}]$$
然而，事实上在策略学习过程中仅仅需要正确的奖励模型以及状态转移函数，并没有利用这一重构模型。因此，本文算法提出基于非对比的无监督表征学习方法SwAV结合没有状态重构的Dreamer算法。有效降低计算复杂度以及提高鲁棒性。

Method

算法首先基于聚类得到k个经过L2-归一化的原型 { c 1 , c 2 , . . . c k } \{c_1,c_2,...c_k \} {c1​,c2​,...ck​}。然后从Buffer中抽样观测序列并基于双线性插值的random shifts得到两种增强序列$o_{1:T}^{(1)},o_{1:T}^{(2)}$。对于每一个增强序列的观测状态，将通过RSSM模型得到表征$z_{1:T}^{(i)}$,然后线性映射为同样L2-归一化的$x_t^i$。与k原型点积以及softmax操作后得到相似程度(分类)：
$$(u_{t,1}^{(i)},\dots ,u_{t,K}^{(i)})=\text{softmax}\left(\frac{x_t^{(i)}\cdot c_1}{\tau },\dots ,\frac{x_t^{(i)}\cdot c_K}{\tau }\right)$$

类似的为了预测增强观测序列匹配的原型，共享RSSM中的卷积编码器，然后经过线性映射$f$得到$y_t$。与原型的相似程度同理
$$(v_{t,1}^{(i)},\dots ,v_{t,K}^{(i)})=\mathrm{softmax}\left(\frac{y_t^{(i)}\cdot c_1}{\tau },\dots ,\frac{y_t^{(i)}\cdot c_K}{\tau }\right),$$

上述两个值预测的是相似度，具体属于哪个标签将采用Sinkhorn-Knopp算法确定。首先，$\bar{\theta }$通过$\theta$ soft update得到。然后给出增强序列的score：$({\bar{y}}_t^{(i)}\cdot c_1,\dots ,{\bar{y}}_t^{(i)}\cdot c_K),\mathrm{where}{\bar{y}}_t^{(i)}=f_{\theta }(o_t^{(i)})$。最后Sinkhorn-Knopp算法分别应用于两个增强序列，得到目标$(w_{t,1}^{(i)},\dots ,w_{t,K}^{(i)})$, 实现聚类目标的均匀分配

综上，算法的表征学习的目标便是最大化下式
$$\begin{array}{c}{\mathcal{J}}_{\mathrm{SWAV}}^t=\frac{1}{2}\sum _{k=1}^K\left(w_{t,k}^{(1)}\log v_{t,k}^{(2)}+w_{t,k}^{(2)}\log v_{t,k}^{(1)}\right),\\ {\mathcal{J}}_{\mathrm{Temp}}^t=\frac{1}{2}\sum _{k=1}^K\left(w_{t,k}^{(1)}\log u_{t,k}^{(1)}+w_{t,k}^{(2)}\log u_{t,k}^{(2)}\right).\end{array}$$
第一项通过对齐不同增强方法后的图像提高表征能力，然而，它忽略了强化学习中至关重要的时间结构。因此，添加了一个第二项 $J_{Temp}^t$，它改进了来自同一状态的预测。使原型接近那些总结了过去观测以及动作的状态（增加w与u的关联），从而将时间结构提炼到原型中。这两项将替代Dreamer中状态重构的损失函数，因此，最终优化函数为
$${\mathcal{J}}_{\mathrm{DREAMERPRO}}=\sum _{t=1}^T{\mathbb{E}}_q[{\mathcal{J}}_{\mathrm{SWAV}}^t+{\mathcal{J}}_{\mathrm{Temp}}^t+{\mathcal{J}}_{\mathrm{R}}^t-{\mathcal{J}}_{\mathrm{KL}}^t],$$

结果

两个表征学习损失函数的消融