🤖 AI

Imitation from Heterogeneous Demonstrations using Grounded Latent-Action World Models

This paper introduces Grounded Latent-Action World Models (GLAM), a framework that learns a shared, prediction-grounded latent action space to enable robust imitation learning from heterogeneous data sources with varying or missing action labels, significantly outperforming existing baselines in both simulation and real-world manipulation tasks.

Original authors: Tianyou Wang, Anson Lei, Joe Watson, Ingmar Posner

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tianyou Wang, Anson Lei, Joe Watson, Ingmar Posner

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

Imagine you want to teach a robot arm how to pick up a cup and put it in a drawer. The traditional way is to have a human physically guide the robot's arm through the motion hundreds of times. This is slow, expensive, and tedious.

This paper proposes a smarter way: Let the robot learn from a mix of "cheap" practice and "expensive" real training.

Here is the breakdown of their solution, GLAM, using simple analogies.

The Problem: Speaking Different Languages

Imagine you have two types of training data:

The "Gold" Data: Real robot recordings where we know exactly how the motors moved (the action labels). This is rare and hard to get.
The "Silver" Data: Lots of videos from other sources—simulations, videos of humans moving objects, or data from a different robot arm. These are abundant and free, but they don't have motor instructions, and they look different (different cameras, different robot shapes).

If you just mix these two together and try to teach the robot, it gets confused. It's like trying to teach a student by mixing a textbook written in English with a video lecture in Japanese, without a translator. The robot doesn't know that "moving the arm up" in the simulation is the same goal as "moving the arm up" in the real world.

The Solution: The "Universal Translator" (GLAM)

The authors built a system called GLAM (Grounded Latent-Action World Model). Think of GLAM as a Universal Translator that speaks a secret, abstract language that both the "Gold" and "Silver" data understand.

Here is how it works in two steps:

Step 1: Learning the Secret Language (The World Model)

Instead of trying to match the robot's specific motor movements, GLAM focuses on cause and effect.

The Core Idea: If an action (like pushing a block) makes the block slide across the table, that action has a specific "meaning." It doesn't matter if the push came from a human hand, a simulation, or a different robot. If the result on the object is the same, the "meaning" of the action is the same.
The Mechanism: GLAM creates a shared "dictionary" (a latent space).
- It looks at the "Silver" data (videos without motor instructions) and asks: "What invisible action caused this object to move?" It guesses the answer and writes it in the secret language.
- It looks at the "Gold" data (real robot) and asks: "What motor movement caused this?" It translates that into the same secret language.
- The Magic: It forces these two translations to match. It ensures that the "push" in the video and the "push" by the robot mean the exact same thing in this secret language.

Step 2: Teaching the Robot (Behavioral Cloning)

Once the secret dictionary is built, the robot can learn much faster.

The system takes all the "Silver" data (the cheap, abundant videos) and re-labels them using the secret language. Now, the robot has a massive library of instructions, even though it never saw the original motor commands.
It then trains the robot to look at a scene, guess the "secret action" needed, and then translate that secret back into real motor commands to execute the task.

A Creative Analogy: The Dance Instructor

Imagine you want to teach a specific dance move to a student (the robot).

Old Way: You stand next to the student and physically move their limbs 1,000 times. (Expensive, slow).
GLAM Way: You have a few videos of professional dancers (Gold data) and thousands of videos of people dancing in their living rooms, on stage, or in cartoons (Silver data).
- The Silver data doesn't have step-by-step instructions, but you can see what the dancers are doing to the floor and the air.
- GLAM acts like a choreographer who ignores the specific shoes or the room size. Instead, it focuses on the rhythm and the flow.
- It realizes that "stepping forward to the beat" is the same concept whether it's done by a pro in a studio or a kid in a living room.
- It creates a "Rhythm Code" that represents the move.
- Now, the student learns by watching all those videos, understanding the "Rhythm Code," and then applying it to their own body. They learn the dance much faster because they had access to thousands of examples, not just a few.

