offline RL | ABM：从 offline dataset 的好 transition 提取 prior policy

ICLR 2020，6 6 6。
材料：
- 论文题目：Keep Doing What Worked: Behavior Modelling Priors for Offline Reinforcement Learning
- 项目网站：https://sites.google.com/view/behavior-modelling-priors
- pdf 版本：https://arxiv.org/pdf/2002.08396.pdf
- html 版本：https://ar5iv.labs.arxiv.org/html/2002.08396
- open review：https://openreview.net/forum?id=rke7geHtwH
总结：
- 是 2020 年的 offline RL 算法。
  - 一些 offline RL 编年史： BCQ 2019 ICLR，CQL 2020 NeurIPS，Advantage-Weighted Regression（AWR） 2020 2021 ICLR 连续被拒，TD3+BC 2021 NeurIPS，IQL 2022 ICLR，Pessimistic Bootstrapping（PBRL） 2022 ICLR。（链接是本站博客）
- 算法框架：Policy Iteration，即先搞出一个 policy，再算 policy 的 Q function，再 a = argmax Q(s,a) 搞出新 policy，再算 policy 的 Q function，...
  - ① 在 policy evaluation 时，使用 offline dataset 里的 transition 来更新。
  - ② 在 policy improvement 时，maximize Q(s, \(\pi\)(s)) ，同时约束 \(\pi\) 与一个 prior policy 的 KL 散度，这是为了不要与 in-distribution 的 state-action 离太远，离得越远，Q 的 over-estimate 越严重。
- 核心 idea 是 prior policy 如何得到。在先前的 BCQ 和 BEAR-QL 中，直接去学 offline dataset 的总的 policy，maximize Σ_{offline dataset} log(a|s)。
- 这篇文章的核心 trick 叫做 ABM（Advantage-weighted Behavior Model）。
  - 在从 offline dataset 里提取 policy 的过程中，对 log(a|s) 进行一个基于 advantage 的加权。对 (s0, a0) 的 advantage = Q - V = Return(s0,a0,s1,...,sn,V(sn)) - V(s0) 。
  - 若 advantage ≥ 0，则加权 = 1，否则加权 = 0。即，只在具有正 advantage 的环境转变的"好"数据上拟合模型。（这个加权函数还尝试了 exp 之类，但效果差异不大）。
思考：
- advantage weighted 的思想，感觉好像 performance 非常好...
- offline RL 有两种主要思想，一种是 Q update 时拉低 OOD 的 Q value（CQL PBRL），或在 policy update 添加不要离 behavior policy 太远的正则项（TD3+BC），另一种是对 offline dataset 的 policy 进行 advantage-weighted（AWR）。IQL 是一股清流，用 (s,a,r,s',a') 来做 Q update。

0 abstract

background：
- off-policy RL 适用于只有固定数据集（batch）且无法获得新 experience 的 setting，对机器人控制等现实世界的问题很有吸引力。然而在实践中，标准的 off-policy RL 在 continuous control 的 batch 设置中失败。
method：
- 在本文中，我们提出了简单解决方案：承认使用任意 behavior policy 生成的数据，并使用学习的先验（优势加权行为模型（ABM））将 RL 策略偏向于先前已执行并可能在新任务上取得成功的操作。我们的方法可以看作是最近 batch RL（offline RL）工作的扩展，可以从冲突的数据源中稳定地学习。
results：
- 在各种 RL 任务中，performance 比起强 baseline 都有所改进，包括在标准的 continuous control benchmark 以及 simulation 和真机的 multi-task learning。

1 intro

Stay close to the relevant data：① 学习一个先验，哪些候选策略可能得到数据支持（同时确保先验 focus on 相关轨迹），② 进行 policy improvement，stay close to the learned prior policy。
提出了一种 policy iteration 算法，学习先验，建立 behavior data 的优势加权模型，使 RL 策略偏向于以前经历过的、并且在当前任务中也很可能成功的 action。
我们还发现，利用适当的先验足以稳定学习；证明当使用 policy iteration 算法时， policy evaluation 步骤是 implicitly stabilized 的 ------ 只要忠实评估 TD-error 式 update 的值函数。这导致了比以前的工作更简单的算法（Fujimoto 等人，2018；Kumar 等人，2019 年）。
advantage-weighted behavior model（ABM）。

3 method

A learned prior for offline off-policy RL from imperfect data - 从不完美数据中学习 offline RL 的先验。

使用策略迭代（policy iteration）框架，在 policy improvement 步骤中有约束（参见 appendix A 的 Algorithm 1）。

首先，从迭代 i 中的给定策略 \(\pi_i\) 开始，去找一个近似的 action-value 函数 \(Q^{\pi_i}(s,a)\) 。
然后，使用 \(\hat Q^{\pi_i}\) 来优化 \(\pi_{i+1}\) ，同时约束确保接近 empirical state-action distribution。

notation：D_μ 是 offline dataset，θ 是 policy 的参数，φ 是 Q function 的参数。

3.1 Policy Evaluation

同时学 value function 和 Q function。
\(\arg\min \big[ r(s_t,a_t) + γ\hat V(s_{t+1})-\hat Q_{new}(s_t,a_t)\big]^2\) ，其中 V hat 为前一个 Q hat (s, \(\pi\)(s)) 。（公式 1）

3.2 Prior Learning and Policy Improvement

把 policy learning 当作约束优化问题：maximize Q hat(s, \(\pi\)(s))，同时约束 policy 与 prior policy 的 KL 散度＜ε。（公式 2）
考虑两种算法来得到 prior policy：
- 类似于 BCQ 和 BEAR-QL 的方法：直接 argmax Σ log(a | s) ，这里面的 s a 来自于 offline dataset。（公式 3）
- weighted advantage 方法：argmax Σ log(a | s) f(Return - V hat) 。（公式 4）
  - Return 用 \(Σγ^{N-i} r(s,a) + γ^{N-t}\hat V\) 来算，f 是一个递增的非负函数。
  - 这里采用 f(x) = 1 if x≥0 else 0。还尝试了 f(x) = exp x 之类，但实验结果差不多。
  - 只在"好"数据上拟合模型（即具有正 advantage 的环境转变，其中 advantage 是使用当前策略估计的）
（如果 offline dataset 质量足够高，可以直接在 dataset 里学 \(\pi_{abm}\) 和 Q abm ，作为输出策略：用 \(\pi_{abm}\) 计算 advantage，然后用 advantage-weight 计算新的 Q abm）
如何优化公式 2 | EM-style optimization：（两步方法）
- 首先，公式 2 的 \(\arg\max\pi\) 最优解可以写为， \(\hat \pi(a|s)\propto \pi_{prior}(a|s)\exp(\hat Q^{\pi_i}\big(s,a\big)/\eta)\) ，其中 η 可用凸优化确定。没有听懂，非常神秘......
- 然后，去用 \(\pi_{prior}\) 来查询 \(\hat Q^{\pi_i}\) 的值，等价于 maximize 一个 weighted log likelihood，用梯度下降 + 约束 KL 散度（trust-region constraint）来实现。（没懂）
如何优化公式 2 | Stochastic value gradient optimization：
- 把 KL 散度的约束用 Lagrange 松弛塞到目标函数里。
算法流程：

5 experiment

实验环境：DeepMind control suite，（multi-task setting）Mujoco 的 Sawyer robot arm（以及对应的真机）。