offline RL | Pessimistic Bootstrapping (PBRL):在 Q 更新中惩罚 uncertainty,拉低 OOD Q value

0 abstract

Offline Reinforcement Learning (RL) aims to learn policies from previously collected datasets without exploring the environment. Directly applying off-policy algorithms to offline RL usually fails due to the extrapolation error caused by the out-of-distribution (OOD) actions. Previous methods tackle such problem by penalizing the Q-values of OOD actions or constraining the trained policy to be close to the behavior policy. Nevertheless, such methods typically prevent the generalization of value functions beyond the offline data and also lack precise characterization of OOD data. In this paper, we propose Pessimistic Bootstrapping for offline RL (PBRL), a purely uncertainty-driven offline algorithm without explicit policy constraints. Specifically, PBRL conducts uncertainty quantification via the disagreement of bootstrapped Q-functions, and performs pessimistic updates by penalizing the value function based on the estimated uncertainty. To tackle the extrapolating error, we further propose a novel OOD sampling method. We show that such OOD sampling and pessimistic bootstrapping yields provable uncertainty quantifier in linear MDPs, thus providing the theoretical underpinning for PBRL. Extensive experiments on D4RL benchmark show that PBRL has better performance compared to the state-of-the-art algorithms.

  • background:
    • offline RL 从之前收集的 dataset 中学习策略,而无需探索环境。由于 OOD actions 导致的 extrapolation error,将 off-policy RL 直接应用于 offline RL 通常会失败。
    • 先前工作通过 penalize OOD action 的 Q value,或去约束 trained policy 接近 behavior policy 来解决此类问题。
    • 然而,这些方法通常阻止了 value function generalize 到 offline dataset 之外,并且也缺乏对 OOD data 的精确表征(characterization)。
  • method:
    • 我们提出了 offline RL 的悲观引导(Pessimistic Bootstrapping,PBRL),它是一个纯粹的 uncertainty-driven 的 offline 算法,没有明确的 policy constraint。
    • 具体的,PBRL 通过 bootstrapped Q functions 的 disagreement 进行 uncertainty 的量化,并根据所估计的 uncertainty,对 value function 进行惩罚,从而实施 pessimistic updates。
    • 对于 extrapolation error 的处理,我们进一步提出了一种新的 OOD sampling 方法。
    • 理论:上述 OOD sampling + pessimistic bootstrapping,在 linear MDP 中形成了一个 uncertainty 的量化器,是可以证明的。
  • 实验:
    • 在 D4RL 基准测试上的大量实验表明,与最先进的算法相比,PBRL 具有更好的性能。

3 method

3.1 使用 bootstrapped-Q function 进行 uncertainty 的量化

  • 维护 K 个各自 bootstrap 更新的 Q-function。
  • uncertainty \(U(s,a)=\mathrm{std}(Q^k(s,a))=\sqrt{\frac1K\sum(Q^k-\bar Q)^2}\) 。(看 figure 1(a),感觉定义是有道理的)

3.2 pessimistic learning - 悲观学习

  • idea:基于 uncertainty 来惩罚 Q function。

  • PBRL 的 loss function 由两部分组成:① ID 数据的 TD-error、② OOD 数据的伪 TD-error。

  • ① ID 数据的 TD-error,见公式 (4),大概就是 \(\hat T^{in}Q^k(s,a):=r+\gamma \hat E\big[Q^k(s',a')-\beta_{in}U(s',a')\big]\) ,对所转移去的 (s',a') 的 uncertainty 进行惩罚。

    • (上文的 ID (s, a, r, s', a') 由 offline dataset 得到)
  • ② OOD 数据的伪 TD-error,s' 好像是 ID 的 state,a' 是 policy 生成的(可能是 OOD 的)action。

    • 惩罚方式的 idea: \(\hat T^{ood}Q^k(s^{ood},a^{ood}):=Q^k(s^{ood},a^{ood})-\beta_{ood}U(s^{ood},a^{ood})\) ,直接减去它的 uncertainty。
    • (如果 (s,a) 是 ID state-action,那么 uncertainty 会很小)
    • 相关的实现细节:早期 Q function 的截断 \(\max[0, \hat T^{ood}Q^k(s,a)]\) ,在训练初期使用大的 β ood 实现对 OOD action 的强惩罚,在训练过程中不断减小 β ood 的值。
    • (感觉也算是使用 sarsa 式更新...)
  • loss function:

\[L_{critic}=\hat E_{(s,a,r,s')\sim D_{in}}\bigg[(\hat T^{in}Q^k-Q^k)^2\bigg] + \hat E_{s^{ood}\sim D_{in},~a^{ood}\sim\pi(s^{ood})}\bigg[(\hat T^{ood}Q^k-Q^k)^2\bigg] \]

  • policy: policy 希望最大化 Q function,具体的,最大化 ensemble Q 中的最小值。

3.3 是理论。