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
(
τ
)
I(\tau)
I(τ),此外定义对轨迹数据中(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
Φ
(
τ
t
)
I^\Phi(\tau_t)
IΦ(τt)。将信息匹配(Information Match,IM)问题定义为学习一个条件策略π(a|s, z),执行该策略所产的轨迹满足一些期望的信息统计值z:
min
π
E
z
∼
p
(
z
)
,
τ
∼
ρ
z
π
(
τ
)
[
D
(
I
Φ
(
τ
)
,
z
)
]
\min_\pi\mathbb{E}_{z\sim p(z),\tau\sim\rho_z^\pi(\tau)}\begin{bmatrix}D(I^\Phi(\tau),z)\end{bmatrix}
πminEz∼p(z),τ∼ρzπ(τ)[D(IΦ(τ),z)]
上式说明当D=0时,
z
∗
=
I
Φ
z^* = I^\Phi
z∗=IΦ, 此时轨迹中状态动作对于z便是最优,那就可以通过BC得到策略或者利用样本加速RL策略学习。
GDT
该方法基于DT算法,其中输入量中引入轨迹数据的信息统计值 I Φ I^\Phi IΦ。优化该模型预测动作的MSE损失,还需最小化信息值与z的先验分布间距离。不同的特征提取函数 Φ \Phi Φ以及 I I I将造就不同的GDT变体。下面讨论两种变体CDT与BDT
CDT
启发于Distribution RL,采用直方图的作为特征空间表达,然后加入到序列模型的输入,并且随着时间步采用类似bellman的递归更新方式。从代码中可以看出,这里对奖励作为变量
z
z
z
z
Φ
h
i
s
t
(
t
)
∝
ϕ
~
t
+
γ
(
1
−
1
[
t
=
T
]
)
z
Φ
h
i
s
t
(
t
+
1
)
z_{\Phi_{\mathsf{hist}}}(t)\propto\tilde{\phi}_t+\gamma(1-\mathbb{1}[t=T])z_{\Phi_{\mathsf{hist}}}(t+1)
zΦhist(t)∝ϕ~t+γ(1−1[t=T])zΦhist(t+1)
DECISION TRANSFORMER WITH LEARNED Φ
CDT 假设提供了一些低维 Φ 用于易于处理的分布近似,但文章也研究了没有提供 Φ 或 r 的情况,而是通过自动编码来学习 。文章提出使用auto-encoding或者对比学习的方法得到 Φ \Phi Φ:
- CDT-AE:
min ψ E s ∼ D [ ∥ s − decoder ψ ( encoder ψ ( s ) ) ∥ 2 ] \min_\psi\mathbb{E}_{s\sim D}\left[\|s-\text{decoder}_\psi(\text{encoder}_\psi(s))\|^2\right] ψminEs∼D[∥s−decoderψ(encoderψ(s))∥2] - CDT-CL:
min ψ E ϵ , ϵ ′ ∼ N ( 0 , σ ) s ∼ D , [ log exp ( encoder ψ ( s ) T W encoder ψ ~ ( s + ϵ ) ) ∑ ϵ ′ ≠ ϵ exp ( encoder ψ ( s ) T W encoder ψ ~ ( s + ϵ ′ ) ) ] , \min_\psi\mathbb{E}_{\overset{s\sim D,}{\epsilon,\epsilon^{\prime}\sim\mathcal{N}(0,\sigma)}}\left[\log\frac{\exp\left(\text{encoder}_\psi(s)^TW\text{encoder}_{\tilde{\psi}}(s+\epsilon)\right)}{\sum_{\epsilon^{\prime}\neq\epsilon}\exp\left(\text{encoder}_\psi(s)^TW\text{encoder}_{\tilde{\psi}}(s+\epsilon^{\prime})\right)}\right], ψminEϵ,ϵ′∼N(0,σ)s∼D, log∑ϵ′=ϵexp(encoderψ(s)TWencoderψ~(s+ϵ′))exp(encoderψ(s)TWencoderψ~(s+ϵ)) ,
其中 ψ ~ ← m ψ ~ + ( 1 − m ) ψ \tilde{\psi}\leftarrow m\tilde{\psi}+(1-m)\psi ψ~←mψ~+(1−m)ψ
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反序后编码,然后作为序列模型的输入