ICLR 2024 poster
paper
Intro
大型预训练模型结合强化学习,通过自适应参数微调实现对new-task的快速适应。
Method
演示数据由K个不同Task构成的集合
{
T
1
,
…
,
T
K
}
.
Each task
T
k
=
(
μ
k
0
,
g
k
)
\{\mathcal{T}_{1},\ldots,\mathcal{T}_{K}\}.\text{ Each task }\mathcal{T}_{k}=(\mu_{k}^{0},g_{k})
{T1,…,TK}. Each task Tk=(μk0,gk)。
μ
k
0
\mu_{k}^{0}
μk0描述初始状态分布,g则表示目标,本文用语言指令描述。单个任务可获得N条轨迹,TAIL基于BC优化智能体策略:
θ
^
=
min
θ
∑
k
=
1
K
E
s
t
,
a
t
∼
D
k
[
∑
t
=
0
l
k
L
(
π
(
a
∣
s
≤
t
,
T
k
;
θ
)
,
a
k
t
)
]
\hat{\boldsymbol{\theta}}=\min_{\boldsymbol{\theta}}\sum_{k=1}^K\mathbb{E}_{s_t,a_t\sim\mathcal{D}_k}\left[\sum_{t=0}^{l_k}\mathcal{L}\left(\pi(a|s_{\leq t},\mathcal{T}_k;\theta),a_k^t\right)\right]
θ^=θmink=1∑KEst,at∼Dk[t=0∑lkL(π(a∣s≤t,Tk;θ),akt)]
TAIL的整体流程可描述为:(1)输入任务及其数据集,初始化task-specific adapter ω k \omega_k ωk (2)adapter结合预训练的参数模型,采用BC损失函数训练优化adapter,此时需要冻结预训练模型的参数。
那么如何获得预训练的模型?如图Fig1(a)所示的预训练模型架构方法(更具体见Fig6)。TAIL使用预训练的CLIP实现语言指令编码(instruction encoder)以及图像数据的空间编码(spatial encoder),并通过一个input fusion model将二者融合,实更好的指令跟随能力。fusion结果将联合空间编码所得到的token序列,输入到transformer架构的Temporal decoder得到序列动作预测,结合BC实现对整个模型的参数优化。
完成上述步骤,TAIL提出使用一个额外的小参数量的、可被优化的插件adapter,联合参数冻结的预训练模型进行微调。结构如图Fig1(b)的最下方的结构。而Adapter的结构如图所示
结构上依旧基于transformer,增加三个模块:Parallel , Sequential , 以及Prefix Token。
-
Parallel Integration (LoRA)
h o u t = W ⊤ h i n + α W u p ⊤ W d o w n ⊤ h i n h_{out}=W^{\top}h_{in}+\alpha W_{up}^{\top}W_{down}^{\top}h_{in} hout=W⊤hin+αWup⊤Wdown⊤hin -
Sequential Integration (Bottleneck Adapter) h o u t = W u p T ϕ ( W d o w n T ( W T h i n ) ) h_{out}=\boldsymbol{W}_{up}^{\mathrm{T}}\phi\left(\boldsymbol{W}_{down}^{\mathrm{T}}(\boldsymbol{W}^{\mathrm{T}}h_{in})\right) hout=WupTϕ(WdownT(WThin))
其中 ϕ \phi ϕ为非线性激活函数 -
Prefix Token Integration (Prefix & Prompt-Tuning)
对Q,K以及V添加可学习token,可以被视为任务描述符,旨在引导模型朝向所需的特定于任务的行为。