Pretrain: TAIL: Task-specific Adapters for Imitation Learning with Large Pretrained Models

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​)。${\mu }_k^0$描述初始状态分布，g则表示目标，本文用语言指令描述。单个任务可获得N条轨迹，TAIL基于BC优化智能体策略：
$$\hat{\mathbf{\theta }}=\underset{\mathbf{\theta }}{\min }\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_{\le t},{\mathcal{T}}_k;\theta ),a_k^t\right)\right]$$

TAIL的整体流程可描述为:（1）输入任务及其数据集，初始化task-specific adapter ${\omega }_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。

1. Parallel Integration (LoRA)
   $$h_{out}=W^{\top }{h}_{in}+\alpha W_{up}^{\top }{W}_{down}^{\top }{h}_{in}$$

2. Sequential Integration (Bottleneck Adapter)$$h_{out}={\mathbf{W}}_{up}^{\mathrm{T}}\varphi \left({\mathbf{W}}_{down}^{\mathrm{T}}({\mathbf{W}}^{\mathrm{T}}{h}_{in})\right)$$
   其中$\varphi$为非线性激活函数

3. Prefix Token Integration (Prefix & Prompt-Tuning)
   对Q，K以及V添加可学习token，可以被视为任务描述符，旨在引导模型朝向所需的特定于任务的行为。