What They Found

The researchers tested this on five different tasks (like stacking blocks or knocking over a bottle) in both computer simulations and the real world.

The Result: The robot trained with GLAM was significantly better than robots trained only on the expensive "Gold" data.
The Gain: In some difficult tasks, the GLAM approach improved success rates by nearly 50% compared to standard methods.
The "Object Mask" Trick: They found that if they told the AI to focus only on the object being moved (ignoring the background and the robot's body), the learning became even more accurate. It's like telling the student, "Don't worry about the room; just watch the ball."

The Bottom Line

GLAM solves the problem of "data scarcity" by teaching robots to understand the physics of movement rather than just memorizing specific motor commands. It allows robots to learn from cheap, messy, and diverse data sources (like simulations or videos) and apply that knowledge to real-world tasks, making them smarter and faster to train without needing thousands of expensive human demonstrations.

Technical Summary: Imitation from Heterogeneous Demonstrations using Grounded Latent-Action World Models

Problem Statement

Imitation learning has achieved significant success in visuomotor manipulation, yet its scalability is constrained by the high cost and scarcity of high-quality, action-labelled demonstrations on target robots. While abundant heterogeneous data sources exist (e.g., simulations, portable handheld devices, human videos), they present three critical challenges:

Action Space Mismatch: Different sources often possess incompatible action spaces or lack action labels entirely.
Visual and Embodiment Variation: Data sources exhibit substantial differences in camera viewpoints, robot morphologies, and scene contexts that do not reflect the underlying task physics.
Limitations of Existing Solutions: Current approaches to co-train on heterogeneous data rely on heuristic, hand-engineered alignment techniques (e.g., optimal transport losses, empirically tuned mixing ratios, or specific tokenizers) rather than principled representation learning.

The paper posits that a more principled approach is required: actions should be represented based on their physical effect on the environment rather than their source. Two actions driving an object along the same trajectory should share a representation, regardless of whether they originate from a real robot, a simulation, or a handheld device.

Methodology: Grounded Latent-Action World Models (GLAM)

The authors propose GLAM, a framework that unifies heterogeneous data into a shared, source-invariant latent action space using a pair of coupled generative models. The pipeline operates in two stages:

Stage 1: GLAM Pretraining (World Model)

The core innovation is the learning of a shared latent action space ( $z_t$ ) grounded in environment dynamics. This is achieved through two generative models trained jointly:

Heterogeneous Model ( $p(X, O, Z)$ ): Trained on the mixture of target data ( $D_{tar}$ ) and auxiliary data ( $D_{aux}$ ). Since $D_{aux}$ lacks action labels, this model treats actions as latent variables inferred solely from state transitions.
- It utilizes an Inverse Dynamics Model (IDM), $q_\phi(z_t | x_t, x_{t+1})$ , to infer latent actions directly from observed transitions (where $x_t$ is a latent state encoded from observations).
- It uses a shared forward dynamics model, $p_\theta(x_{t+1} | x_t, z_t)$ , to predict future observations.
- Principle: By sharing the IDM and forward dynamics across all data sources, the model enforces that latent actions are defined by how they affect the environment, ensuring source invariance.
Target Model ( $p(S, A, Z)$ ): Trained only on the labelled target data ( $D_{tar}$ ).
- It includes an Action Encoder, $q_\eta(z_t | s_t, a_t)$ , which infers latent actions from executable robot actions and proprioceptive states.
- It includes an action reconstruction term, $\log p_\theta(a_t | z_t)$ , ensuring the latent space remains grounded in the target robot's specific control semantics.

Alignment Mechanism:
To bind these two models, the authors introduce an asymmetric KL divergence constraint ( $C_{KL}$ ). This constraint aligns the IDM posterior (from the heterogeneous model) with the Action Encoder posterior (from the target model) specifically on the target dataset.

The asymmetry (using a stop-gradient on the target encoder) ensures that the IDM absorbs the "privileged" control semantics from the target data without being distorted by the ambiguous signals from the auxiliary data.
The combined objective minimizes the negative Evidence Lower Bounds (ELBOs) of both models plus the alignment constraint, creating a unified latent space where any transition can be relabelled with a shared, control-aware latent action.

Note: The authors also investigate a variant using object segmentation masks as observations instead of raw RGB frames to further sharpen the focus on manipulated object dynamics.

Stage 2: World-Model-Aligned Behavioural Cloning (BC)

Once the GLAM model is pretrained and frozen, it serves as a latent-action anchor:

Relabelling: Every transition in both $D_{tar}$ and $D_{aux}$ is relabelled with the inferred latent action $z_t$ using the IDM posterior mean.
Policy Training: A downstream Behavioural Cloning policy (specifically MIP, a lightweight regression policy) is trained to map observations to these latent action chunks.
Execution: The policy's output (latent actions) is decoded back into executable robot actions using the target model's action decoder.

This approach allows unlabelled auxiliary data to serve as direct supervision for the BC policy, effectively substituting for expensive target-robot teleoperation.

Key Contributions

Principled Representation Learning: The paper shifts the paradigm from heuristic alignment to learning action representations grounded in physical prediction. Actions are aligned if they produce the same environmental transition, regardless of the source.
GLAM Architecture: The introduction of a dual-model generative framework (Heterogeneous + Target) linked by asymmetric KL alignment and shared forward dynamics. This creates a latent space that is both source-invariant (learned from mixed data) and control-aware (grounded in target robot actions).
Data Efficiency: The method demonstrates that heterogeneous, unlabelled data can effectively substitute for target-robot demonstrations, significantly reducing the data requirements for reliable deployment.
Object-Centric Grounding: The exploration of using object masks as inputs to further improve transferability by focusing on the physics of the manipulated object.

Experimental Results

The authors evaluated GLAM on five manipulation tasks: three in the real world (lifting, pick-and-place, knock-down) and two in simulation (2-cube and 3-cube stacking).

Performance Gains: GLAM-aligned policies significantly outperformed Behavioural Cloning baselines (MIP) and prior latent-action methods (CLAM, LAPA).
- Achieved an average +48% improvement in task success rate across tasks with the same data-scarce setting.
- On the complex bimanual "stack-three" task (which typically requires hundreds of demonstrations), GLAM achieved a 72.7% success rate, whereas all baselines failed (≤4%).
Cross-Source Transfer: Qualitative experiments confirmed that latent actions inferred from unlabelled sources (e.g., UMI handheld simulations) could be successfully decoded and executed on the target robot (Kinova) without any action labels, validating the alignment of the latent space.
Data Efficiency: In scaling experiments, GLAM-O (using object masks) matched the performance of a standard BC policy trained on 500 target trajectories using only 100 target trajectories augmented with 400 auxiliary trajectories. This suggests auxiliary data can substitute for target data on a near 1:1 basis in terms of performance gain.
Object Masks: Using object masks as observations consistently improved performance, particularly in complex tasks like "knock-down" (+25%) and "stack-three" (+16%), supporting the hypothesis that object-centric grounding enhances transfer.

Significance and Claims

The paper claims that GLAM provides a paradigm for learning from heterogeneous data that overcomes the data bottleneck in imitation learning. By treating actions as latent variables grounded in environment transitions, the method eliminates the need for hand-crafted alignment recipes.

The authors emphasize that GLAM:

Unifies labelled and unlabelled data into a single, source-invariant latent space.
Enables the direct use of cheap, abundant auxiliary data (like simulations or handheld device data) as supervision for target robots.
Achieves state-of-the-art performance in data-scarce settings without requiring large-scale pretraining or massive online interaction.

The work is presented as a step toward scaling capable policies to the "long tail" of real-world tasks by leveraging diverse, imperfect data sources through a principled, physics-grounded representation learning framework.