Offline RL ： Generalized Decision Transformer for Offline Hindsight Information Matching

ICLR 2022 Spotlight
paper
code

Intro

基于序列模型架构的离线强化学习引入累计奖励丰富输入信号(HER 以及DT中的“returns-to-go：)，从而帮助序列模型的样本效率提升。本文证明了所有这些方法都是事后信息匹配 (hindsight information matching，HIM)—通过训练策略，输出与未来状态信息的在一些统计数据相匹配的其他轨迹。因此，本文提出广义DT算法(GDT)解决离线任务中的HIM问题。同时基于GDT推出提出Categorical DT (CDT)解决多任务state-marginal matching (SMM)问题以及提出Bi-directional DT (BDT) 解决multi-task imitation learning问题。

Method

HIM

对于轨迹数据的信息统计表示为$I(\tau )$，此外定义对轨迹数据中(s,a)通过函数$\Phi$进行特征提取后的得到 τ t Φ = { ϕ t , ϕ t + 1 , … , ϕ T } , ϕ t = Φ ( s t , a t ) ∈ F \tau_t^\Phi=\{\phi_t,\phi_{t+1},\ldots,\phi_T\},\phi_t=\Phi(s_t,a_t)\in F τtΦ​={ϕt​,ϕt+1​,…,ϕT​},ϕt​=Φ(st​,at​)∈F，以及该集合的信息统计量为$I^{\Phi }({\tau }_t)$。将信息匹配(Information Match，IM)问题定义为学习一个条件策略π(a|s, z)，执行该策略所产的轨迹满足一些期望的信息统计值z：
$$\underset{\pi }{\min }{\mathbb{E}}_{z\sim p(z),\tau \sim {\rho }_z^{\pi }(\tau )}\left[\begin{array}{c}D(I^{\Phi }(\tau ),z)\end{array}\right]$$
上式说明当D=0时，$z^{\ast }=I^{\Phi }$, 此时轨迹中状态动作对于z便是最优，那就可以通过BC得到策略或者利用样本加速RL策略学习。

GDT

该方法基于DT算法，其中输入量中引入轨迹数据的信息统计值$I^{\Phi }$。优化该模型预测动作的MSE损失，还需最小化信息值与z的先验分布间距离。不同的特征提取函数$\Phi$以及$I$将造就不同的GDT变体。下面讨论两种变体CDT与BDT

CDT

启发于Distribution RL，采用直方图的作为特征空间表达，然后加入到序列模型的输入，并且随着时间步采用类似bellman的递归更新方式。从代码中可以看出，这里对奖励作为变量$z$
$$z_{{\Phi }_{\mathsf{hist}}}(t)\propto {\tilde{\varphi }}_t+\gamma (1-1[t=T])z_{{\Phi }_{\mathsf{hist}}}(t+1)$$

DECISION TRANSFORMER WITH LEARNED Φ

CDT 假设提供了一些低维 Φ 用于易于处理的分布近似，但文章也研究了没有提供 Φ 或 r 的情况，而是通过自动编码来学习 。文章提出使用auto-encoding或者对比学习的方法得到$\Phi$:

1. CDT-AE:
   $$\underset{\psi }{\min }{\mathbb{E}}_{s\sim D}\left[\Vert s-{\text{decoder}}_{\psi }({\text{encoder}}_{\psi }(s)){\Vert }^2\right]$$
2. CDT-CL:
   $$\underset{\psi }{\min }{\mathbb{E}}_{\stackrel{s\sim D,}{ϵ,{ϵ}^{\prime }\sim \mathcal{N}(0,\sigma )}}\left[\log \frac{\exp \left({\text{encoder}}_{\psi }(s{)}^{T}W{\text{encoder}}_{\tilde{\psi }}(s+ϵ)\right)}{{\sum }_{{ϵ}^{\prime }\ne ϵ}\exp \left({\text{encoder}}_{\psi }(s{)}^{T}W{\text{encoder}}_{\tilde{\psi }}(s+{ϵ}^{\prime })\right)}\right],$$
   其中$\tilde{\psi }\leftarrow m\tilde{\psi }+(1-m)\psi$

BI-DIRECTIONAL DECISION TRANSFORMER FOR ONE-SHOT IMITATION LEARNING

除了上述方法对$\Phi$参数化处理，还有通过参数聚合器(parameterized aggregator)。启发于One-shot IL 方法GPT-1提出Bi-directional Decision Transformer (BDT). 其中anti-causal aggregator将z反序后编码，然后作为序列模型的输入