Pretrain: Learning to Act without Actions Learning to Act without Actions

ICLR 2024 Spotlight
paper

Intro

利用海量网络视频数据对模型预训练是提升算法效率的有效方法。然而此类数据缺乏offline下的任务动作标签，因此应用于RL困难。本文提出LAPO则是从只包含观测的演示中推断潜在动作在以及隐空间策略表现形式，为在线RL以及离线RL提供预训练的策略。

Method

Latent action representation

LAPO预训练阶段包含两个模型:逆动力学模型$P_{IDM}(z_t|o_t,o_{t+1})$和前向动力学模型$P_{FDM}(o_{t+1}|o_t,z_t)$,

而在实际使用中，LAPO采样一个k-steps sequence的历史数据以及下一时刻的观测预测latent action $z_t$
$$z_t\sim p_{\mathrm{IDM}}(\cdot |o_{t-k},\dots ,o_t,o_{t+1})$$
FDM将继续使用历史数据以及预测的动作来预测下一时刻的观测
$${\hat{o}}_{t+1}\sim p_{\mathrm{FDM}}(\cdot |o_{t-k},\dots ,o_t,z_t)$$
然后通过一个BC损失函数联合优化两个模型：$||{\hat{o}}_{t+1}-o_{t+1}|{|}^2$
为了防止IDM直接将$o_{t+1}$copy到$z$,然后FDM直接预测${\hat{o}}_{t+1}=o_{t+1}$。LAPO对z设置信息瓶颈，强制压缩IDM到FDM的信息量，防止IDM与FDM共有的($o_{t-k},\dots ,o_t$)包含的信息被忽略。同时，由于IDM和FDM都获得历史观测，学习到的潜在行为可能成为这些观察的条件。在不同的状态下，一些潜在的动作z可能对应于不同的真实动作。

因此，在将每个潜在动作传递给FDM之前，采取矢量量化（Vector Quantization，VQ）（EMA-based update），从而迫使IDM在不同状态空间重复使用有限数量的离散latent，实现解耦特征表示。

BC a latent policy

在相同数据集，利用BC方法学习latent policy $\pi :\mathcal{O}\to \mathcal{Z}$
$$||\pi (o_t)-z_t|{|}^2\mathrm{where}{z}_t\sim p_{\mathrm{IDM}}(\cdot |o_{t-k},\dots ,o_t,o_{t+1})$$

Decoding latent action

latent policy 将观测值映射到隐空间，而对 RL，则需要从隐空间decoding得到具体的action与环境交互。对于在线RL只需将latent policy head替换为智能体策略的head。而对于离线RL，已知action-label，只需要单独训练一个decoder $d:\mathcal{Z}\to \mathcal{A}$

Limitations

1. 在观测中具有延迟效应的行动将同样被预测延迟发生，即潜在政策实际上模拟了一个行动的可见效应，而不是行动本身。然而，在大多数环境中，对环境状态有任何影响的行为将会引起观测变量的某种程度的即时变化。
2. 此外，可以通过扩展IDM和FDM体系结构来考虑过去和未来的多个时间步长来部分解决延迟操作，例如通过使用Transformer结构。
3. 显著的随机性会使IDM难以压缩噪声中的有用的信息，从而降低latent action质量。可以通过在更大的数据集上进行训练来缓解。
4. 在更大的数据集上进行训练，将需要扩大模型架构，这在平衡FDM的强度和latent action表征能力方面带来了新的挑战