Technical Summary: MaskLAM – Guiding Latent Action Models in the Presence of Distractors

Problem Statement

Latent Action Models (LAMs) offer a pathway to pre-train embodied agents on large-scale, action-free video by inferring latent actions between consecutive observations. These latent actions are subsequently decoded into ground-truth controls using a small set of labeled examples. However, standard LAMs, such as LAPO, rely on global reconstruction objectives that force the model to encode the entire scene to predict future frames.

In realistic environments containing visual distractors (e.g., dynamic backgrounds, camera shake, or moving objects), this approach fails. The standard reconstruction loss drives the latent action to encode exogenous motion (distractors) rather than agent-controlled (endogenous) dynamics. This phenomenon, formalized by the Exogenous Block MDP (Ex-BMDP) framework, results in latent spaces that do not recover true actions, causing downstream policies to underperform significantly when fine-tuned. Prior attempts to solve this have required architectural changes, auxiliary losses, or privileged action supervision during pre-training, increasing complexity and data requirements.

Methodology: MaskLAM

The authors propose MaskLAM, a lightweight modification to the LAPO training objective that addresses the exogenous contamination problem without altering the model architecture or requiring action labels during pre-training.

Core Insight

The method relies on the observation that while endogenous (agent) and exogenous (distractor) factors are entangled in the latent space, they are typically spatially separable in pixel space. Agent-controlled changes are concentrated on the agent's pixels, while distractor motion occurs elsewhere.

Technical Implementation

1. Zero-Shot Segmentation: MaskLAM utilizes off-the-shelf segmentation foundation models (specifically SAM 2.1) to generate binary masks ($M_t$) identifying agent pixels. These masks are obtained zero-shot from a single prompt (e.g., a bounding box) on the first frame and propagated through the video sequence.
2. Masked Forward Dynamics Loss: Instead of reconstructing the full observation, the forward dynamics model (FDM) loss is restricted to agent pixels. The loss function is defined as:
   $$L_{MaskLAM}^{FDM} = \frac{1}{\|M_{t+1}\|_1} \| M_{t+1} \odot (\hat{o}_{t+1} - o_{t+1}) \|_2^2$$
   where $\odot$ denotes the Hadamard product. This ensures that pixels outside the agent mask (distractors) contribute zero gradient to the latent action $z_t$.
3. Mask-Conditioned Inputs: The segmentation masks are concatenated as an additional input channel to both the Inverse Dynamics Model (IDM) and the FDM. This allows the models to utilize the full scene context (including objects the agent interacts with that may fall outside the strict agent mask) while ensuring the learning signal for $z_t$ is derived only from agent-controlled dynamics.
4. No Architectural Changes: The intervention modifies only the loss calculation and input channels, leaving the underlying LAPO architecture (IDM/FDM pair) and the three-stage training pipeline (pre-training, relabeling, decoder fine-tuning) intact.

Key Contributions

1. Spatial Separability as an Intervention: The paper identifies the spatial separability of agent and distractor pixels as the key property that allows the abstract Ex-BMDP factorization to be implemented as a concrete, training-time intervention.
2. MaskLAM Proposal: A method that eliminates exogenous gradients on the latent action by restricting the reconstruction loss to agent pixels, requiring no auxiliary losses, no architectural changes, and no action labels during pre-training.
3. Empirical Validation: Extensive evaluation across two benchmarks (Distracting Control Suite and Distracting Meta-World) demonstrates that MaskLAM significantly outperforms label-free baselines and narrows the performance gap to methods that rely on privileged action supervision (LAOM-Labels).

Results

The authors evaluated MaskLAM on 14 tasks across two benchmarks, comparing it against LAPO, LAOM-Labels (upper bound with action supervision), LAOF (flow-based), and LAPO-Slots (object-centric).

• Latent Action Quality: MaskLAM reduced the normalized linear-probe Mean Squared Error (NMSE) by up to 3.51× compared to LAPO. In visualizations (UMAP projections), MaskLAM produced smooth, monotonic manifolds aligned with ground-truth actions, whereas LAPO collapsed into entangled clusters.
• Downstream Performance: MaskLAM improved normalized return by up to 4.97× over LAPO in distractor settings. It matched or exceeded all label-free baselines and approached the performance of LAOM-Labels, which uses privileged action supervision.
• Sample Efficiency: MaskLAM demonstrated superior sample efficiency in the decoder fine-tuning stage. On Distracting Meta-World, it reached near-peak performance with as few as 2,000 action labels, whereas LAPO and LAPO-Slots required nearly two orders of magnitude more labels to close the gap.
• Robustness:
  • Occlusion: MaskLAM maintained low error rates even with up to 75% agent occlusion, whereas unmasked models degraded significantly.
  • Imperfect Masks: The method remained effective even when segmentation masks were perturbed (mIoU dropped to 0.4–0.6), indicating that pixel-perfect segmentation is not required.
  • Compact Latent Spaces: MaskLAM allowed for smaller latent dimensions (e.g., 64-dim) to achieve the same alignment as larger unmasked latents (256-dim), as the latent space is not burdened with encoding exogenous variance.

Significance and Claims

The paper claims that the failure of standard LAMs in distractor-rich environments is driven by the prediction target (global reconstruction) rather than the architecture. By re-allocating the supervision burden from scarce action labels to cheap, zero-shot available segmentation masks, MaskLAM makes latent action pre-training viable on noisy, real-world video.

The authors emphasize that this approach:

• Decouples the need for action labels from the ability to learn disentangled dynamics.
• Preserves the simplicity of the LAPO framework while solving a critical robustness issue.
• Demonstrates that spatial priors derived from foundation models can effectively guide representation learning in the absence of explicit control annotations.

The work suggests that for embodied agents operating in complex visual environments, the supervision required to disentangle agent dynamics from distractors can be sourced from segmentation masks rather than action labels, potentially enabling large-scale pre-training on internet-scale video data.