Segment to Focus: Guiding Latent Action Models in the Presence of Distractors

TL;DR

MaskLAM tackles the problem of latent action models being misled by action-correlated distractors in unlabelled videos. It introduces a lightweight, plug-and-play approach that uses segmentation masks to weight the LAM reconstruction loss, biasing the latent space toward agent-centric dynamics without architectural changes. Empirically, MaskLAM achieves up to 4x improvements in acyclic rewards and 3x gains in latent action quality across four MuJoCo tasks with distractors, while enabling smaller latent spaces and better generalisation to out-of-distribution noise. This method enhances the practicality of unsupervised pretraining from video data for imitation learning and robotics, with broad implications for data-efficient RL/IL pipelines.

Abstract

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.

Figures (13)

• Figure 1: Visualised architecture of LAPO and our proposed modification --- MaskLAM. MaskLAM employs agent-centric segmentation masks, $M_{t+1}$, to weigh the reconstruction objective, discouraging the representation of background information and prioritising action-relevant features in the latent actions.
• Figure 2: Observations and segmentation masks. Examples of the augmented observations in Cheetah and Hopper from the Distracting Control Suite used in our experiments (left images) and the segmentation masks used in MaskLAM (right images).
• Figure 3: Aggregate downstream performance. We report the average normalised returns across the four tested tasks. (a) MaskLAM consistently outperforms standard LAPO across all latent sizes, recovering a large portion of the performance lost to distractors compared to the distractor-free LAPO. (b) Sample efficiency: MaskLAM demonstrates superior sample efficiency, achieving higher returns with fewer labelled trajectories. Notably, it matches LAPO's peak performance with only $\sim$4 labels.
• Figure 4: Sample efficiency in downstream control. We compare the normalised evaluation returns of the downstream policy vs. the number of labelled trajectories used to train the action decoder (latent dimension is fixed at 8192). MaskLAM consistently outperforms the LAPO baseline in the presence of distractors, recovering a significant portion of the performance gap relative to the distractor-free oracle. MaskLAM similarly demonstrates superior sample efficiency, achieving higher returns with fewer labelled demonstrations.
• Figure 5: Impact of latent action dimension on performance. We analyse the sensitivity of the downstream policy to the size of the latent action space ($d \in [64, 8192]$), using a fixed budget of 128 labelled trajectories. MaskLAM consistently outperforms LAPO across the latent dimensionality and is competitive against the distractor-free LAPO.