Unsupervised RL:METRA: Scalable Unsupervised RL with Metric-Aware Abstraction

ICLR 2024 Oral
paper

Intro

无监督RL旨在发现潜在的行为帮助提高下游任务效率以往方法集中于探索以及基于互信息的技能发现(skill)。然而去前者在高危复杂空间实现困难，后者也容易因为缺乏激励导致探索能力不足。本文提出METRA核心观点认为与其在复杂状态空间处理，不如构造一个更紧凑的隐空间z，类似于PCA将复杂状态空间简化:$\varphi :\mathcal{S}\to \mathcal{Z}$，而z可以通过时间距离（temporal distances）度量metric d链接到状态空间。

采用时间距离度量(状态间转换的最小交互步数)的原因是其对状态表征是不变的，因此该度量适用于pixel-based的任务。因此，通过最大化Z空间中的覆盖，可以获得近似覆盖整个状态空间的各种行为，便可实现扩展到高维、复杂的环境。

Method

METRA的目标是基于Z空间实现状态空间的最大化覆盖。类似于Skill-based的无监督方法，提出如下优化目标
$$I_{\mathcal{W}}(S;Z)=\mathcal{W}(p(s,z),p(s)p(z)),$$
上式计算了状态与skills之间的Wasserstein dependency measure (WDM)。为了让上述目标简介且易于计算，采用Kantorovich-Rubenstein duality提供一种可计算的方式最大化WDM。
$$\begin{array}{r}{I}_{\mathcal{W}}(S;Z)=\underset{\Vert f{\Vert }_L\le 1}{\sup }{\mathbb{E}}_{p(s,z)}[f(s,z)]-{\mathbb{E}}_{p(s)p(z)}[f(s,z)]\end{array}v$$
直观来说更希望采样联合分布的(s,z)而非边际分布的(s,z)。上述目标是可计算的，可以设置f为1阶李普希兹连续的函数，通过梯度下降优化，而z-condition的策略$\pi (a|s,z)$可以通过RL算法优化，其奖励函数表示为：$r(s,z)=f(s,z)-N^{-1}{\sum }_{i=1}^{N}f(s,z_i)$。而对每一个状态都需要从z的先验分布中采集N个$z_i$，这样计算量增加。

因此本文简化计算量提出参数化$f(s,a)=\varphi (s{)}^{\top }\psi (z)$。上述优化目标转化为：
$$I_{\mathcal{W}}(S;Z)\approx \underset{\Vert \varphi {\Vert }_L\le 1,\Vert \psi {\Vert }_L\le 1}{\sup }{\mathbb{E}}_{p(s,z)}[\varphi (s{)}^{\top }\psi (z)]-{\mathbb{E}}_{p(s)}[\varphi (s){]}^{\top }{\mathbb{E}}_{p(z)}[\psi (z)].$$
原文给出证明，在隐空间维度D趋于正无穷时$f(s,a)$与$\varphi (s{)}^{\top }\psi (z)$等价。然后考虑状态为最后T时刻的设定
$$\begin{array}{rl}{I}_{\mathcal{W}}(S_T;Z)& \approx \underset{\Vert \varphi {\Vert }_L\le 1,\Vert \psi {\Vert }_L\le 1}{\sup }{\mathbb{E}}_{p(\tau ,z)}[\varphi (s_T{)}^{\top }\psi (z)]-{\mathbb{E}}_{p(\tau )}[\varphi (s_T){]}^{\top }{\mathbb{E}}_{p(z)}[\psi (z)]\\ & =\underset{\varphi ,\psi }{\sup }\sum _{t=0}^{T-1}\left({\mathbb{E}}_{p(\tau ,z)}\right[(\varphi (s_{t+1})-\varphi (s_t){)}^{\top }\psi (z)]-{\mathbb{E}}_{p(\tau )}[\varphi (s_{t+1})-\varphi (s_t){]}^{\top }{\mathbb{E}}_{p(z)}[\psi (z)]),\end{array}$$
其中$p(s_0)$与$p(z)$独立分布，设$\psi (z)=z$, 约简优化目标为
$$I_{\mathcal{W}}(S_T;Z)\approx \underset{\Vert \varphi {\Vert }_L\le 1}{\sup }{\mathbb{E}}_{p(\tau ,z)}\left[\sum _{t=0}^{T-1}(\varphi (s_{t+1})-\varphi (s_t){)}^{\top }(z-\bar{z})\right]$$
且其中$\bar{z}={\mathbb{E}}_{p(z)}[z]$，若是z均值为0那该问题可看作奖励函数$\begin{array}{r}r(s,z,s^{\prime })=(\varphi (s^{\prime })-\varphi (s){)}^{\top }z\end{array}$的RL问题，联合优化$\varphi$与$\pi (a|s,z)$

METRA

到目前为止，上述问题没有指定距离函数$d$, 本文提出两个状态时间距离$d_{\mathrm{temp}}(s_1,s_2)$, 即从$s_1$到$s_2$最小环境交互步数。那么优化问题为
$$\begin{array}{r}\underset{\pi ,\varphi }{\sup }{\mathbb{E}}_{p(\tau ,z)}\left[\sum _{t=0}^{T-1}(\varphi (s_{t+1})-\varphi (s_t){)}^{\top }z\right]\\ s.t.\Vert \varphi (s)-\varphi (s^{\prime }){\Vert }_2\le 1,\forall (s,s^{\prime })\in {\mathcal{S}}_{\text{adj}},\end{array}$$
其中${\mathcal{S}}_{\text{adj}}$为相邻状态对的集合，因此d=1。直观来说，上述目标迫使策略$\pi (a|s,z)$尽可能向着由z指定的方向进行探索，但是由于$\Vert \varphi (s_1)-\varphi (s_2){\Vert }_2$存在一个upper bound。这样潜在空间应该将其(有限的)维度分配给原始状态空间中流形最大限度“展开”的状态。从某种意义上说，状态集中的最短路径应该尽可能长

Zero-shot goal-reaching with METRA

得益于$\varphi (s)$考虑时间距离上的状态抽象，可以方便的实现Zero-shot goal-reaching 。只需将方向设置$z=(\varphi (g)-\varphi (s))/\Vert \varphi (g)-\varphi (s){\Vert }_2$（连续技能）或者$z=\arg {\max }_{\dim }\left(\varphi (g)-\varphi (s)\right)$(离散技能)