Segment to Focus: Guiding Latent Action Models in the Presence of Distractors
This addresses a critical challenge in reinforcement learning from unlabelled videos by improving LAM robustness to distractors, though it is incremental as a lightweight modification to existing methods.
The paper tackled the problem of Latent Action Models (LAMs) capturing spurious correlations from action-correlated noise like background motion, and introduced MaskLAM, which uses segmentation masks to weight the reconstruction loss, resulting in up to a 4x increase in rewards and a 3x improvement in latent action quality on MuJoCo tasks.
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.