O2O : Finetuning Offline World Models in the Real World

CoRL 2023 Oral
paper
code

Intro

算法基于TD-MPC，利用离线数据训练世界模型，然后在线融合基于集成Q的不确定性估计实现Planning。得到的在线数据将联合离线数据共同训练目标策略。

Method

TD-MPC

TD-MPC由五部分构成:

1. 状态特征提取$z=h_{\theta }(s)$
2. 隐动力学模型$z^{\prime }‘=d_{\theta }(z,a)$
3. 奖励模型$\hat{r}=R_{\theta }(z,a)$
4. planning policy $\hat{a}={\pi }_{\theta }(z)$
5. 终止状态下的$\hat{q}=Q_{\theta }(z,a)$

通过联合训练进行优化，损失函数为：
$$\mathcal{L}(\theta )={\mathbb{E}}_{(\mathbf{s},\mathbf{a},r,{\mathbf{s}}^{\prime }{)}_{0:h}\sim \mathcal{B}}\lfloor \sum _{t=0}^h\left(\underset{\text{Latent dynamics}}{\underbrace{\Vert {\mathbf{z}}_t^{\prime }-\mathrm{sg}(h_{\varphi }({\mathbf{s}}_t^{\prime })){\Vert }_2^2}}+\underset{\text{Reward}}{\underbrace{\Vert {\hat{r}}_t-r_t{\Vert }_2^2}}+\underset{\text{Value}}{\underbrace{\Vert {\hat{q}}_t-q_t{\Vert }_2^2}}-\underset{\text{Action}}{\underbrace{Q_{\theta }({\mathbf{z}}_t,{\hat{\mathbf{a}}}_t)}}\right)\rfloor (1)$$
在Offline 设定下，分布偏移将导致Q估计以及隐模型以及价值函数的错误估计。启发于IQL，通过只对in-sample的动作尽心TD-backups来估计，缓解过估计问题。因此对模型价值函数利用离线数据进行训练时，此时Q函数采用IQL中的期望回归方法优化
L V ( θ ) = ∣ τ − 1 { Q ϕ ( z t , a t ) − V θ ( z t ) < 0 } ∣ ( Q ϕ ( z t , a t ) − V θ ( z t ) ) 2 , \mathcal{L}_{V}(\theta)=|\tau-1_{\{Q_{\phi}(\mathbf{z}_{t},\mathbf{a}_{t})-V_{\theta}(\mathbf{z}_{t})<0\}}|(Q_{\phi}(\mathbf{z}_{t},\mathbf{a}_{t})-V_{\theta}(\mathbf{z}_{t}))^{2}, LV​(θ)=∣τ−1{Qϕ​(zt​,at​)−Vθ​(zt​)<0}​∣(Qϕ​(zt​,at​)−Vθ​(zt​))2,
同时对planning policy采用AWR的更新，即$\exp (\beta (Q_{\varphi }({\mathbf{z}}_t,{\mathbf{a}}_t)-V_{\theta }(\widehat{{\mathbf{z}}_t})))\log {\pi }_{\theta }({\mathbf{a}}_t|{\mathbf{z}}_t)$

Uncertainty Estimation as Test-Time Behavior Regularizatio

离线训练的模型依旧存在OOD数据过估计，需要在线微调。文章提出基于不确定性估计的planning实现在线交互过程中的动作选择。planning一定程度缓解基于约束的离线算法导致的在现阶段探索能力不足。进而导致算法样本效率低的问题。

首先构建集成Q函数模型，计算基于标准差的不确信度，作为惩罚项对奖励进行调整，实现保守的在线planning。
R ^ = γ h ( Q θ ( z h , a h ) − λ u h ) + ∑ t = 0 h − 1 γ t ( R θ ( z t , a t ) − λ u t ) , u t = s t d ( { Q θ ( i ) ( z t , a t ) } i = 1 N ) \hat{\mathcal{R}}=\gamma^{h}\left(Q_{\theta}(\mathbf{z}_{h},\mathbf{a}_{h})-\lambda u_{h}\right)+\sum_{t=0}^{h-1}\gamma^{t}\left(R_{\theta}(\mathbf{z}_{t},\mathbf{a}_{t})-\lambda u_{t}\right),\quad u_{t}=\mathrm{std}\left(\{Q_{\theta}^{(i)}(\mathbf{z}_{t},\mathbf{a}_{t})\}_{i=1}^{N}\right) R^=γh(Qθ​(zh​,ah​)−λuh​)+t=0∑h−1​γt(Rθ​(zt​,at​)−λut​),ut​=std({Qθ(i)​(zt​,at​)}i=1N​)

除此外，还维护两个buffer分别存储离线数据于在线数据，通过balance sampling数据训练模型、策略以及价值函数。

结果