💻 computer science

Latent Actions from Factorized Transition Effects under Agent Ambiguity

This paper introduces Observed Transition Factorization (OTF) and its variants, OTF-LAM and OTF-LAM-Dino, which decompose ambiguous observation transitions into reusable primitives to learn robust, action-like latent representations that outperform baselines in complex, distractor-rich environments without supervision.

Original authors: Heejeong Nam, Chandradithya S Jonnalagadda, Harshit Aggarwal, Eric Xu, Randall Balestriero

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Heejeong Nam, Chandradithya S Jonnalagadda, Harshit Aggarwal, Eric Xu, Randall Balestriero

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Noisy Kitchen"

Imagine you are trying to learn how to cook by watching a video of a chef. However, the video is messy. While the chef is chopping onions (the action you care about), the camera is shaking, a dog is running across the floor, and the wind is blowing the curtains.

In the world of AI, this is called Agent Ambiguity.

The Agent: The chef (the thing you want to control).
The Action: Chopping onions.
The Ambiguity: The video shows all the changes at once (onions moving, dog moving, curtains swaying).

Older AI models tried to guess the "chopping" action by looking at the whole messy video. They often got confused, thinking the dog running was part of the cooking recipe, or they couldn't tell the difference between the chef moving and the camera shaking.

The Solution: Breaking the Video into "Lego Bricks"

The authors propose a new method called OTF-LAM. Instead of trying to guess the whole action at once, they break the video changes down into small, reusable pieces, like sorting a pile of mixed-up Lego bricks.

They call these pieces Observed Transition Primitives.

Instead of seeing "Chef chopping," the AI sees: "A small patch of pixels moving left," "An edge shifting down," and "A background blur."
These are the Lego bricks. They are simple, reusable patterns of movement that can happen anywhere, regardless of whether it's a chef, a robot, or a cartoon character.

How It Works: The Two-Step Process

Step 1: The "Brick Sorter" (OTF)
First, the AI watches thousands of videos and learns to sort every tiny change in the picture into a specific "brick" from its vocabulary.

Analogy: Imagine a librarian who doesn't care what the book is about, only how the pages are turning. They sort every page-turn into a category: "Fast flip," "Slow slide," or "Rip."
The AI learns that a "Fast flip" looks the same whether it's a human hand turning a page or a robot arm turning a dial. This makes the "bricks" reusable.

Step 2: The "Chef's Assistant" (OTF-LAM)
Once the bricks are sorted, the AI needs to figure out which bricks were caused by the chef (the agent) and which were caused by the dog or the wind.

Analogy: The AI looks at the current scene (the kitchen) and asks, "Given that the chef is holding a knife, which of these 'moving bricks' are likely the chef's doing?"
It picks the relevant bricks, ignores the noise (the dog), and combines them into a single "Action Signal." This signal tells the AI: "The chef is chopping."

The "Frozen DINO" Trick (OTF-LAM-Dino)

The paper also introduces a smarter version called OTF-LAM-Dino.

Analogy: Imagine trying to predict the next frame of a video. Usually, the AI tries to redraw every single pixel (the color of the onion, the texture of the table). This is hard because the lighting might change or the table might look different.
The Trick: Instead of redrawing the pixels, this version uses a "frozen" expert (DINOv2) that already understands the shape and structure of things. The AI only needs to predict how the structure changes, not the exact colors.
This is like predicting the next move in a chess game by looking at the board positions (structure) rather than trying to paint the wooden pieces (pixels). It makes the AI much better at ignoring distractions.

What Did They Prove?

The authors tested this on two main things:

Moving MNIST (The "Digit" Test): They showed the AI videos of numbers moving. They trained it on numbers 0–4 and then tested it on numbers 5–9 (which it had never seen).
- Result: Because the AI learned the movement patterns (the bricks) rather than the specific look of the numbers, it worked perfectly on the new numbers. It proved the "bricks" are reusable.
DCS (The "Robot" Test): They used a complex robot simulation where the robot has to run or walk while a distracting background moves.
- Result: The new method (OTF-LAM) was better at learning how to control the robot than older methods. It successfully ignored the distracting background and focused on the robot's movement.

The Bottom Line

This paper says: "Don't try to guess the whole action from a messy video. Instead, break the video changes down into small, reusable 'movement bricks.' Sort these bricks first, then figure out which ones belong to the agent you want to control."

By doing this, the AI becomes much better at learning to act in messy, real-world environments where there are lots of distractions.

Technical Summary: Latent Actions from Factorized Transition Effects under Agent Ambiguity

1. Problem Statement

Latent Action Models (LAMs) aim to learn action-like proxies directly from observation transitions (e.g., video frames $x_t$ to $x_{t+1}$ ) without explicit action labels. While effective in controlled settings where the agent is the dominant source of change, LAMs face significant agent ambiguity in complex scenes. In environments with distractors, moving backgrounds, camera dynamics, or articulated objects, the visual transition $\Delta x_t$ is a mixture of the agent's controlled motion and uncontrolled exogenous factors.

Standard LAMs compress this entire transition into a monolithic latent vector. Without supervision, the model lacks a basis to distinguish which observed changes are attributable to the agent versus distractors. This ambiguity leads to latent spaces that entangle controllable motion with environment-specific noise, hindering transferability and downstream policy learning. Existing approaches often rely on object-centric priors or segmentation, which impose rigid assumptions and limit scalability.

2. Methodology

The authors propose Observed Transition Factorization (OTF), a bottom-up approach that treats the visual transition not as a single event to be compressed, but as a sparse composition of reusable observed-transition primitives. The framework consists of two main stages:

2.1. Observed Transition Factorization (OTF)

The first stage learns a shared vocabulary of local motion effects.

