Offline RL :When Demonstrations meet Generative World Models: A Maximum Likelihood Framework...

NIPS 2023 oral
Keywords: Model-based Offline IRL
paper

1 Introduction

离线IRL旨在利用离线专家数据恢复奖励函数以及环境动态结构。由于分布偏移问题，从固定数据集恢复的环境模型不准确，容易导致奖励函数估计存在误差。本文采用双层优化(Bi-level optimization)对二者进行交替优化。基于下层专家策略实现上层保守模型似然最大化。由图1可以看出方法分为两阶段：利用离线数据估计世界模型，并采用不确定估计技术不确定性度量，对高模型不确定性和低覆盖率的(s,a)进行惩罚，避免在不熟悉区域探索。第二阶段则是恢复奖励函数.

2 Preliminaries

给定数据集$D:=(s,a,s^{\prime })$，训练世界模型$\hat{P}(s^{\prime }|s,a)$，基于此，model-based IRL的优化问题为：
$$\begin{array}{rl}\underset{\theta }{\max }L(\theta )& :={\mathbb{E}}_{{\tau }^{\mathrm{E}}\sim (\eta ,{\pi }^{\mathrm{E}},P)}\left[\sum _{t=0}^{\infty }{\gamma }^t\log {\pi }_{\theta }(a_t|s_t)\right]\\ s.t.{\pi }_{\theta }& :=\arg \underset{\pi }{\max }{\mathbb{E}}_{{\tau }^{\mathrm{A}}\sim (\eta ,\pi ,\hat{P})}\left[\sum _{t=0}^{\infty }{\gamma }^t\right(r(s_t,a_t;\theta )+U(s_t,a_t)+\mathcal{H}(\pi (\cdot |s_t)\left)\right)]\end{array}$$
其中$\mathcal{H}(\pi (\cdot |s)):={\sum }_{a\in \mathcal{A}}-\pi (a|s)\log \pi (a|s)$；U代表对世界模型$\hat{P}$不确定性度量的惩罚项；

上述问题优化为Bi-level问题。s.t.代表的low-level问题固定奖励函数优化策略；high-level则是固定策略优化奖励函数，使得最优策略${\pi }_{\theta }$在专家数据的对数似然最大。

其次，将对数似然作为优化目标是合理的，因为它在对世界有限的认知范围下（$\hat{P}$由D得到）搜索最有奖励函数，来解释专家行为。

low-level问题下得到的策略是保守的，因为包含代表不确定性惩罚以及正则化的U。确保最优策略不会再不熟悉区域探索。

3 Method

基于奖励函数以及世界模型，构造Q以及V
$$\begin{array}{rl}& Q_k(s,a):=r(s,a;{\theta }_k)+U(s,a)+\gamma {\mathbb{E}}_{s^{\prime }\sim \hat{P}(\cdot |s,a)}[V_k(s^{\prime })]\\ & V_k(s):={\mathbb{E}}_{\tau \sim (\eta ,{\pi }_k,\hat{P})}[\sum _{t=0}^{\infty }{\gamma }^t(r(s_t,a_t;{\theta }_k)+U(s_t,a_t)+\mathcal{H}({\pi }_k(\cdot |s_t))\left)\right|s_0=s]\end{array}$$
类似于SAC构造基于熵的策略形式：
$${\pi }_{\theta }(a|s)=\frac{\exp Q_{\theta }(s,a)}{{\sum }_{\tilde{a}\in \mathcal{A}}\exp Q_{\theta }(s,\tilde{a})},V_{\theta }(s)=\log (\sum _{a\in A}\exp Q_{\theta }(s,a))$$

将该策略带入到原问题，可以将原问题分解为$\hat{L}$与动力学模型误差项。文章接下来假设奖励函数以及模型惩罚项U存在bound，那么$\hat{L}$与$L$之间的差距是被专家状态动作分布下的模型误差bound

上述引理说明随着在线收集数据逐渐覆盖专家(s,a)的分布，那么$\hat{L}$的最优解可以近似原问题的解。

