Offline RL : Context-Former: Stitching via Latent Conditioned Sequence Modeling

paper

基于HIM的离线RL算法，解决基于序列模型的离线强化学习算法缺乏对序列拼接能力。

Intro

文章提出了ContextFormer，旨在解决决策变换器（Decision Transformer, DT）在轨迹拼接（stitching）能力上的不足。轨迹拼接是离线RL中一个重要的能力，它允许算法通过组合次优的轨迹片段来获得更优的策略。ContextFormer通过集成基于上下文信息的模仿学习（Imitation Learning, IL）和序列建模，模仿有限数量专家轨迹的表示，来实现次优轨迹片段的拼接。实验结果表明，ContextFormer在多模仿学习设置下具有竞争力，并且在与其他DT变体的比较中表现出色。

两个定义

上述两个定义分别给出基于隐变量的条件序列模型建模方式，以及使用专家序列，通过度量经过embedding后的变量距离，使得待优化策略应满足靠近专家策略，远离次优轨迹策略。对于定义二有如下形式化的目标来优化上下文隐变量表征
$$\begin{gathered} {\mathcal{J}}_{{\mathbf{z}}^{\ast }}=\underset{{\mathbf{z}}^{\ast },I_{\varphi }}{\min }{\mathbb{E}}_{{\tau }^{\ast }\sim {\pi }^{\ast }(\tau )}[\Vert {\mathbf{z}}^{\ast }-I_{\varphi }({\tau }^{\ast })\Vert ] \\ -{\mathbb{E}}_{\hat{\tau }\sim \hat{\pi }}[\Vert {\mathbf{z}}^{\ast }-I_{\varphi }(\hat{\tau })\Vert ], \end{gathered}$$

Method

ContextFormer的训练过程包括两个关键模型：Hindsight Information Extractor $I_{\varphi }$和Contextual Policy。Hindsight Information Extractor使用BERT作为编码器，并采用VQ-VAE（Vector Quantization Variational Autoencoder）损失来训练。Contextual Policy则是一个基于潜在条件的序列模型（DT），通过上下文信息作为目标来优化策略接近专家策略。

根据定义4.1建模序列模型以及$I_{\varphi }$，通过监督学习方式优化上下文策略${\pi }_z$以及HI extractor。
$${\mathcal{J}}_{{\pi }_{\mathbf{z}},I_{\varphi }}={\mathbb{E}}_{\tau \sim ({\pi }^{\ast },\hat{\pi })}[\Vert {\pi }_{\mathbf{z}}(\cdot |I_{\varphi }(\tau ),{\mathbf{s}}_0,{\mathbf{a}}_0,\cdots ,I_{\varphi }(\tau ),{\mathbf{s}}_t)-{\mathbf{a}}_t\Vert ],(4)$$
其中$\hat{\pi }\mathrm{and}{\pi }^{\ast }$分别表示次优策略以及专家策略。同时，基于定义4.2对$I_{\varphi }$以及上下文embedding $z^{\ast }$进行优化。
$${\mathcal{J}}_{{\mathbf{z}}^{\ast },I_{\varphi }}=\underset{{\mathbf{z}}^{\ast },I_{\varphi }}{\min }{\mathbb{E}}_{\hat{\tau }\sim \hat{\pi }(\tau ),{\tau }^{\ast }\sim {\pi }^{\ast }(\tau )}[\Vert {\mathbf{z}}^{\ast }-I_{\varphi }({\tau }^{\ast })\Vert -||{\mathbf{z}}^{\ast }-I_{\varphi }(\hat{\tau })||]（5）$$

除此外，对于$I_{\varphi }$还需VQ-loss进行优化，三者联合构成了VQ-VAE的训练损失函数。

伪代码

（伪代码Training部分的第二步，VQ-loss应对应公式20）

结果