Input Construction: A motion observation $o_t$ is constructed by applying a gradient transform (or Sobel) to consecutive frames and computing the difference. This emphasizes observable change while reducing dependence on static appearance and object identity.
Patchwise Vector Quantization: The motion input is partitioned into spatial patches. Each patch is encoded and quantized into a discrete codebook of size $K$ using a VQ-VAE style objective.
Primitives: The codebook learns reusable observed-transition primitives (e.g., local displacement, edge shifts, rotation-like residuals, contact deformation). Crucially, these primitives represent visual effects rather than pure physical motion or object identities.
Factor Representation: For each transition, the model outputs a structured set of factors $E_t = \{e_{t,k}\}_{k=1}^K$ , where each factor $e_{t,k}$ consists of the codebook entry, an occupancy map (spatial support), and an activation strength.
Training Objective: The factorizer is trained to reconstruct the motion observation $o_t$ using the selected primitives and the current frame $x_t$ as a reference for appearance, minimizing reconstruction loss alongside standard VQ losses and an orthogonality regularizer to ensure codebook diversity.

2.2. OTF-LAM: Action Abstraction

Once the OTF module is pretrained, it is frozen, and a Latent Action Model is trained on top of the extracted factors.

State-Aware Abstraction: The model takes the factors $E_t$ and the current state $x_t$ to form a compact latent action $z_t^{act}$ .
Gating Mechanism: An "Action Abstraction Module" converts each factor into a state-aware token. A gating network assigns a relevance weight $\alpha_{t,k}$ to each factor based on the current context, effectively selecting which observed-transition primitives are relevant to the agent's action.
Prediction: The aggregated latent $z_t^{act}$ is used in a standard inverse-forward dynamics framework to predict the next observation (or a residual update).

2.3. OTF-LAM-Dino

A decoder-free variant is introduced to further separate reusable components from environment-specific details.

Frozen Representation: Instead of predicting pixel space, the model predicts future states in a frozen DINOv2 representation space.
Architecture: Both the visual encoder (DINOv2) and the OTF motion primitive extractor are frozen. The learnable components are limited to the action aggregator and the forward dynamics predictor within the DINO feature space. This forces the model to focus on abstracting transition structure relevant to control rather than reconstructing pixel-level textures.

2.4. Policy Learning

To connect the learned latent space to real-world control, a small amount of action supervision is used. A behavioral cloning policy maps the latent action $z_t^{act}$ to true actions, and an action decoder is trained to map the predicted latent back to ground-truth actions, utilizing the current observation to resolve ambiguity regarding which object the motion applies to.

3. Key Contributions

Observed Transition Factorization (OTF): A novel method to decompose complex visual transitions into a sparse set of reusable, localized primitives before attempting to infer latent actions. This shifts the learning target from "recovering the true action" to "structuring the observed transition effects."
OTF-LAM and OTF-LAM-Dino: Two architectures that utilize the factorized transition interface. OTF-LAM-Dino specifically demonstrates that combining reusable transition vocabularies with frozen visual state representations (DINOv2) can improve downstream performance by acting as a regularizer against pixel-level nuisance factors.
Zero-Shot Transferability: Empirical evidence showing that the learned motion primitives transfer zero-shot across different visual carriers (e.g., different digit classes in Moving MNIST) and embodied morphologies (e.g., from Walker to Cheetah in Distracting Control Suite), indicating that the codebook captures general transition structures rather than environment-specific templates.
Robustness to Ambiguity: The method demonstrates that factorizing transitions allows for more robust latent action learning in the presence of distractors, outperforming or matching monolithic LAM baselines in complex settings.

4. Experimental Results

The paper evaluates the method on the Distracting Control Suite (DCS) (Cheetah-Run and Walker-Run) and Moving MNIST.

Reusability: In cross-morphology and cross-carrier transfer experiments, the OTF vocabulary showed significantly lower performance degradation compared to monolithic VQ-VAE baselines. For instance, in DCS, the OTF vocabulary trained on Walker-Run maintained high reconstruction quality when applied zero-shot to Cheetah-Run, whereas monolithic baselines suffered large error increases (up to 71% degradation vs. ~20% for OTF).
Downstream Policy Learning:
- On DCS, OTF-LAM and OTF-LAM-Dino achieved competitive or superior returns compared to baselines like LAPO, FLAM, and HiLAM.
- OTF-LAM-Dino consistently outperformed the decoder-based OTF-LAM, suggesting that predicting in a frozen representation space is more effective for control tasks than pixel reconstruction.
- The method successfully distilled latent actions into policies that matched or exceeded baselines, even when the visual transition contained significant distractor noise.
Ablation: The study on codebook size ( $K$ ) indicated that while larger vocabularies generally helped in complex environments (Cheetah-Run), the performance was not strictly monotonic, suggesting an optimal capacity dependent on the environment and architecture.

5. Significance and Claims

The paper claims that factorized observed effects serve as a superior intermediate representation for observation-only latent action learning under agent ambiguity.

Shift in Perspective: The authors argue that under observation-only learning, the most reliable target is not the true action itself, but the structure of the observed transition effects. By treating the transition as a composition of reusable primitives, the model can separate controllable motion from distractors without relying on rigid object-centric priors.
Generalization: The results suggest that learning a "vocabulary" of visual transition effects allows for better generalization across different agents and visual domains compared to learning a single, monolithic latent action per transition.
Modest Scope: The authors explicitly state that their work evaluates these approaches on controlled settings and DCS distractor environments. They acknowledge that extending this to richer, more diverse video-game domains (e.g., scrolling cameras, diverse non-agent dynamics) remains future work. They do not claim to solve the general inverse problem of action recovery without any supervision, but rather demonstrate that a factorized intermediate representation significantly improves the robustness and transferability of the learned latent action interface.