Offline: Look Beneath the Surface: Exploiting Fundamental Symmetry for Sample-Efficient Offline RL

ICML 2023 Workshop SCIS Poster
paper
code

Intro

研究在有限的且覆盖范围狭窄的数据集中学习到一个可靠的策略。具体来说，通过结合时间反演对称性（T-对称性），在隐空间中对前向和反向的动力学过程的一致性建立动力学模型（TDM）。文章发现TDM为小样本数据集提供了良好可靠的的表示空间，并可以基于T-对称性为衡量指标为OOD样本进行度量。基于此构造了一个新的离线RL算法（TSRL），该算法具有较少保守的策略约束和一个可靠数据增强方法。在D4RL数据上即使只有原始样本的1%，该算法的性能显著优于现有的离线RL算法

Metho

时间反演对称

文章指出大多数的现实环境模型中，可以用非线性一阶常微分方程近似动力学模型，若存在可逆变换$\Gamma :\Omega \mapsto \Omega$使得方程$d\Gamma (\mathbf{x})/dt=-F(\Gamma (\mathbf{x}))$，则称为ODE系统是时间反转对称。定义时间演化算子$U_{\Delta t}:\mathbf{x}(t)\mapsto U_{\Delta t}(\mathbf{x}(t))=\mathbf{x}(t+\Delta t)$。T-对称表明$\Gamma \circ U_{\tau }=U_{-\tau }\circ \Gamma$（复合映射），即任意状态的前向时间演化的反转应该等于反向状态的后向时间演化。

在时间离散的MDP中，$x=(s,a)$,方程为$\dot{s}=\frac{ds}{dt}=s^{\prime }-s$。此时前向模型为$F(s,a)=\dot{s}$而反向模型为$-\tilde{G}(s^{\prime },a^{\prime }).$二者相等。为了避免下一时刻状态产生不可逆的动作或者特殊的动力学过程，设置反向模型中的动作为$a$.

TDM

TDM为encoder-decoder的架构，其中解码器额外添加${\delta }_s=0$是指引其解码结果目标为状态值，$\delta =1$为解码目标为$\dot{s}$。encoder-decoder的损失函数由两个MSE构成
$${\ell }_{rec}(s,a)=\Vert s-{\psi }_s(z_s,0){\Vert }_2^2+\Vert a-{\psi }_a(z_a){\Vert }_2^2$$

隐变量前向模型$f(z_s,z_a)={\dot{z}}_s$中，目标是让${\dot{z}}_s=z_{s^{\prime }}-z_s{\dot{z}}_s=z_{s^{\prime }}-z_s$。而链式法则得到${\dot{z}}_s=\frac{d{z}_s}{dt}=\frac{\partial z_s}{\partial s}\cdot \frac{ds}{dt}={\nabla }_s{z}_s\cdot \dot{s}.$因此，构造优化函数：
$${\ell }_{fwd}(s,a,s^{\prime })=\Vert ({\nabla }_s{z}_s)\dot{s}-{\dot{z}}_s{\Vert }_2^2=\Vert \frac{\partial \varphi (s,a)}{\partial s}\dot{s}-f(\varphi (s,a)){\Vert }_2^2$$
同时要求解码器解码器从${\dot{z}}_s$中解码得到$\dot{s}$，构造损失函数
$${\ell }_{ds}(s,a,s^{\prime })=\Vert \dot{s}-{\psi }_s({\dot{z}}_s,1){\Vert }_2^2=\Vert \dot{s}-{\psi }_s(f(\varphi (s,a)),1){\Vert }_2^2$$

同理。反方向动力学模型的损失函数为
$${\ell }_{rvs}(s,a,s^{\prime })=\Vert ({\nabla }_{s^{\prime }}{z}_{s^{\prime }})(-\dot{s})-(-{\dot{z}}_s){\Vert }_2^2=\Vert \frac{\partial \varphi (s^{\prime },a)}{\partial s^{\prime }}(-\dot{s})-g(\varphi (s^{\prime },a)){\Vert }_2^2$$

上述两个模型满足T对称$g(z_s+f(z_s,z_a),z_a)=-{\dot{z}}_s=-f(z_s,z_a)$，因此最小化损失函数
$${\ell }_{T-sym}(z_s,z_a)={\Vert f(z_s,z_a)+g(z_s+f(z_s,z_a),z_a)\Vert }_2^2$$

综上所述，TDM的联合优化目标为
$${\mathcal{L}}_{TDM}=\sum _{(s,a,s^{\prime })\in \mathcal{D}}[{\ell }_{rec}+{\ell }_{ds}+{\ell }_{fwd}+{\ell }_{rvs}+{\ell }_{T-sym}](s,a,s^{\prime })+{\lambda }_{L1}[{\mathcal{L}}_1(f)+{\mathcal{L}}_1(g)]$$
其中${\mathcal{L}}_1$为L1-正则化。

TSRL

完成TDM的优化后便是将其用于离线学习。其中价值函数通过最小化基于TDM表征的均方bellman误差
$$Q=\underset{Q}{\mathrm{argmin}}{\mathbb{E}}_{(s,a,s^{\prime })\sim \mathcal{D}}\left[\right(r(s,a)+\gamma \hat{Q}(\varphi (s^{\prime },\pi (\cdot |s^{\prime })))-Q(\varphi (s,a)){)}^2]$$

对于策略则是采用如下优化
$${\mathrm{argmax}}_{\pi }{\mathbb{E}}_{(s,a)\sim \mathcal{D}}[\alpha Q(\varphi (s,\pi (\cdot |s)))-{\lambda }_1\Vert z_{a^{\pi }}-z_a{\Vert }_2^2-{\lambda }_2{\ell }_{T-sym}(\varphi (s,\pi (\cdot |s)))]$$
其中添加了两种正则化项。第一项是限制隐变量偏差，第二项则是使得样本符合T-对称，这使得即使面对OOD的数据，如果与数据集中的一些潜在动作具有相似的隐空间特征表示，并且与满足TDM 中的 T 对称属性一致，那么可以期望这些操作之间存在一定程度的等价性。通过在可靠的 OOD 区域上实现策略学习和泛化，这导致了更宽松的策略约束，这对于小数据集设置至关重要。

数据增强

对隐空间变量进行数据增强即添加扰动项$(z_s,z_a)\mapsto (z_s+ϵ,z_a),$。前向模型${\dot{z}}_s=f(z_s,z_a)$则满足$z_{s^{\prime }}+{ϵ}^{\prime }=z_s+ϵ+f(z_s+ϵ,z_a).$。反向模型中则是$-{\dot{z}}_s=g(z_{s^{\prime }},z_a)\mathrm{as}:z_s+{ϵ}^{\prime \prime }=z_{s^{\prime }}+{ϵ}^{\prime }+g(z_{s^{\prime }}+{ϵ}^{\prime },z_a)$。为了满足T-对称则需要使得${ϵ}^{\prime \prime }=ϵ$
$${ϵ}^{\prime \prime }-ϵ=f(z_s+ϵ,z_a)+g(z_s+ϵ+f(z_s+ϵ,z_a),z_a)=0$$
该式等价于原始${\ell }_{T-sym}$输入项添加扰动即${\ell }_{T-sym}(z_s+ϵ,z_a)=0$.。在实际操作中，设置${\ell }_{T-sym}(z_s+ϵ,z_a)<h$，该值为所有(s,a)的${\ell }_{T-sym}(\varphi (s,a))$的分位数值(实际取50%或者70%)

伪代码

Results

能够看出技术在小数据集上，算法也有不错的效果。