O2O : MOTO: Offline to Online Fine-tuning for Model-Based Reinforcement Learning

CoRL 2023 \ RRL 2023
paper

Intro

文章研究了在现实机器人任务中，从高维观测数据进行强化学习的离线预训练和在线微调问题。最近的离线无模型方法成功地使用在线微调来提高智能体在数据收集策略上的性能，或适应新任务。同时，基于模型的强化学习算法在样本效率和任务复杂性方面取得了显著进展，但在微调设置中仍未得到充分利用。文章认为，现有的基于模型的离线强化学习方法由于分布偏移、非动态数据和非平稳奖励等问题，在高维领域不适合离线到在线的微调。因此，提出了一种on-policy的基于模型的方法，该方法可以通过基于模型的价值扩展和策略正则化高效地重用先前数据，同时通过控制认知不确定性来防止模型过度利用。

Method

World Model

采用VAE based的RSSM模型建模

通过联合优化ELBO对模型训练
$${\mathcal{L}}_{p_{\theta },q_{\theta }}^{\mathrm{model}}={\mathbb{E}}_{\tau \sim \mathcal{D}}\left[\sum _t-\ln p_{\theta }(x_t\mid s_t)-\ln p_{\theta }(r_t\mid s_t)+{\mathbb{D}}_{KL}[q_{\theta }(s_t|x_t,s_{t-1},a_{t-1})||p_{\theta }^{i_t}(s_t|s_{t-1},a_{t-1})]\right].$$
其中latent state采用离散形式。对动力学模型采用集成模型 { p θ i ( s t + 1 ∣ z t ) } i = 1 M \{p_\theta^i(s_{t+1}|z_t)\}_{i=1}^M {pθi​(st+1​∣zt​)}i=1M​，执行时随机采样一个模型去优化世界模型。同时，为了防止模型被过度开发使用，采用基于认知不确定的奖励正则化
u θ ( s t , a t ) = s t d ( { l θ i ( z t + 1 ) } i = 1 M ) , r ^ θ ( s t , a t , s t + 1 ) = r θ ( s t + 1 ) − α u θ ( s t , a t ) u_\theta(s_t,a_t)=\mathrm{std}(\{l_{\theta^i}(z_{t+1})\}_{i=1}^M),\\ \widehat{r}_\theta(s_t,a_t,s_{t+1})=r_\theta(s_{t+1})-\alpha u_\theta(s_t,a_t) uθ​(st​,at​)=std({lθi​(zt+1​)}i=1M​),r θ​(st​,at​,st+1​)=rθ​(st+1​)−αuθ​(st​,at​)

Variational Model-Based Value Expansion

