TWM论文阅读笔记

• 这是ICLR2023的一篇world model论文，用transformer来做世界模型的sequence prediction。
• 文章贡献是transformer-based world model（不同于以往的如transdreamer的world model，本文的transformer-based world model在inference 的时候可以丢掉）两个损失，一个采样策略。

WM

• TWM用的仍然是经典的world model框架：
  • encoder-decoder用的是VAE，不过输入是四帧而不是一帧
  • dynamic model用过去 $l$ 步的 $z$ $a$ $r$ 和当前的$z$ $a$作为输入，用transformer预测 $h_t$，再用$h_t$预测$r_t$，${\gamma }_t$和 $z_{t+1}$，如下：
    
  • 上面的三个$p$都是MLP，$f$是transformerXL，$3l-1$个token输入，预测一个token：
    
  • z，r，gamma的MLP的输出分别是：a vector of independent categorical distributions, a normal distribution,
    and a Bernoulli distribution
• 提的两个损失，一个是如下的encoder-decoder的损失，由三项组成，第一项是VAE的损失，第二项是对z的熵损失，第三项是与sequence model的一致损失：
  
• 第二个损失是用来train sequence model的，第一项其实跟上一个损失的第三项一样，但是上一个损失是train VAE的，这个损失是train sequence model的；第二项第三项不用说，就是正常的reward和discount的损失：
  
• 相比LSTM GRU之类模型，transformer的好处在于长序列建模，sequence model总是能看到过去$l$步发生的确切的事情，而非仅能观察到一个压缩的状态$h_t$

RL

• 这里可以看到，dreamerv3等模型预测的是奖励 $r_t$ 和terminate $d_t$，但是TWM预测的是discount factor $\gamma$，在这里就可以派上用场了，预测的$\gamma$用来train RL模型（而其他的WM，RL模型的$\gamma$用的是固定值）。那么训练的时候$\gamma$怎么监督呢，文章定义的label是${\gamma }_t=0$和${\gamma }_t=\gamma$，即当terminate的时候$\gamma =0$而其他时候$\gamma$是固定值，label是这样，而模型应该会灵活预测？不知道
• 这里有一些新的损失，比如对策略的熵的损失，不能低于一个阈值：
  
• 这里还要关注的是policy model的输入，一般policy model的输入是$z_t$和$h_t$的concate,如dreamerv3和STORM，文章试了发现decoder的输出也可以（IRIS就这么干的），$o_t$也可以，本文用的是$z_t$，比较轻量快速，只需要encoder而不需要sequence model。并且，训练的时候用的是sequence model预测的zt，而测试的时候则用的是encoder编码的zt加上frame stacking操作（这里有点疑问，维度？）
• train的时候还是常规的三步走：用RL model采样，train world model，用world model train RL model。
• training的时候有个sampling的stategy，如下， 是为了让模型更关注后面采样得到的sample，但vt的公式也没给，之说是incremented every time an entry is sampled：