Imitation from Heterogeneous Demonstrations using Grounded Latent-Action World Models

Abstract: Imitation learning has emerged as a powerful paradigm for learning visuomotor policies, but its generalisation and stability are limited by the scale and quality of demonstration data needed. A promising direction is to leverage more abundant but heterogeneous data sources, which differ in action space and often lack action labels altogether. Existing co-training approaches that combine heterogeneous data sources rely on heuristic and hand-engineered alignment techniques. In contrast, we argue that action representations should be grounded in prediction: actions that produce the same effect on the environment should share the same representation, regardless of their sources. To this end, we instantiate this principle by using a grounded latent-action world model (GLAM), a pair of generative models with a shared latent action space across data sources that is grounded by predicting future observations consistently across sources. This latent action space is used to train downstream behavioural cloning (BC) policies which map observations to latent actions and decode them back to robot actions, providing a paradigm for learning from heterogeneous data. Empirically, we demonstrate that GLAM successfully learns an aligned latent action space that facilitates action transfer across data sources with and without action labels. Across five manipulation tasks in simulation and in the real world, GLAM-aligned policies significantly outperform BC baselines and prior latent-action methods, achieving an average of +48% improvement in task success rate with the same data-scarce setting. Videos and code are available at https://viccccciv.github.io/glam/.

• The paper presents GLAM, a novel framework that aligns latent actions with environment transitions to enable effective imitation from heterogeneous demonstrations.
• It employs a dual generative model architecture, using an inverse dynamics model for auxiliary data and an action encoder for target robot data to ensure source-invariant, control-aware representations.
• Experimental results demonstrate significant improvements in task success rates, data efficiency, and motion smoothness, significantly reducing the need for expensive target-robot teleoperation.

Grounded Latent-Action World Models for Imitation from Heterogeneous Demonstrations

Motivation and Challenges in Heterogeneous Imitation Learning

The paradigm of imitation learning for visuomotor policies has demonstrated scalability bottlenecks due to the dependence on high-quality action-labelled demonstrations. The collection of such data on target robots remains resource-intensive, motivating research into leveraging more abundant but heterogeneous sources, such as simulated trajectories, portable handheld devices, or human videos. These auxiliary data sources often diverge significantly in embodiment, action space, and lack action labels, rendering naive pooling ineffective. Prior works in co-training (e.g., [14, 15, 16, 17]) utilize empirical alignment heuristics, suboptimal for task generalization. The paper proposes a principled framework where actions are represented and aligned according to their effect on environment transitions, independent of source, transforming the alignment problem into one of representation learning anchored in physical dynamics.

GLAM: Model Architecture and Latent Action Grounding

The Grounded Latent-Action World Model (GLAM) consists of two tightly coupled generative models:

• Heterogeneous Model: This operates across the union of target and auxiliary data, treating actions as latent variables inferred via an inverse dynamics model (IDM) conditioned on environment transitions.
• Target Model: This leverages the explicit action-labelled target robot data, employing an action encoder to infer latent actions directly from state-action pairs.

Both models share a continuous latent action space, with their alignment governed by an asymmetric KL constraint and forward dynamics consistency. Critically, only the target model’s action encoder receives privileged access to ground-truth action labels, grounding latent actions in executable robot control while keeping the IDM source-invariant. The IDM posterior absorbs control semantics only through target transitions, facilitating cross-domain action transfer.

The unified latent space enables relabelling of all trajectories—regardless of origin—with control-aware latent actions. Notably, the optional incorporation of binary object masks as observation input further sharpens the grounding of action representations to manipulated object movement, leading to demonstrable improvements in transferability.

World-Model-Aligned Behavioural Cloning Pipeline

The GLAM framework is operationalized in a two-stage pipeline:

1. Pretraining: GLAM is trained end-to-end to produce source-invariant and control-aware latent actions by minimizing negative ELBO losses of both generative models, supplemented with the posterior alignment constraint.
2. Imitation Learning: Using the pretrained GLAM, all available observation transitions are relabelled with latent actions, serving as supervision for a downstream behaviour cloning (BC) policy. Any BC architecture is compatible; the implementation utilizes a MIP-based BC policy, regressing from visual and proprioceptive inputs to latent actions and decoding them to executable robot actions.

Auxiliary demonstrations without action labels become instantly usable as BC supervision, drastically reducing the reliance on expensive target-robot teleoperation.

Experimental Evaluation and Numerical Results

Experiments span five manipulation tasks, three real-world and two simulated, using a Kinova arm for real tasks and MuJoCo for simulated tasks. Auxiliary data is collected with a floating UMI gripper or simulated trajectories, unlabelled in action space. Each task uses 100 target-robot demonstrations and 400 auxiliary, unlabelled trajectories.

GLAM is compared against BC, CLAM [46], LAPA [38], and their object-masked variants, controlling for architecture and data scale. Key empirical findings include:

• Success Rate Improvement: GLAM(-O) policies demonstrate an average +35% improvement over the best baseline in real-world tasks, +44% on simulated stack-two, and achieve +69% improvement on bimanual stack-three, where other methods fail (<4% success).
• Latent Action Transfer: Qualitative analysis shows latent actions inferred from auxiliary sources can be executed on the target robot, consistently reproducing behaviours in simulation.
• Data Efficiency: Scaling experiments reveal that GLAM-O matches BC performance with only one-fifth of the target data required. Furthermore, increasing auxiliary trajectories yields performance gains parity with increasing target data, validating the utility of aligned latent actions for data efficiency.
• Motion Smoothness: GLAM-O policies achieve superior smoothness metrics and more reliable joint-space convergence compared to BC at equal training scale, corroborating generalization trends at both task and motion level.

Limitations and Implications

GLAM's architecture and alignment strategy enable effective generalization from heterogeneous sources but exhibit several limitations:

• Auxiliary data share deployment camera viewpoints and tabletop scenes; transfer to truly in-the-wild video (with scene and viewpoint drift) is untested.
• The framework has not been validated across morphologically distinct end-effectors (e.g., multi-fingered hands); integration with human-hand demonstrations remains an open direction.
• Non-task-matched auxiliary data scalability, the use of only skill primitives, and the extension to multi-modal supervision (incorporating tactile or force sensing) are promising future avenues.

Practically, GLAM reduces teleoperation costs and enhances generalization in manipulation by transforming auxiliary data into actionable supervision, opening pathways for robust policy learning with sparse target labels. Theoretically, it reframes heterogeneous imitation as a grounded latent-action representation problem, shifting focus from heuristic alignment to principled dynamics-based grounding.

Conclusion

The GLAM framework advances imitation learning by leveraging a grounded latent-action world model to unify demonstrations across heterogeneous sources, with and without action labels. By aligning latent actions through environment transitions and injecting control semantics via action-encoder alignment, GLAM serves as a scalable, source-invariant anchor for downstream behavioural cloning. Empirical results substantiate strong improvements in task success rates, particularly under data-scarce regimes. The framework’s modularity and principled alignment strategy present significant implications for efficient real-world robotic learning, with substantial opportunities for extension into multimodal and morphologically diverse tasks.