ICLR 2023
paper

Intro

采用Transformer架构的Planning方法对马尔可夫序列重构,(et. TT)在面对高维状态动作空间，容易面对计算复杂度高的问题。本文提出TAP算法，基于Transformer的VQ-VAE，利用提取的状态动作在隐空间的低微特征进行Planning，然后使用latent codes经过decoder得到重构序列，在Offline下取较好的结果。

Method

VQ-VAE

训练VQ-VAE使用离线数据$\tau =({\mathbf{s}}_1,{\mathbf{a}}_1,r_1,R_1,{\mathbf{s}}_2,{\mathbf{a}}_2,r_2,R_2,\dots ,{\mathbf{s}}_T,{\mathbf{a}}_T,r_T,R_T)$。以上图为例，经过encoder得到T个特征(图中T=9)，然后步长为L的一维卷积以及最大池化得到向量$({\bar{x}}_1,{\bar{x}}_2,{\bar{x}}_3)$。在由最近邻找到对应的codebook中的$e_i$作为latent code。
$${\mathbf{z}}_i={\mathbf{e}}_k,\mathrm{where}k={\mathrm{argmin}}_j||{\mathbf{x}}_i-{\mathbf{e}}_j|{|}_2$$
解码阶段，首先将latentcode扩展，与输入等维度。concat初始状态，经过decoder得到重构的序列。损失函数则是由原序列与重构序列的均方误差。除此外还最小化特征向量、latent code分别与codebook的距离:$||{\mathbf{x}}_i-{\mathbf{e}}_k|{|}_2\mathrm{and}||{\mathbf{z}}_i-{\mathbf{e}}_k|{|}_2$

得到latent code后，还需要训练其先验分布用于后续的Planning过程。TAP采用Transformer架构的自回归模型$p({\mathbf{z}}_t|{\mathbf{z}}_{<t},{\mathbf{s}}_1)=p({\mathbf{z}}_t|{\mathbf{s}}_1,{\mathbf{z}}_1,{\mathbf{z}}_2,...,{\mathbf{z}}_{t-1})$构建更加紧凑的latent code

Planning

使用先验模型，生成当前state在隐空间中的latent code序列，再用decoder进行解码就得到预测的轨迹。对每条生成轨迹有如下评价函数
$$g({\mathbf{s}}_1,{\mathbf{z}}_1,{\mathbf{z}}_2,...,{\mathbf{z}}_M)=\sum _t{\gamma }^t{\hat{r}}_t+{\gamma }^T{\hat{R}}_T+\alpha \ln \left(\min (p({\mathbf{z}}_1,{\mathbf{z}}_2,...,{\mathbf{z}}_M|{\mathbf{s}}_1),{\beta }^M)\right)$$
前两项衡量轨迹累计折扣奖励，后一项则是惩罚项，如果轨迹有概率大于阈值则相信累计奖励。而若小于阈值，则后项由于权重$\alpha$取值为大于折扣回报的最大值，使得此时后项对评分的主导远大于累计奖励，即选择高概率的轨迹。

Beam Serach

基于初始状态以及先验模型，采样生成latent code z采用Beam search

可以看出，首先利用先验模型采样得到n个$z_1$，然后对每个$z_1$由先验模型得到概率最高的排序为前E个的$z_2$拼接，然后由decoder解码并根据评价函数得到轨迹分数，选取Top-B的序列的$z_1,z_2$, 重复上述过程选取最大评分的轨迹。

总结

对于高维复杂环境，通过encoder到低维度隐空间进行推理学习好的特征表示，在decoder生成是一个好的框架。对于下游任务，就可以直接采用学习好的特征以及decoder实现zero-shot。这是一个不错的思路。后续ICLR2024有个工作使用在隐空间进行Diffusion：《Efficient Planning with Latent Diffusion》