Pretrain RL:Pre-Training Goal-based Models for Sample-Efficient Reinforcement Learning

ICLR 2024 Oral
paper

Introduction

在大型任务无关数据集上预训练可以帮助RL提高在复杂任务上的样本采样效率。本文提出的PTGM便是构造为下游任务生成goal的上层策略，指导下游任务的goal-conditioned 低级高效探索利用。而为了解决高维空间产生的高维goal以及保持技能表示，提出基于clustering形成离散化的high-level策略的动作空间。此外，在线RL过程中，引入一个预训练的goal prior模型，通过KL散度对high-level的策略进行正则化，保证在线训练的稳定高效。

Method

假设任务无关数据集来自相同环境、不同任务且由状态动作对构成 D = { τ = { ( s i , a i ) } i = 0 T } D=\{\tau=\{(s_{i},a_{i})\}_{i=0}^{T}\} D={τ={(si​,ai​)}i=0T​}。PTGM首先预训练一个goal-conditioned的低级策略$P_{\varphi }(a_t|s_t,s^g)$。PTGM从数据集中采样k-steps的连续数据，将最后一步的状态作为goal，然后通过最小化负对数似然训练参数：
$$\mathcal{L}(\varphi )={\mathbb{E}}_D\left[-\log P_{\varphi }(a_i|s_i,s^g)\right]$$

CLUSTERING IN THE GOAL Space

为了解决高维goal引入高维状态空间导致RL的采样效率降低，算法从数据集采用t-SNE对状态降维，利用K-means得到N个聚类中心 G = { i : s i g } i = 1 N , G=\{i:s_{i}^{g}\}_{i=1}^{N}, G={i:sig​}i=1N​,。

PRE-TRAINING THE GOAL PRIOR MODEL

由于高级策略缺乏goal相关的先验知识，因此文章基于数据集训练一个先验模型。
goal先验模型${\pi }_{\psi }^p(a^h|s)$与高级策略结构相同，其中$a^h$表示goal聚类中心的index。根据index选择对应的聚类中心，便实现预测当前状态下未来goal的分布。

模型训练时，类似于训练low-level策略时，选择k-steps后的状态作为goal $s^g$，采用余弦相似度得到监督训练的标签：$a^h=\arg {\max }_{i\in [N]}\left(\frac{s_i^g\cdot s^g}{\Vert s_i^g\Vert \cdot \Vert s^g\Vert }\right)$

而连续空间下，通过最小化负对数似然优化：
$$\mathcal{L}(\psi )={\mathbb{E}}_D\left[-\log {\pi }_{\psi }^p(a^h|s_t)\right].$$
goal先验模型将作为 RL期间高级策略的正则化器，提供内在奖励来指导代理对数据集中可能目标的探索。

REINFORCEMENT LEARNING WITH PTGM

基于goal cluster G、预训练low-level policy $P_{\varphi }$与goal prior model ${\pi }_{\psi }^p$，便开始训练高级策略。

每一步根据高级策略${\pi }_{\theta }(a^h|s)$输出的center index，选择对应的goal state，fixed low-level policy将基于goal与环境进行k-steps的交互。最大化累计奖励优化高级策略参数。下：
$$J(\theta )=\mathbb{E}{\pi }_{\theta }\left[\sum _{t=0}^{\infty }{\gamma }^t\left(\sum _{i=kt}^{(k+1)t}R(s_i,a_i)-\alpha D_{\text{Kl.}}\left({\pi }_{\psi }^p(a^h|s_{kt})\Vert {\pi }_{\theta }(a^h|s_{kt})\right)\right)\right]$$