Offline RL : StARformer: Transformer With State-Action-Reward Representations for Robot Learning

IEEE PAMI 2023
ECCV 2022
paper
code
额外添加Transformer捕捉单时间步中的关键信息，并作为goal加入到T序列模型中学习策略

Intro

1. StARformer由Step Transformer和Sequence Transformer两个基本组件组成，并通过交错连接进行交互。
2. Step Transformer学习局部表示（即StAR-representations），通过在单个时间步长的窗口内自注意力机制处理状态-动作-奖励token。
3. Sequence Transformer结合StAR-representations和纯图像状态表示（卷积特征提取），对整个序列进行自注意力操作以进行动作预测。

method

Step Transformer

将一段轨迹按$(a_{t-1},r_{t-1},s_t)$分成多个group。其中状态图片分解为N个Patch，

State-Action-Reward Embeddings

对状态，动作以及奖励分别按照如下进行embedding
$$\begin{array}{rl}{z}_{s_t^i}& =\mathrm{FC}(\mathrm{Flatten}(s_t^i))+e_i^{\mathrm{spatial}}\\ z_{a_{t-1}}& =\mathrm{FC}(a_{t-1}),z_{r_t}=\mathrm{Tanh}(\mathrm{FC}(r_t)).\end{array}$$
其中$e_i^{\mathrm{spatial}}$表示每个patch的空间位置编码信息。这样，得到轨迹的embedding的集合 Z t 0 = { z a t − 1 , z r t , z s t 1 , z s t 2 , … , z s t n } . Z_t^0=\{z_{a_{t-1}},z_{r_t},z_{s_t^1},z_{s_t^2},\ldots,z_{s_t^n}\}. Zt0​={zat−1​​,zrt​​,zst1​​,zst2​​,…,zstn​​}..

Step Transformer Layer

将 Z t 0 = { z a t − 1 , z r t , z s t 1 , z s t 2 , … , z s t n } . Z_t^0=\{z_{a_{t-1}},z_{r_t},z_{s_t^1},z_{s_t^2},\ldots,z_{s_t^n}\}. Zt0​={zat−1​​,zrt​​,zst1​​,zst2​​,…,zstn​​}.编码作为下一个Step Transformer Layer的输入。
$$Z_t^l=F_{\mathrm{step}}^l(Z_t^{l-1}).$$
同时该层还需输出所有group的StAR-representations给到 Sequence Transformer
$$g_t^l=\mathrm{FC}([Z_t^l])+e_t^{\mathrm{temporal}}.$$
其中$[\cdot ]$表示将所有group的特征信息进行concatenate，${}_t^{\mathrm{temporal}}$表示时间编码信息。

Sequence Transforme

对状态进行卷积
$$h_t^0=\mathrm{Conv}(s_t)+e_t^{\mathrm{temporal}},$$
联合StAR-representations得到Sequence Transformer 的输入序列 : Y i n l = { g 1 l , h 1 l − 1 , g 2 l , h 2 l − 1 , … , g T l , h T l − 1 } . Y_{\mathrm{in}}^l=\{g_1^l,h_1^{l-1},g_2^l,h_2^{l-1},\ldots,g_T^l,h_T^{l-1}\}. Yinl​={g1l​,h1l−1​,g2l​,h2l−1​,…,gTl​,hTl−1​}.
然后输入到序列模型得到2T个输出，为了保持下一个layer的输入为T个，如图中所示选择index为$2i$的预测结果。
Y o u t l = { y o u t ; 1 l , y o u t ; 2 l , … , y o u t ; 2 T l } h i l : = y o u t ; 2 i l . \begin{aligned}Y_{\mathrm{out}}^{l}&=\{y_{\mathrm{out};1}^{l},y_{\mathrm{out};2}^{l},\ldots,y_{\mathrm{out};2T}^{l}\}\\h_{i}^{l}&:=y_{\mathrm{out};2i}^l.\end{aligned} Youtl​hil​​={yout;1l​,yout;2l​,…,yout;2Tl​}:=yout;2il​.​
最后一层的Layer将以$h$为输入，通过一层线性映射直接预测动作$:{\hat{a}}_t=\varphi (h_t^l).$

results

然后对grouping以及patch进行消融实验

对sequence transformer如何利用step transformer表征信息进行结构上的消融实验