LAOF: Robust Latent Action Learning with Optical Flow Constraints

AI summary

Overview

Research area: Robot learning / embodied AI — specifically latent action learning (learning action representations from large-scale video without action labels), evaluated on imitation learning and reinforcement learning benchmarks.

Technical level: Advanced. The paper assumes familiarity with the Latent Action Policies (LAPO) autoencoding framework, inverse and forward dynamics models, VQ-VAE discretization, DINOv2 features, and optical flow estimation (RAFT, LangSAM).

Scope: The paper proposes LAOF, a framework that uses an agent's pixel-level optical flow as a pseudo-supervision signal to make latent action representations robust to action-irrelevant distractors, and evaluates it on the LIBERO and PROCGEN benchmarks.

What This Paper Is About

Latent action methods like LAPO learn compact "action-like" codes by jointly training an inverse dynamics model (IDM) and forward dynamics model (FDM) to reconstruct the next frame. This assumes every change between two frames comes from the agent. In real large-scale video, backgrounds move, distractors drift, and visual appearance gets entangled with motion, so the learned latent actions become noisy and unstable.

The paper's goal is to constrain latent actions with the agent's optical flow — pixel-level motion between consecutive frames that naturally suppresses static background and emphasizes the moving agent — so that latent actions stay anchored to real physical motion, especially when action labels are extremely scarce. A dedicated flow decoder maps latent actions directly into optical flow, providing a supervision signal that requires no manual annotation.

Key Contributions

1. Optical flow constraints for latent action learning. LAOF adds a flow decoder $d_{\text{flow}}: \mathcal{Z} \rightarrow \mathcal{F}{\text{rgb}}$ that maps latent actions to optical flow features, trained with an $\mathcal{L}{\text{flow}}$ loss added to the standard next-state reconstruction loss. The framework is validated on downstream imitation learning tasks on LIBERO and reinforcement learning tasks on PROCGEN.

2. LAOF-Action for sparse action labels. When a small set of labels exists, an action decoder (a lightweight MLP) applies action supervision to labeled samples while optical flow supervises the unlabeled ones, with the combined objective $\mathcal{L}{\text{reconstruction}} + (1-\lambda)\cdot\mathcal{L}{\text{flow}} + \lambda\cdot\mathcal{L}_{\text{action}}$ and $\lambda = \frac{M}{N+M}$. The approach is evaluated as the action-label proportion grows up to 10%, and LAOF without any action supervision is shown to match or outperform action-supervised methods trained with a 1% action ratio.

3. Ablation over constraint structure and coupling. Five variants are compared (LAOF, LAOF-FlowFDM, LAOF-Only($z_t$), LAOF-Only($z_t, s_t$), LAOF-AE), and attaching a dedicated flow decoder directly to the latent actions performs best.

4. Optical flow pipeline for both static and dynamic distractors. Global optical flow (estimated with RAFT) is used where distractors are static, and object-centric optical flow (masked with LangSAM using textual prompts, $f^{\text{sam}}_{\text{rgb},t} = \text{mask}t \odot f^{\text{all}}{\text{rgb},t}$) is used where distractors move. Raw flow is converted to RGB format following VideoJAM to match DINOv2 input.

Main Findings

• Optical flow improves unsupervised and supervised latent actions on LIBERO. Table 1 (continuous latent actions, LIBERO, 1000 trials per evaluation) reports average improvement over LAPO of +4.2 success rate and −0.069 MSE for LAOF, and +11.5 success rate and −0.096 MSE for LAOF-Action. LAOF reaches success rates of 82.5 ± 2.3 (SPATIAL), 85.3 ± 1.4 (OBJECT), 87.2 ± 2.2 (GOAL), and 52.0 ± 1.7 (LONG), with MSEs of 0.111, 0.082, 0.118, and 0.088 respectively. LAOF-Action reaches 88.2 ± 1.5, 95.9 ± 1.3, 88.6 ± 1.6, and 63.7 ± 1.9 with MSEs of 0.076, 0.064, 0.081, and 0.068.

• Optical flow improves results on PROCGEN too. Table 2 (discrete latent actions) reports average accuracy improvement of +16.20 and average return improvement of +0.16 for LAOF over LAPO, and +20.73 accuracy and +0.22 return for LAOF-Action. The paper's text describes the PROCGEN gains as increases in normalized episodic rewards of 16% and 22% for the unsupervised and action-supervised settings respectively, and describes LIBERO gains as +4.2% success for unsupervised and +11.5% for action-supervised methods.

