CLAM: Continuous Latent Action Models for Robot Learning from Unlabeled Demonstrations

TL;DR

CLAM tackles the data bottleneck in imitation learning by learning continuous latent actions from unlabeled observations and grounding them to real actions with a joint action decoder. It uses a two-stage approach: first, a Latent Action Model (LAM) pretrains via IDM/FDM to relabel transitions with continuous latent actions, then a latent action policy is trained on relabeled data using imitation learning, with an action decoder grounding $z_t$ to $a_t$ using a small labeled dataset. The key finding is that continuous latent actions plus joint grounding yield 2–3x improvements over strong baselines across DMControl, MetaWorld, CALVIN, and a real WidowX robot, even without action-labeled expert data. This demonstrates scalable policy learning from action-less data, with best results achieved by leveraging large unlabeled datasets and non-expert play data for grounding, enabling practical deployment in real-world robotics.

Abstract

Learning robot policies using imitation learning requires collecting large amounts of costly action-labeled expert demonstrations, which fundamentally limits the scale of training data. A promising approach to address this bottleneck is to harness the abundance of unlabeled observations-e.g., from video demonstrations-to learn latent action labels in an unsupervised way. However, we find that existing methods struggle when applied to complex robot tasks requiring fine-grained motions. We design continuous latent action models (CLAM) which incorporate two key ingredients we find necessary for learning to solve complex continuous control tasks from unlabeled observation data: (a) using continuous latent action labels instead of discrete representations, and (b) jointly training an action decoder to ensure that the latent action space can be easily grounded to real actions with relatively few labeled examples. Importantly, the labeled examples can be collected from non-optimal play data, enabling CLAM to learn performant policies without access to any action-labeled expert data. We demonstrate on continuous control benchmarks in DMControl (locomotion) and MetaWorld (manipulation), as well as on a real WidowX robot arm that CLAM significantly outperforms prior state-of-the-art methods, remarkably with a 2-3x improvement in task success rate compared to the best baseline. Videos and code can be found at clamrobot.github.io.

Figures (11)

• Figure 1: Overview of CLAM. CLAM consists of a latent inverse dynamics model, $f_\phi$, which infers the latent action between a transition and latent forward dynamics model, $g_\psi$, which predicts the future observation conditioned on the latent action. CLAM learns a latent action space through the self-supervised objective of future observation reconstruction. Unlike prior work, CLAM produces continuous latent actions. To ensure the learned latent space is amenable to decoding to real-world actions, CLAM jointly trains the action decoder and the latent action model.
• Figure 2: WidowX Robot Arm Setup and Evaluation Tasks. We evaluate the scalability of CLAM using four manipulation tasks (Right) on a WidowX robot arm (Left) in a toy kitchen setup walke2023bridgedata.
• Figure 3: MetaWorld Image-Based Experiments. Task success rate over 50 evaluation rollouts across 3 random seeds using image-based inputs. All methods are trained using the same amount of action-labeled data (100 trajectories) for fair comparison. Since $\cal{D}_{\text{labeled}}$ is not necessarily expert data, all the baselines struggle to learn a performant downstream policy, while MLP-CLAM and Transformer-CLAM perform significantly better, with a $3 \times$ improvement in task success over the best baseline. We denote baselines that use a discrete latent action space with hashed markers. We report results for BC-Expert which trains on the same amount of expert trajectories to represent the ideal performance BC achieves with ground truth action labels.
• Figure 4: CLAM Design Choices(a) There are diminishing returns as we increase the latent action dimension. However, setting it too small limits the model's expressivity. (b) Policy performance continues to scale with $|\cal{D}_{\text{unlabeled-expert}}|$. (c) Latent action policy performance improves with more non-expert$\cal{D}_{\text{labeled}}$ while BC plateaus. (d) CLAM is robust to expertise level of $\cal{D}_{\text{labeled}}$.
• Figure 5: Evaluation environments in simulation. We evaluate our approach on both locomotion tasks from the DMControl benchmark (Hopper and HalfCheetah) and manipulation tasks (Assembly, Bin Picking, Peg Insert Side, and Shelf Place) from the MetaWorld benchmark. We also evaluate in CALVIN with the Close Drawer and Slider Left tasks.