待优化问题为：
$$\underset{\theta }{\max }\hat{L}(\theta ),\text{s.t.}{\pi }_{\theta }:=\arg \underset{\pi }{\max }{\mathbb{E}}_{{\tau }^{\Lambda }\sim (\eta ,\pi ,{\hat{I}}^{\lambda })}\left[\sum _{\ell =0}^{\infty }{\gamma }^{\ell }\right(r(s_{\ell },a_{\ell };\theta )+U(s_{\ell },a_{\ell })+\mathcal{H}(\pi (\cdot |s_{\ell })\left)\right)]$$
其中$\hat{L}(\theta ):={\mathbb{E}}_{{\tau }^{\mathrm{E}}\sim (\eta ,{\pi }^{\mathrm{E}},P)}\left[{\sum }_{t=0}^{\infty }{\gamma }^t\right(r(s_t,a_t;\theta )+U(s_t,a_t)\left)\right]-{\mathbb{E}}_{s_0\sim \eta (\cdot )}[V_{\theta }(s_0)]$
上层问题优化奖励函数，下层问题优化策略函数。双层优化分两步：1）保守策略改进2）奖励函数更新

3.1 policy improvement

为了执行策略改进步骤，首先使用${\hat{Q}}_k$来近似软q函数，该估计满足以下条件
$$\Vert \widehat{Q_k}-Q_k{\Vert }_{\infty }:=\underset{s\in \mathcal{S},a\in \mathcal{A}}{\max }|\widehat{Q_k}(s,a)-Q_k(s,a)|\le {ϵ}_{\mathrm{app}}.$$
得到近似的${\hat{Q}}_k,$ 利用soft policy iteration得到策略表示${\pi }_{k+1}(a|s)\propto \exp \left({\hat{Q}}_k(s,a)\right),\forall s\in \mathcal{S},a\in \mathcal{A}.$

在实践中，可以采用soft Q-learning和SAC，以获得精确逼近的软q函数，从而实现策略迭代的稳定更新。

3.2 Reward optimization

第k次迭代，根据奖励${\theta }_k$、改进后策略${\pi }_{k+1}$，优化得到改进后的奖励${\theta }_{k+1}$。原问题为最大化$\hat{L}(\theta )$，因此可采用梯度上升更新。首先计算目标问题梯度$\nabla \hat{L}({\theta }_k)$。

而原high-level问题梯度可以表示为（证明过程见原文的Appendix F）
$$\nabla \hat{L}(\theta )={\mathbb{E}}_{{\tau }^{\mathrm{E}}\sim (\eta ,{\pi }^{\mathrm{E}},P)}[\sum _{t=0}^{\infty }{\gamma }^t{\nabla }_{\theta }r(s_t,a_t;\theta )]-{\mathbb{E}}_{{\tau }^{\mathrm{A}}\sim (\eta ,{\pi }_{\theta },\hat{P})}[\sum _{t=0}^{\infty }{\gamma }^t{\nabla }_{\theta }r(s_t,a_t;\theta )]$$
在实践中，由于在每个奖励函数下反复求解最优策略${\pi }_{\theta }$在计算量上太大，因此采用基于${\pi }_{k+1}$构造精确梯度估计量。特别的，采用两部进行数据收集：1）选择一个专家轨迹${\tau }_k^E$ 2)基于世界模型以及${\pi }_{k+1}$得到轨迹${\tau }_k^A$。$\nabla \hat{L}({\theta }_k)$可以进一步简写为：$$g_k:=h({\theta }_k;{\tau }_k^{\mathrm{E}})-h({\theta }_k;{\tau }_k^{\mathrm{A}})$$

其中$h(\theta ;\tau ):={\sum }_{t=0}^{\infty }{\gamma }^t{\nabla }_{\theta }r(s_t,a_t;\theta )$。进一步，通过梯度上升更新奖励函数
$${\theta }_{k+1}={\theta }_k+\alpha g_k$$