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}$$