Latent Action Models (LAMs) learn to extract action-relevant representations solely from raw observations, enabling reinforcement learning from unlabelled videos and significantly scaling available training data. However, LAMs face a critical challenge in disentangling action-relevant features from action-correlated noise (e.g., background motion). Failing to filter these distractors causes LAMs to capture spurious correlations and build sub-optimal latent action spaces. In this paper, we introduce MaskLAM -- a lightweight modification to LAM training to mitigate this issue by incorporating visual agent segmentation. MaskLAM utilises segmentation masks from pretrained foundation models to weight the LAM reconstruction loss, thereby prioritising salient information over background elements while requiring no architectural modifications. We demonstrate the effectiveness of our method on continuous-control MuJoCo tasks, modified with action-correlated background noise. Our approach yields up to a 4x increase in accrued rewards compared to standard baselines and a 3x improvement in the latent action quality, as evidenced by linear probe evaluation.

翻译：潜在动作模型（LAMs）能够仅从原始观测中学习提取与动作相关的表征，从而支持从未标记视频中进行强化学习，并显著扩展可用的训练数据。然而，LAMs面临一个关键挑战：如何将动作相关特征与动作相关的噪声（例如背景运动）解耦。若未能过滤这些干扰物，会导致LAMs捕捉到虚假相关性并构建次优的潜在动作空间。本文提出MaskLAM——一种对LAM训练的轻量级改进方法，通过引入视觉智能体分割来缓解此问题。MaskLAM利用预训练基础模型产生的分割掩码对LAM重建损失进行加权，从而在无需修改模型架构的情况下，优先考虑显著信息而非背景元素。我们在添加了动作相关背景噪声的连续控制MuJoCo任务上验证了该方法的有效性。实验表明，与标准基线相比，我们的方法使累计奖励最高提升4倍，并通过线性探针评估证明潜在动作质量提高了3倍。