完成世界模型训练后，便是利用他进行价值函数的优化。这里采用的时类似$TD(\lambda )$的MC估计方法。每一次训练时，采样真实状态为初始状态，利用model进行rollout。
$${\hat{a}}_j^t\sim {\pi }_{\psi }(a|{\hat{s}}_j^{t-1}),{\hat{s}}_j^{t+1}\sim p_{\theta }(s|{\hat{a}}_j^t,{\hat{s}}_j^t),{\hat{r}}_j^t\sim p_{\theta }(r|{\hat{s}}_j^t),$$
然后按照一般的off-policy的算法，采用两个Q函数以及target Q函数，对价值函数建模
V 0 π ψ ( s ^ j t ) = min ⁡ { Q ψ 1 ( s ^ j t , a ^ j t ) , Q ψ 2 ( s ^ j t , a ^ j t ) } , V K π ψ ( s ^ j t ) = ∑ k = 1 K γ k − 1 r ^ j k + t + γ K V 0 π ψ ( s ^ j t + K ) V_{0}^{\pi_{\psi}}(\hat{s}_{j}^{t})=\min\{Q_{\psi^{1}}(\hat{s}_{j}^{t},\hat{a}_{j}^{t}),Q_{\psi^{2}}(\hat{s}_{j}^{t},\hat{a}_{j}^{t})\},\quad V_{K}^{\pi_{\psi}}(\hat{s}_{j}^{t})=\sum_{k=1}^{K}\gamma^{k-1}\hat{r}_{j}^{k+t}+\gamma^{K}V_{0}^{\pi_{\psi}}(\hat{s}_{j}^{t+K}) V0πψ​​(s^jt​)=min{Qψ1​(s^jt​,a^jt​),Qψ2​(s^jt​,a^jt​)},VKπψ​​(s^jt​)=k=1∑K​γk−1r^jk+t​+γKV0πψ​​(s^jt+K​)
然后通过GAE对价值函数估计
$$V^{{\pi }_{\psi }}({\hat{s}}_j^t)=(1-\lambda )\sum _{k=1}^{H-t-1}{\lambda }^{k-1}{V}_k^{{\pi }_{\psi }}({\hat{s}}_j^t)+{\lambda }^{H-t-1}{V}_{H-t}^{{\pi }_{\psi }}({\hat{s}}_j^t)$$
设${\hat{V}}^{{\pi }_{\psi }}(s):=\lambda V^{{\pi }_{\psi }}(s)+(1-\lambda )V_0^{{\pi }_{\psi }}$，策略优化目标为
$${\mathcal{L}}_{{\pi }_{\psi }}^{\mathrm{model}}=-\frac{1}{HT}{\mathbb{E}}_{\tau \sim \mathcal{D}{\pi }_{\psi },p_{\theta }}\left[\sum _{t=0,j=1}^{H-1,T}{\hat{V}}^{{\pi }_{\psi }}({\hat{s}}_j^t)\right]$$
对Critic的训练由两个均方误差构成
$$\begin{gathered} {\mathcal{L}}_{Q_{{\psi }^i}}^{\mathrm{model}}=\frac{1}{HT}\underset{\tau \sim \mathcal{D}}{E}\underset{{\pi }_{\psi },p_{\theta }}{E}\left[\sum _{t=0,j=1}^{H-1,T}({\bar{V}}^{{\pi }_{\psi }}({\hat{s}}_j^t)-Q_{{\psi }^i}({\hat{s}}_j^t,{\hat{a}}_j^t){)}^2\right] \\ {\mathcal{L}}_{Q_{{\psi }^i}}^{\mathrm{data}}=\frac{1}{T-1}\underset{\tau \sim \mathcal{D}}{E}\underset{{\pi }_{\psi }}{E}\left[\sum _{j=1}^{T-1}{\left(r_{j+1}^0+\gamma {\hat{V}}^{{\pi }_{\psi }}(s_{j+1}^0)-Q_{{\psi }^i}(s_j^0,a_j^0)\right)}^2\right] \end{gathered}$$
其中${\bar{V}}^{{\pi }_{\psi }}({\hat{s}}_j^t)$为target Q所求得。第一个损失数据为模型rollout数据，第二项针对真实环境数据。

Behaviour Prior Policy Regularization

现实的机器人学习数据集通常由狭隘的数据组成。因此，在训练的初始阶段，动力学模型可能非常不准确，代理策略可以从更强的数据正则化中受益。为了避免建模行为分布的额外复杂性，遵循类似于47的方法，该方法部署了$f$形式的正则化项, 其中f可以采用阈值函数,例如对优势函数为正的样本执行行为克隆。
$${\mathcal{L}}_{{\pi }_{\psi }}^{\mathrm{reg}}=-\underset{\tau \sim \mathcal{D}}{E}\left[\sum _{t=1}^T\log {\pi }_{\psi }({\mathbf{a}}_t\mid s_t)f\left(\underset{\text{Advantage over trajectory snippet }{\mathbf{s}}_t:{\mathbf{s}}_{t+H}}{\underbrace{{\gamma }^H{V}^{{\pi }_{\psi }}(s_{t+H})+\sum _{j=1}^H{\gamma }^j{r}_{t+j}-V^{{\pi }_{\psi }}(s_t)}}\right)\right]$$

在这项工作中，算法专注于具有稀疏奖励的真实机器人操作任务，并简单地根据轨迹是否在环境中实现目标来作为$f$。因此，加入正则后得策略优化目标为；
$${\mathcal{L}}_{{\pi }_{\psi }}={\mathcal{L}}_{{\pi }_{\psi }}^{\mathrm{model}}+\beta {\mathcal{L}}_{{\pi }_{\psi }}^{\mathrm{reg}}$$

伪代码

results

实验环境为Meta World以及Kitchen