• Adding flow to an existing baseline also helps. CoMo w/ OF improves over CoMo by +4.0 success on LIBERO and by +15.18 accuracy / +0.12 return on PROCGEN, showing the constraint is not specific to LAOF's architecture.

• Continuous representations beat discrete ones. Figure 3(a) and 3(b), comparing solid (continuous) with dashed (discrete) curves in both the action-free and 1% action-ratio settings, show continuous representations outperform discrete across all downstream tasks.

• The evaluation metric tracks downstream performance. Averaged across all tasks, the mean Pearson correlation coefficient between the latent action evaluation metric and downstream performance is 0.8288 for PROCGEN and −0.7311 for LIBERO.

• Gains persist up to a 10% action ratio. The paper recommends applying optical flow constraints when the dataset contains less than 10% action-labeled data. Beyond that, positive effects diminish, and at a 100% action ratio performance degradation is observed in BIGFISH and CHASER, attributed to noise in the optical flow pseudo-labels.

• Flow without any action labels can rival 1%-label supervision. The paper reports that under the extreme 1% action ratio, LAOF without any action supervision matches or even outperforms the action-supervised LAOM-Action.

• Training stability improves. The original LAPO shows large standard deviations in CHASER and HEIST, and LAOM-Action shows instability and overfitting in LEAPER and CHASER at a 1% action ratio; LAOF-Action shows faster convergence, smaller variance, and higher final action accuracy across all tasks.

• Temporal-difference inputs are not a substitute for motion. Replacing the IDM's future-frame input with inter-frame differences, as in CoMo, does not consistently improve performance and degrades it in SPATIAL and GOAL, because feature differences lack pixel-level direction, magnitude, and locality.

• Architecture hierarchy. Averaged over all LIBERO and PROCGEN tasks (Table 3), the ranking is LAOF (LIBERO MSE −0.069, Succ +4.2; PROCGEN Acc +16.20, Return +0.16) > LAOF-AE (−0.060, +3.9, +13.31, +0.14) > LAOF-Only($z_t$) (−0.053, +3.6, +13.2, +0.14) > LAOF-FlowFDM (−0.046, +3.3, +11.95, +0.13) > LAOF-Only($z_t, s_t$) (−0.012, +0.8, +6.26, +0.03).

Methodology in Plain English

The pipeline has three stages. Pre-training uses unlabeled video: an encoder turns consecutive observations into features, an IDM (a spatial-temporal transformer) infers a latent action from the pair of frames, an FDM (a spatial transformer) predicts the next state from the current state plus that latent action, and a new flow decoder predicts the optical flow that should have occurred. Training minimizes next-state reconstruction error plus the error between predicted and pseudo-labeled optical flow. The optical flow pseudo-labels come from RAFT; where distractors move, LangSAM masks the agent first. Flow vectors are turned into RGB images (direction to Hue, magnitude to Saturation and Value, with magnitude normalized by $\min(1.0, m/(\sigma\sqrt{H^2+W^2}))$) so DINOv2 can encode them.

Distillation trains a policy $\pi$ by behavior cloning on a dataset where the frozen pre-trained IDM has annotated each observation pair with its latent action, so the policy can infer latent actions from the current state

Authors’ abstract

Learning latent actions from large-scale videos is crucial for the pre-training of scalable embodied foundation models, yet existing methods often struggle with action-irrelevant distractors. Although incorporating action supervision can alleviate these distractions, its effectiveness is restricted by the scarcity of available action labels. Optical flow represents pixel-level motion between consecutive frames, naturally suppressing background elements and emphasizing moving objects. Motivated by this, we propose robust Latent Action learning with Optical Flow constraints, called LAOF, a pseudo-supervised framework that leverages the agent's optical flow as an action-driven signal to learn latent action representations robust to distractors. Experimental results show that the latent representations learned by LAOF outperform existing methods on downstream imitation learning and reinforcement learning tasks. This superior performance arises from optical flow constraints, which substantially stabilize training and improve the quality of latent representations under extremely label-scarce conditions, while remaining effective as the proportion of action labels increases to 10 percent. Importantly, even without action supervision, LAOF matches or surpasses action-supervised methods trained with 1 percent of action labels.