Latent Action Learning Requires Supervision in the Presence of Distractors
Authors: Alexander Nikulin
Ilya Zisman
Denis Tarasov
Lyubaykin Nikita
Andrei Polubarov
Igor Kiselev
Vladislav Kurenkov
Venue: ICML 2025 poster (2025)
Keywords: latent action learning
imitation learning
learning from observations
learning from videos
latent action model
Primary area: reinforcement_learning->batchoffline
TL;DR: We empirically investigate the effect of distractors on latent action learning.

Abstract: Recently, latent action learning, pioneered by Latent Action Policies (LAPO), have shown remarkable pre-training efficiency on observation-only data, offering potential for leveraging vast amounts of video available on the web for embodied AI. However, prior work has focused on distractor-free data, where changes between observations are primarily explained by ground-truth actions. Unfortunately, real-world videos contain action-correlated distractors that may hinder latent action learning. Using Distracting Control Suite (DCS) we empirically investigate the effect of distractors on latent action learning and demonstrate that LAPO struggle in such scenario. We propose LAOM, a simple LAPO modification that improves the quality of latent actions by **8x**, as measured by linear probing. Importantly, we show that providing supervision with ground-truth actions, as few as 2.5% of the full dataset, during latent action learning improves downstream performance by **4.2x** on average. Our findings suggest that integrating supervision during Latent Action Models (LAM) training is critical in the presence of distractors, challenging the conventional pipeline of first learning LAM and only then decoding from latent to ground-truth actions.

Latent Action Learning Requires Supervision in the Presence of Distractors
Alexander Nikulin
Ilya Zisman 134 Denis Tarasov Nikita Lyubaykin 15 Andrei Polubarov 134
Igor Kiselev
Vladislav Kurenkoy
Abstract
Recently, latent action learning, pioneered by Latent Action Policies (LAPO)), have shown re- markable pre-training efficiency observation- only data, offering potential for leveraging vast amounts video available on the web for em- bodied AI. However, prior work has focused on distractor-free data, where changes between obser- vations are primarily explained by ground-truu truth actions. Unfortunately, real-world videos contain action-correlated distractors that may hinder la- tent action learning. Using Distracting Control Suite (DCS) we empirically investigate the ef- feet of distractors on latent action learning and demonstrate that LAPO struggle in such scenario. We propose LAOM, a simple LAPO modification that improves the quality of latent actions by 8x, measured by linear probing. Importantly, show that providing supervision with ground-truth actions, as few as 2.5% of the full dataset, dur- ing latent action learning improves downstream performance by 4.2x on average. Our findings suggest that integrating supervision during Latent Action Models (LAM) training is critical in the presence of distractors, challenging the conven- tional pipeline of first learning LAM and only then decoding from latent to ground ground-truth actions.
Lspp-supoiiitionn (ours)
LAOM (ours)
Normalized Retum
0.4
LAPO (original)
0.3
BC
0.2
0.1
Labeled trajectories
Figure I. We show that in the presence of distractors, LAPO strug- gles to Iearn latent actions useful for pre-training and that simple BC IDM are effective. We propose LAOM, simple modification that doubles the performance but still underperforms. Thus, we propose reuse available ground-truth action labels to supervise latent action learning, which significantly improves the performance, achieving normalized score 0.44. It recovers almost half the performance of BC with access to the full action- labeled dataset, while having access only 2.5%. Results are averaged over four environments from Distracting Control Suite, three random seeds each. We provide per-environment plots on Figure 8. See Section 3 for the evaluation protocol, Sections 4 and for method details.
els (LAM) infer latent actions between successive obser- vations, effectively compressing observed changes. Under certain conditions, latent actions ground truth action space (Schmidt & Jiang, 2023; Bruce et al., 2024). After training, LAM can be utilized for imi- tation learning on latent actions to obtain useful behavioral priors. For example, LAPA (Ye et al., 2024) showed that latent action learning can be used to pre-train large model on only human manipulation videos, and despite the huge coosseembodiment gap, still outperform Open (Kim et al., 2024b), which pre- pre-trained expert in-domain data with available action labels.
1. Introduction
rediscover the
Recently, a new wave of approaches based on latent action learning has emerged (Edwards al., 2019), demonstrating superior pre-training efficiency datasets without action labels in large-scale robotics (Ye et al., 2024; Chen et al., 2024b;a; Cui et al., 2024; Bruce et al., 2024) and reinforce- ment learning (Schmidt & Jiang, 2023). Latent Action Mod-
'AIRI 'MITT' "Skolteh 'Research Center for Trusted Artificial Intelligence, ISP RAS IInoopolis University 'Accenture. Corre- spondence to: Alexander Nikulin <nikulin @ net>, Work done by dunnolab.ai.
Despite the initial success and the promise of unlocking vast amounts of video available on the web (Schmidt Jiang, 2023; Ye et al., 2024), there is a critical shortcom- ing of previous work it uses distractor-fer free data, where
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Latent Action Learning Requires Supervision in the Presence of Distractors
chestahhrun
wolker-run
hopper-hop
humansid-walk
Latent Action Model
Latent Action Pretraining
Action Finetuning
C
d2
FDM
BC
#: IDM &
BC
T
IDM
Ot Op-1
Ox
Figure 2. Visualization of the environments from the Distracting Control Suite (DCS) used in our work. Top row: without any distractors, identical to the original DeepMind Control Suite. Bot- tom row: with distractors, which consists of dynamic background videos, agent color change and camera shaking, See Section 3 for additional details.
Figure 3. Overview of the latent action learning pipeline. In the first stage, the Latent Action Model (LAM) is pre-trained to infer latent actions between consecutive observations. In the second stage, the LAM is used to relabel the entire dataset with latent actions, which are then used for behavioral cloning. Finally, decoder is trained to map from latent to true actions using a small number of labelled trajectories. In our work, we do not modify this pipeline in any way: we only examine the LAM architecture itself (see Figure 4).
all changes between observations are mainly and most ef- ficiently explained by ground truth actions only, such robot manipulation on a static background (Khazatsky et al., 2024). Unfortunately, this is not true for real-world web- scale data, as it contains a lot of action-correlated noise (Misra et al., 2024), e.g. people moving in the background. Such noise may better explain video dynamics and thus lead to latent actions unrelated real actions. The phenomenon of overfitting task-irrelevant information is not new and has been studied in model-based (Wang al., 2024) and representation learning (Lamb et al., 2022; Zhang et al., 2020; Zhou et al., 2023). However, the effect of distractors on latent action learning, which we aim to address in this work, has not been similarly investigated.
tion, we show that latent action learning with supervision generalizes better in contrast to approaches based on inverse dynamics models (Baker al., 2022; Zhang et al., 2022a; Zheng et al., 2023) but does not learn control-endogenous minimal state (Lamb et al., 2022).
2. Preliminaries
Learning from observations. Most methods in reinforce-
ment learning require access to the dataset D = 1.W of N trajectories, where each tn == {(o ",""""""" con- tains observations, actions and rewards. Similarly, imitation learning requires access to trajectories Tn = {(o]", )== that contain actions. Unfortunately, most expert demonstra- tions in the real world, such as YouTube videas of some human activity (Aytar et al., 2018; Baker et al., 2022; Zhang et al., 2022a; Ghosh et al., 2023), do not include rewards or action labels. Thus, researchers are actively exploring how to most effectively use the data Tn {o without action labels to accelerate the learning of embodied agents at scale (Torabi et al., 2019). Still, we can often assume that very small number of action labels are available. For exam- ple, previous work has explored ratios of up to 10% (Zheng et al., 2023), whereas in our work we allow a maximum of ~ 2.5% of labeled transitions.
In this work we empirically investigate the effect of action- correlated distractors on latent action learning using Dis- tracting Control Suite (Stone et al., 2021). We demonstrate that naive latent action learning based on quantization and reconstruction objectives, such as LAPO (Schmidt & Jiang, 2023), struggle in the precense of distractors (see Section 4). We propose LAOM, simple LAPO modification that im- prove the quality of latent actions by 8x, as measured by linear probing, and double the downstream performance (see Figure 8). However, even after this, the resulting perfor- mance is only slightly better than simple Behavioral Cloning on available -ruun actions. Thus, as our core con- tribution, we show that providing supervision with a small number, as little as 2.5% of the complete dataset, of action labels during LAOM training improves the downstream per- formance by 4.3x on average (see Section 5), outperforming all baselines (see Figure 8). Our findings suggest that the pipeline used in most current work (Ye et al., 2024; Cui et al., 2024; Chen et al., 2024b) to first learn LAM and only then decode to ground-truth actions is suboptimal when dis- tractors are present, as with supervision better result can be achieved using the same budget of actions labels. In addi-
Latent action learning. Latent action learning approaches (Edwards et al., 2019; Schmidt & Jiang, 2023; Chen et al., 2024b; Cui et al., 2024; Ye et al., 2024) aim to infer latent actions 2, such that they are maximally informative about each observed transition (o,, or+1) 1) while being minimal. After the latent action model (LAM) is pre- pre-trained, we can train policies to imitate latent actions on full data to obtain
2

---

Latent Action Learning Requires Supervision in the Presence of Distractors
81
FDM
Encoder
O
71
8--1
O0
Latent IDM
Latent F FDM
Encoder
Ot-X
Quamnizer
Oct1
a
IDM
Linear
C
8t+1
C
0t+1
LAPO
LAOM
Figure 4. Simplified architecture visualization of LAPO, and LAOM our proposed modification. LAPO consists of IDM and FDM, both with separate encoders, uses latent action quantization and predict next ohservation image space via the decoder in FDM. LAOM incorporates multi step IDM, removes quantization and does not reconstruct images, relying latent temporal consistency loss. Images are encoded by shared encoder, while IDM and FDM operate in compact latent space. When small number of ground ground-truth action labels available, we them for supervision, linearly predicting from latent actions. For detailed description see Section 4.
useful behavioral priors. Finally, small decoder heads can be learned from latent to real actions of domain of interest.
dg from states St and S1+k, where k (1.2.3.... However, as showed by Misra et al. (2024), in the pres- ence of exogenous noise, i.e. non-iid noise that is tempo- rally action-correlated, the sample complexity of learning control-endogenous minimal state from video data can be exponentially worse than from action labeled data. They hypothesized that this is true for latent action learning as well but did not provide any analysis regarding the quality of latent actions, which tried to empirically address in this work.
We base our work on LAPO (Schmidt & Jiang, 2023), which is used in recent work (Chen et al., 2024b; Cui et al., 2024; Ye et al., 2024). LAPO uses two models in combination to infer latent actions. First is inverse dynamics model (IDM), which is given two consecutive observations pre- diets latent action žt PIDM(-lOO, 0t+1). Second is for- ward dynamics model (FDM), which observes current ob- servation and latent action, and predicts the next obser- vation 0i+1 ~ PFDM(-1o0 ()) =)). Both models are trained jointly to minimize the next observation prediction loss Figure 4.
3. Experimental Setup
|00+11 1121 We illustrate the model architecture in
Environments and datasets. To decouple the effects of latent action quality and exploration on performance, we work in offline setting. For our purposes, it is essential that the Behavior Cloning (BC) agent should recover most of the experi performunce when truined on the full dataset with ground-trin truth actions revealed, otherwise it would be difficult to understand the effect of latent action quality on pre-training.
Given the information bottleneck latent actions, e.g. quantization via the VQ-VAE (Van Den Oord et al., 2017), IDM cannot simply copy the next observation into the FDM as is, so it will be forced to compress and encode the differ- ence between observations to be most predictive of the next observation. Without the distractors, through simplicity bias (Shah al., 2020), the latent actions will recover the ground truth actions they most predictive of the dynamics. However, due to the presence of distractors, it may be not true for real-world data. In this work, we empirically exam- ine how well current LAMs can recover true actions in such circumstances.
As currently existing benchmarks with distractors (Stone et al., 2021; Ortiz et al., 2024) new datasets with custom difficulty, based on Distracting Control Suite (DCS) (Stone et al., 2021). DCS uses dy- namic background videos, camera shaking and agent color change as distractors (see Figure for visualization). The complexity determined by the number of videos as well as the scale for the magnitude of the camera and the color change. We empirically found that using 60 videos and scale of 0.1 is the hardest setting when BC can still recover expert performance. We collect datasets with five thousand trajectories for four tasks: cheetah-run, walker-run, hopper- hop and humanoid-walk listed in the order of increasing
yet solved, collect
Control-endogenous minimal state. Lamb et al. (2022) de- fines control-endognouus minimal state as representation that contains all the information necessary to control the agent, while completely discarding all irrelevant informa- tion. Lamb et al. (2022); Levine et al. (2024), theoretically and practically show that to learn such minimal state multi- step IDM should be used, i.c IDM that predicts action
3

---

Latent Action Learning Requires Supervision in the Presence of Distractors
difficulty. See Appendix C for additional details.
Evaluation. To access the quality of the latent actions, we use two methods. First, we follow the approach of Zhang et al. (2022b) and use linear probing (Alain, 2016), which is a common technique used to evaluate the quality of learned representations by training a simple linear classifier or re- gressor on top the representations. Since we include ground truth actions in our datasets for debugging purposes, we train linear probes predict them from latent actions simul- taneously with the main method, e.g. LAPO (Schmidt & Jiang, 2023). We do not pass the gradient through the latent actions, so this does not affect the training. Second, follow- ing the most commonly used three-stagee Stage pipeline (Schmidt & Jiang, 2023; Chen et al., 2024b; Ye et al., 2024), we first pre-train LAM, then train BC model to predict latent actions on the full dataset, and finally, we reveal small number of labeled trajectories to train small two-layer MIP decoder from latent to real actions (see Figure 3). Using this decoder, we then evaluate the resulting agent in the environment for 25 episodes. To access scaling properties with different budgets of real actions, similar to Schmidt & Jiang (2023), we repeat this process for a variable number of labeled tra- jectories, from 2 128. All experiments averaged over three random seeds.
29.17
With distractors
Median action probe mse
Without: distractors
25
22.09
20
16.41
10
5
2.55
LAPO
LAPO without FSQ
Figure 5. Quality of latent actions learned by LAPO. We show that quantization of latent actions significantly reduces the quality of actions, even on data without distractors, where LAPO should work without problems. Removing the quantization recovers the latent action quality, but additional modifications are needed improve LAPO performance with distractors. Results are averaged across all four environments, each with three random seeds.
interested in estimating the upper bound performance of each method in controlled setting, we believe that it is appropriate. For exact hyperparameters see Appendix F.
Baselines. We use BC on true actions as our main baseline, since the main goal of latent action learning is to pre-train useful behavioral policies (Edwards et al., 2019; Schmidt & Jiang, 2023), which can be achieved by recovering true ac- tions as accurately as possible. We use it in two ways. First, we try get the best performance for each full dataset with true actions to use the final return for normalization. With such normalization, we can quantify how much performance we have recovered compared to we had access to fully action-labeled dataset. Second, we train BC from scratch on the same number of labels available to LAM, to evaluate the benefit of pre-training on large unlabeled data. Our last of the most successful and simplest approaches to learn from action-free data at scale (Baker et al., 2022; Zhang et al., 2022a; Zheng et al., 2023). For additional details, see Appendix E.
4. Latent Action Learning Struggle in the Presence of Distractors
To access the effect of distractors on latent action learn- ing we start by carefully reproducing and adapting LAPO (Schmidt & Jiang, 2023) for our domain. We use similar architecture (see Figure 4) with ResNet (He et al., 2016) as observation encoders, borrowed from the open-source official LAPO implementation. Similar to Schmidt & Jiang (2023) we resize observations to 64 height and width, stack- ing 3 consecutive frames.
Quantization hinders latent action Icarning. To vali- date our implementation, we first measured performance on distractor-free datasets, which should not cause any dif- ficulty. Contrary to previous research (Schmidt & Jiang, 2023; Chen al., 2024b; Ye al., 2024; Bruce al., 2024), we found that commonly used latent action quantization during training significantly hindered the resulting latent action quality. We initially hypothesized that the problem might be with the VQ- used for quantizing. In con- versation with Schmidt & Jiang (2023) we confirmed that VQ-VAE is indeed susceptible to codebook collapse and requires extensive tuning. We tried the more modern FSQ (Mentzer et al., 2023), which has already been used suc- cessfully in RL (Scannell et al., 2024) and does not suffer from codebook collapse. Unfortunately, even after tuning, we were unable to improve the results significantly, so we simply removed it. To our surprise, this resulted in large positive improvement (see Figure 5), but only for datasets
baseline is IDM, as it remains
We do not consider other possible types of unsupervised pre-training, as was already extensively explored by other researchers (Tomar et al., 2021: Zhang et al., 2022b; Kim et al., 2024a), even with distractors (Misra et al., 2024; Ortiz et al., 2024). Our aim is not to compare latent action learning with existing approaches, but to investigate whether it works at all in the presence of action-correlated distractors.
On hy byerpaarametess tuning, We tune the hyperparameters based on online performance for BC, on MSE to real actions on the full dataset for IDM, and on final linear probe MSE to real actions for latent action learning. In more practical tasks, we usually do not have this luxury, but since we are

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Medi an action probe msee
0.5
0
5
10
15
20
25
LAPO, without distractors
LAPD with distractors
Normalized Retur
0.4
LAPO, with distractors
LAOM (ours), without distractors
(ours), distractors
0.3
muiti-step IDM
9.54
0.2
larger latent action dim
3,89
0.1
in space
0.0
x8 improvement
augmentations ens
2.74
Labeled trajectories
LAPO without distractors
Figure 7. Performance evaluation of the LAPO and the proposed LAOM with and without distractors. As can be seen, large gap in performance remains in the presence of distractors. Results are averaged across all four environments, each with three random seeds.
modification
1.93
Figure 6. The individual effect of cach proposed change in LAOM, our modification of LAPO, which overall improves latent action quality in the presence of distractors by factor of 8. We describe the proposed changes in detail Section and visualize final architecture in Figure 4. Results are averaged across all four environments, each with three random seeds.
demonstrating the importance of the proposed changes to specifically help latent action learning in the presence of distractors. We visualize the resulting architecture, which we called Latent Action Observation Model (LAOM) in Figure 4 and describe changes in detail next:
without distractors, while with distractors the action quality remained almost the same level.
Multi-step IDM. Inspired by research on control-endogenous minimal state discovery (Lamb et al., 2022; Levine et al., 2024) via multi-step IDM, we slightly modify our IDM ob- jective to estimate latent action 2; from and Ot+k. where k E 21, K}. instead of just consecutive observa- tions. During training, we sampled uniformly for each sample and found that K == 10 worked best. Multi-step IDM helps to learn representation which encodes control- endogenous information with respect current latent ac- tions, which in turn helps learn better latent actions. This simple change alone doubled the latent action quality.
One explanation for the result on Figure 5 may be that we are working with continuous actions, unlike the Schmidt & Jiang (2023) which used discrete actions. However, we believe that there are more general reasons. The main moti- vation for quantizing latent actions was to prevent shortcut learning, i.e. IDM copying 0r+-1 FDM is, and to in- centivize IDM to learn simpler latents that capture only action-related changes. We observed no evidence for short- cut learning, suggesting that it is unlikely to occur with high-dimensional observations, similar to the unlikelihood of collapse in Siamese networks (Chen He, 2021). More action-correlater correlated distractors, the information bottleneck may have the opposite effect, incentivising the IDM encode noise into latent actions. This noise can explain the dynamics more easily, so without guidance, the IDM has no way of distinguishing it from real actions. Therefore, we advise against the use of quantization for LAM training on real-world data.
Increasing latent actions capacity. So far we have used latent actions with 128 dimensions, as in the original LAPO. However, for reasons similar to quantization removal, we significantly increased it 8192, it allows better next- observation prediction. Since IDM cannot distinguish control-related features from noise, the best we can hope for in general is to learn the full dynamics of the environment as accurately as possible. In such a case, latent actions will by definition contain true actions and we will be able to extract them via the probe. This change gives an additional 2.5x improvement.
importantly, in the presence of action-
Latent action quality can be significantly improved. As Figure shows, naive LAPO may be able learn good latent actions in the presence of distractors and further im- provements are needed. Thus, we propose simple modi- fications to the LAPO architecture, which in combination improve latent action quality by 8x, almost closing the gap with distractor-fre free setting (see Figure 6). Interestingly, on distractor-free data improvements are marginal, further
Removing observation reconstruction. The need to fully reconstruct the next observation forces latent actions to en- capsulate changes in each pixel, which is not always related to true actions, e.g. video in the background. Thus, we use the latent temporal consistency loss (Schwarzer et al., 2020; Hansen et al., 2022; Zhao al., 2023) to predict next
5

---

Latent Action Learning Requires Supervision in the Presence of Distractors
LAUM-supervision (ours)
LAOM (ours)
LAPO (original)
IDM
BC
cheetah-run
walker-run
hopper-hop
humanoid-wwalk
0.2
Bamalbee Ratur
0.6
Nomaltzed Raturn
Nommllzzed Return
Noralize aenallteed Return
0.5
0.3
0.2
0.1
0.3
0.2
0.1
0.1
0.0
0.0
32 Labeled trajectories
8
16
64
128
2
8
16
32
64
128
16 32 Labeled trajectories
8
64
128
16
32
64
12a
Labeled trajectories
Labeled trajectories
Figure 8. Performance evaluation of latent action learning approaches and baselines across different budgets of ground-truth action labels. As be seen, LAPO struggles in the presence of distractors, being outperformed by simpler baselines. LAOM, our modification of LAPO. performs better, but not significantly. However, when we reuse the same labels used for decoding from latent to true actions to provide supervision during LAOM training (see Section 5), we significantly improve downstream performance, outperforming basclines in all environments. Importantly, all methods were pre-trained on the same unlabeled datasets and had access to exactly the same action labels, differing only in their use. Results are averaged over three random seeds. detailed description the evaluation pipeline, see Section 3.
ICN
LAOH
supprrttion
LAOM
LAOH isugrrion
tormalized Retur
10
16. 6%.2% decrease
Action probe
Action
63. 1 % decrease
128
0192
4208
2045
1024
312
B192
4006
2043
1024
312
Labeled trajectories
Latent action dim
Latent action dim
(a) Generalization to new distractors
(b) Latent actions quality vs dimensionality
(c) Performance vs actions dimensionality
Figure 9. (a) We show that latent action learning with supervision generalizes better than IDM to novel distractors for all considered budgets of ground-truth action labels available for pre-training. (b)-(c) Supervision with a small number of ground-truth actions during latent action learning allows for smaller action dimensionality without major performance degradation. Without supervision, the quality of latent actions, as well performance, quickly degrades.
observation in compact latent space without reconstruction. IDM and FDM now operate on latent representation and consist of MLPs instead of ResNets (see Figure 4). This brings additional benefits, as with such architecture we can get rid of expensive decoder, reducing model size and in- creasing training speed. For target next observation we use simple stop-grad as in Chen & He (2021) or EMA encoder (Schwarzer et al., 2020). These change alone slightly in- creases probe MSE due to the instabilities. We fix them with the next change.
in later stages.
The large gap in downstream performance remains. As Figure shows, our improvements partially transfer to down- stream performance, LAOM outperforms vanilla LAPO on all label budgets, improving performance by up to 2x. LAOM also outperforms LAPO on data without distractors, but not significantly. However, there remains a large gap in final performance with and without distractors. We should emphasize that this gap is not due the fact that setting with distractors is more difficult for BC, for example. We normalize performance by the return achieved by BC trained on each full dataset with ground-truti actions. Thus, the difference in performance relative to BC and is explained by a difference in the quality of the latent actions.
Adding augmentations. Augmentations are commonly used in conjunction with self-supervised objectives to stabilize training and avoid collapse (Schwarzer et al., 2020; Hansen et al., 2022; Zhao et al., 2023). Similarly, we found that augmentations help with stability and improve performance to even smaller probe MSE. We use the subset of augmen- tations from Almuzairee et al. (2024), which consists of random shifts, rotations and changes of perspective. We apply then only during latent actions training and do not use
Unfortunately, linear probing has a major limitation can only tell us whether real actions are contained in latent ac- tions not. For example, by increasing the dimensionality of latent actions in LAOM, we have improved the quality
6

---

Latent Action Learning Requires Supervision in the Presence of Distractors
according the probe, but sacrificed their minimality, i.e. they additionally describe full dynamics, that is mostly un- related to real actions. This can be detrimental as, during the BC stage, not only do we waste capacity predicting ac- tions with higher dimensionality, but we also risk learning spurious correlations. This is probably the main reason for the poor performance, but is the best we can do, otherwise latent actions will not contain true actions at all.
-upM-superviiion (ours)
LAOM (ours)
Normalized Return
LAPO
BC
0.1
5. Latent Action Learning Requires Supervision
16
64
In previous sections, we proposed LAOM, an improved version of LAPO which almost doubled the downstream performance in the presence of distractors for all budgets of true action labels considered. However, overall performance remained quite low. Similar to unlikelihood of recovering the control- endogenous minimal state in the presence of distractors (Misra et al., 2024), our results suggest that with- out any supervision latent action learning may not be able to learn actions useful for efficient pre- pre-training. What if we can provide supervision? Even the smallest number of true actions may ground latent action learning to focus on control-related features. We explore this in the following experiments.
Labeled trajectories
Figure 10. Evaluation of latent action learning approaches in cross- embodied pre-training in the presence of distractors, c.g. pre- training LAM on datascts from three environments and fune-tuning action labeled data from the remaining one. Supervision dur- ing latent action pre-training improves downstream performance. However, overall performance is comparable to that of a simple BC trained from scratch on available action labels. Results are averaged across all four environments, each with three random seeds.
ure 9a measuring action prediction accuracy on evaluation dataset with never seen distractor background videos. IDMM indeed generalizes worse than LAOM+supervision.
Supervision significantly increases downstream perfor- mance. Despite the fact that existing approaches (Schmidt & Jiang. 2023; Ye et al., 2024; Chen et al., 2024b) pre-train still need to have
Supervision enables compact latent actions without large performance degradation. As we mentioned earlier very high dimensional latent actions are not optimal, as they may not be minimal, i.c. contain formation and require larger BC models to imitate accu- rately. Similarly, LAPA (Ye et al., 2024) also reported that more compact latent action space increases pre-training efficiency. Unfortunately, the effectiveness of LAPO and even LAOM degrades dramatically when the dimensionality of latent actions is reduced. In Figure 9b and Figure 9c we show that supervision can partially mitigate this effect. LAOM+supervision loses only 16% of performance when reducing latent actions dimensionality from 8192 to 64, compared to 63% loss for LAOM. We used 128 labeled trajectories for this experiment.
LAM without true actions, in practice some number of labels to learn the action decoder as last stage. We reuse these labels to provide supervision by lin- carly predicting them from latent actions during LAOM training (see Figure 4 for the final architecture). We plot the resulting downstream performance for each environment in Figure 8 and summarize in Figure 1. As can be seen, LAOM+supervisson outperforms all baselines and scales better with larger budget of real actions. It achieves an average normalized score of 0,44, i.e. it recovers almost half the performance of BC with access to the full dataset of true actions, while using only 2.5% of the action labels. Importantly, all methods have access to exactly the same number of action labels, differing only in how they use them. We provide results for distractor-fret data in Appendix A.
undelltddd in-
Supervision improves cross-embodied pre-training, So far we have used homogeneous datasets, which contain data from only one environment. However, in practice our hope is to pre-train LAM on large and diverse dataset from differ- ent embodiments, including humans (McCarthy et al.. 2024; Ye et al., 2024). To access performance in such a scenario, assemble cross-embodied datasets in leave-one-one fashion, e.g. for the cheetah-run, we sample 1666 trajec- tories (to get ~ 5k) from other environments and combine them into a single dataset. We pre-train LAM and BC on them as usual and use the labeled data from the excluded en- vironment for action decoding or supervision during LAOM
Latent action learning with supervision generalizes bet- ter that IDM. Learning to predict true actions with IDM with a small number of labels and then relabeling larger datasets has recently been a quite successful approach (Baker et al., 2022; Zheng et al., 2023). Unfortunately, IDM greatly limited in its generalization capabilites as dataset with labels may not contain some distractors or cover all actions. L.OM*superviiion on other hand pre-trains on full combined dataset and can adapt better to larger variety of distractors and actions. We confirm this intuition in Fig-
7

---

Latent Action Learning Requires Supervision in the Presence of Distractors
large dataset (Torabi al., 2018). Baker et al. (2022) showed that this approach can work on scale, achiev- ing great success in Minecraft (Kanervisto et al., 2022), Zhang et al. (2022a) used similar pipeline, unlocking hours of in-the-wild driving videos for pretraining. Schmeckpeper et al. (2020) used unlabeled human manipulation videos within online RL loop, which supplied labels to IDM for re-labeling. Zheng et al. (2023) conducted large scale anal- ysis of IDM re-labeling in offline RL setup, showing that only 10% of suboptimal trajectories with labels is enough to match performance on fully labeled dataset.
True
LAOM
IDM
In contrast to previous work (Schmeckpeper et al., 2020; Baker et al., 2022; Zheng et al., 2023), we show that while IDM is a strong baseline in setups without distractors (see Figure 15 Appendix A), generalizes poorly when dis- tractors are present. Our results show that when small number of action labels are available, it is much better to combine IDM and latent action learning to achieve much stronger performance and generalization (see Figure 8), sug- gesting that for web-scale data (Baker et al., 2022; Zhang et al., 2022a) our approach may be better than simple IDM re-labeling.
Figure 11. In contrast to IDM. latent action learning encode of control-unrelated information, such background videos, into the observation representations. This finding suggest that using latent action learning exclusively sentations is not viable in the presence of distractors. We visualize the representations by training a separate decoder to reconstruct original observations.
way to pre-train visual repre-
Latent action learning. To our knowledge, Edwards et al. (2019) was the first to propose the task of recovering latent actions and imitating latent policies from observation, with limited success on simple problems. However, the origi- nal objective had scalability issues (Struckmeier & Kyrki, 2023). LAPO (Schmidt & Jiang. 2023) greatly simplified the approach, removed scalability barriers, and for the first time achieved high success on the hard, procedurally gen- erated ProeGen benchmark (Cobbe et al., 2020). Latent action learning was further scaled by Bruce et al. (2024); Cui et al. (2024); Ye et al. (2024); Chen et al. (2024b;a) to larger models, data, and harder, more diverse robotics domains.
training. As Figure 10 shows supervision during LAM pre- training yields a large performance improvement. However, the final performance is better than training BC only the provided labels from scratch. This is slightly concerning and further emphasizes the limitations of LAM methods in the presence of distractors.
In contrast to IDM, latent action learning does not learn minimal state. DynaMo (Cui et al., 2024) used latent action learning only as an objective to pre-train visual represen- tations, not to obtain useful latent actions. To access the viability of such approach, we additionally train decoders to reconstruct original observations from the representations learned by LAM and IDM. What information does LAM en- codes into its representations? As Figure 11 shows decoders were able to reconstruct original observations quite well, indicating that both LAOM and Luprvsiueviiioon encode information, including distractors.
In contrast to our work, all the mentioned approaches (Schmidt & Jiang, 2023; Ye et al., 2024; Cui et al., 2024; Chen et al., 2024b;a) use data without distractors, where all changes in dynamics are mainly explained by ground truth actions only. As we show in our work (see Section 4), naive latent action learning does not work in the presence of distractors. Although we propose improvements that double the performance, is not enough (see Figure 7). Providing supervision with a small number of action labels during LAM training significantly improves performance (see Fig- ure 1), suggesting that the pipeline used in most current work (Ye et al., 2024; Cui et al., 2024; Chen et al., 2024b;a) to first learn LAM and only then decode to ground-truth actions is suboptimal.
lot of In contrast, multi-step IDM truly learns control-endogenous minimal state as predicted by Lamb et al. (2022); Islam et al. (2022); Levine et al. (2024), fully ignoring underel- information, such as background videos or agent color. This result appears to provide compelling evidence that using LAM exclusively as a way to obtain visual representations not a viable approach in the presence of distractors.
6. Related Work
The most closely related to us is the work of Cui al. (2024), which also removes latent action quantization, the reconstruction objective in favor of latent temporal con-
Action relabeling with inverse dynamics models. Sim- plest approach to utilize unlabeled data to pretrain IDM on small number of action labels further re-label much
8

---

Latent Action Learning Requires Supervision in the Presence of Distractors
sistency (Schwarzer et al., 2020; Zhao al., 2023), and provides ablation with ground-truth actions supervision dur- ing LAM training. However, they train LAM only way to pre-train visual representations and do not provide any analysis regarding the effect of their proposed changes on the quality of the resulting latent actions. This also explains why they report that supervision with true actions gives no improvement, while we show that gives significant gains (see Figure 1). Moreover, visually reconstructing rep- resentations, we show that latent action learning methods do not produce control-endogenou state (see Figure 1 11), and thus are probably not suitable as a method of visual representation learning in the presence of distractors.
8. Conclusion
In this work, we empirically investigated the effect of action- correlated distractors on latent action learning. We showed that LAPO struggles to learn latent actions useful for pre- training. Although we proposed LAOM, simple modi- fication of LAPO, which doubled performance, it did not fully close the gap with the distractor-frc setting, Crucially, we found that even minimal supervision reusing as little as 2.5% of the dataset's ground-truth action labels during latent action learning significantly improved downstream performance, challenging the conventional pipeline of first pre tre-training LAM and only then decoding from latent to real actions. Our findings suggest that integrating supervision is essential for robust latent action learning in real-world scenarios, paving the way for unlocking the vast amounts of video data available on the web for embodied AL. We discuss the limitations of our work the Section 7.
7. Limitations
There are several notable limitations to our work. First, al- though we used the Distracting Control Suite (Stone et al., 2021), which allows us to precisely control the difficulty of distractors in a convenient way and clearly access gen- cralization to new distractors, the overall distribution and noise patterns may be quite different compared to real-world videos on the web. Thus, our conclusions may not be fully applicable, e-g. it is possible that supervision is not im- portant for relevant to embodied AI data, or vice versa, it may turn out to be much more necessary for good results than we have used. Nevertheless, we believe that the overall conclusion about the need for some form of supervision is quite general.
A Acknowledgements
This work was supported by the Ministry of Economic De- velopment of the RF (code 25-139-66879-1-0003)
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Second, the need for supervision for latent action learning is serious limitation, as compared to our setup, which is more reminiscent of Minecraft (Kanervisto et al., 2022) or Nethack (Hambro et al., 2022). where both labeled and unlabeled data are available, we have no chance to get real labels for already existing videos on the web or to fully cover their diversity with hand-crafted labels. Therefore, further research is needed to find out whether pre-training LAM on web data combined with supervision on robot data will achieve a similar effect, although our preliminary ex- periment on cross-embodied pre-training is pessimistic. It is quite possible that supervision can come in other forms than ground -ttntd actions, as we simply need a way to ground latent actions on control-related features of the observations. For example, for egocentric videos (Grauman et al., 2022) we can use hand tracking as proxy action supervise latent action learning.
References
Alain, G. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644, 2016.
Almuzairee, A., Hansen, N., and Christensen, H. I. A recipe for unbounded data augmentation in visual reinforcement learning. arXiv preprint arXiv:2405.1 /7416, 2024.
Aytar, Y., Pfaff, T., Budden, D., Paine, T., Wang. Z., and De Freitas, N. Playing hard exploration games by watch- ing youtube. Advances in neural information processing systems, 31, 2018.
Baker, B., Akkaya, I., Zhokov, P., Huizinga, J., Tang, J., Ecoffet, A., Houghton, B., Sampedro. R., and Clune, J. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639- 2022.
Finally, similar to offline RL (Levine et al., 2020), the prob- lem of hyperparameter tuning remains, since without action labels there currently no way to access the quality of latent actions.
Batra, S. and Sukhatme, G. S. Zero-shot generalization of vision-based rl without data augmentation. arXiv preprint arXiv:2410.0744., 2024.
Bertoin, D., Zouitine, A., Zouitine, M., and Rachelson, E. Look where you look! saliency-guide q qentworrk for

---

Latent Action Learning Requires Supervision in the Presence of Distractors
generalization in visual reinforcement learning. Advances in Neural Information Processing Systems, 35:30693- 30706, 2022.
Ghosh, D., Bhateja, C. A., and Levine, S. Reinforcement learning from passive data via latent intentions. In Inter- national Conference Machine Learning, pp. 11321- 11339. PMLR, 2023.
Bhatt, A., Palenicek, D., Belousov, B., Argus, M., Ami- ranashvili, A., Brox, T., and Peters, J. Crossq: Batch normalization in deep reinforcement learning for greater sample efficiency and simplicity, arXiv preprint arXiv:1902.05500 000000 2019.
Grauman, K., Westbury, A., Byrne, E., Chavis, Z., Furnari, A., Girdhar, R., Hamburger, J., Jiang, H., Liu, M., Liu, X., et al. Ego4d: Around the world in 3,000 hours of egocentric video. In Proceedings of the IEEF/CVE Con- ference on Computer Vision and Pattern Recognition, pp. 18995-19012, 2022.
Bruce, J., Dennis, M. D., Edwards, A., Parker-Hou Holder,, J., Shi, Y., Hughes, E., Lai, M., Mavalankar, A., Steiger- wald, R., Apps, C., et al. Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning, 2024.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018
Burchi, M. and Timofte, R. Mudreamer: Learning predictive world models without reconstruction. arXiv preprint arXiv:2405. 15083, 2024.
Hambro, E., Raileanu, R., Rothermel, D., Mella, V., Rocktäschel, T., Küttler, H., and Murray, N. Dungeons and data: large-scale nethack dataset. Advances in Neu- ral Information Processing Systems, 35:24864-24878. 2022.
Chen, X. and He, K. Exploring simple siamese represen- tation learning. In Proceedings of the IEEF/CVF con- ference on computer vision and pattern recognition, pp. 15750-15758, 2021.
Hansen, N. and Wang, X. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE Inter- national Conference on Robotics and Automation (ICRA), pp. 1361 1-36177 IEEE, 2021.
Chen, X., Guo, J., He, T., Zhang, C., Zhang, P., Yang. D. C., Zhao, L., and Bian, J. Igor: Image-goal representations are the atomic control units for foundation models in embodied ai. arXiv preprint arXiv:2411.00785 2024a.
Hansen, N., Su, H., and Wang. X. Stabilizing deep q- learning with convnets and vision transformers under data augmentation. Advances in neural information pro- cessing systems, 14:3680-36933 2021.
Chen, Y., Ge, Y. Li, Y., Ge, Y., Ding, M., Shan, Y., and Liu, X. Moto: Latent motion token as the bridging language for robot manipulation. arXiv preprint arXiv:2412.0444 2024b.
Hansen, N., Wang, X., and Su. H. Temporal difference learning for model predictive control. arXiv preprint M:Xiv:203304953 2022.
Cobbe, K., Hesse, C., Hilton, J., and Schulman, J. Lever- aging procedural generation to benchmark reinforcement learning. In International conference on machine learn- ing. pp. 2048-2056 2056. PMLR, 2020.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual leam- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Cui, Z. J., Pan, H., Iyer, A., Haldar, S., and Pinto, Dy- namo: In-domain dynamies pretraining for visuo-motor control. arXiv preprint arXiv:2409.12192 2024.
Huang, S., Dossa, R. J., Ye, C., Braga, J., Chakraborty, D., Mehta, K., and Araãšjo, J. G. Cleanrl: High-quality single-file implementations deep reinforcement learn- ing algorithms. Journal of Machine Learning Research, 23(274):1-18, 2022a.
Deng, F., Jang, I., and Ahn, S. Dreamerpro: Reconstruction- free model-based reinforcement learning with prototypi- cal representations. In International conference on ma- chine learning, pp. 4956-4975. PMLR, 2022.
Edwards, A., Sahni, H., Schroecker, Y., and Isbell, C. Imi- tating latent policies from observation. In International conference machine learning, pp. 1755-1763, PMLR, 2019.
Huang. Y., Peng, P., Zhao, Y., Chen, G., and Tian, Y. Spec- trum random masking for generalization in image-based reinforcement learning. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.). Advances in Neural Information Processing Systems, volume 35, pp. 20393-20444 Curran Associates, Inc., 2022b.
Fu, X., Yang, G., Agrawal, P., and Jaakkola, T. Learning task informed abstractions. In International Conference on Machine Learning, pp. 3480-3491. PMLR, 2021.
10

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Isla,, R., Tomar, M., Lamb, A., Efroni, Y., Zang, H., Di- dolkar, A., Mísra, D., Li, X., Van Seijen, H., Combes, R. T. d., et al. Agent-controller representations: Principled offline rl with rich exogenous information. arXiv preprint arXiv:221 10022221000000 1.0164, 2022.
Ma, G., Wang, Z., Yuan, Z.., Wang, X., Yuan, B., and Tao, D. A comprehensive survey of data augmentation in visual reinforcement learning, arXiv preprint arXiv:2210.0456. 2022.
McCarthy, R., Tan, D. C., Schmidt, D., Acero, F., Herr, N., Du, Y. Thuruthel, TT G., and Li, Z. Towards general- ist robot learning from internet video: A survey, arXiv preprint arXiv:2404. 19664, 2024.
Kanervisto, A., Milani, S., Ramanauskas, K., Topin, N., Lin, Z., Li, J., Shi, J., Ye, D., Fu, Q.. Yang, W. Hong, W., Huang, Z., Chen, H., Zeng, G., Lin, Y., Micheli, V., Alonso, E., Fleuret, F., Nikulin, A., Belousov, Y., Svidchenko, O., and Shpilman. A. Minerl diamond 2021 competition: Overview, results, and lessons learned. In Kiela, D., Ciccone, M., and Caputo, B. (eds.), Pro- ceedings of the NeurlPS 202 Competitions and Demon- strations Track, volume 176 of Proceedings of Ma- chine Learning Research, pp. 13-28. PMLR, 06-14 Dec 2022. htt.ttps: httptttpss////coceeeeee ngs. iag.mngs... mlr., wallss/rss////.. #1766kaaervisst2222 html.
Mentzer, F., Minnen, D., Agustsson, E., and Tschannen, M. Finite scalar quantization: Vq-vae made simple. arXiv preprint
2023.
Misra, D., Saran, A., Xie, T., Lamb, A., and Langford, To- wards principled representation learning from videos for reinforcement learning, arXiv preprint arXiv:2403. 13765, 2024.
Ni, T., Eysenbach, B., Seyedsalehi, E., Ma, M., Gehring, C., Mahajan, A., and Bacon, P.-... Bridging state and history representations: Understanding self-predictive rl. arXiv preprint arXiv:2401.08892 2024.
Khazatsky, A., Pertsch, K., Nair, S., Balakrishna, A., Dasari, S., Karamcheti, S., Nasiriany, S., Srirama, M. K., Chen, L. Y., Ellis, K., et al. Droid: A large-scale an-the-wild robot manipulation dataset. arXiv preprint arXiv:2403. 12945, 2024.
Okada, M. and Taniguchi, T. Dreaming: Model-based re- inforcement learning by latent imagination without re- construction. In 2021 ieee international conference on robotics and automation (icra), pp. 4209 4215. IEEE, 2021.
Kim, D., Lee, H., Lee, K., Hwang, D., and Choo, J. In- vestigating pre-trate pre-training objectives for generalization in vision- reinforcement learning. arXiv preprint 240.022406.06037 2024a.
Ortiz, J., Dedieu, A., Lehrach, W., Guntupalli, S., Wen- delken, C., Humayun, A., Zhou, G., Swaminathan, S., Lázaro-Grediilla, M., and Murphy, K. Dme-vh: A bench- mark for representation learning for control with visual distractors. arXiv preprint arXiv:2409, 18330, 2024.
Kim, M. J., Pertsch, K., Karamcheti, S., Xiao, T., Balakr- ishna, A., Nair, S., Rafailov, R., Foster, E., Lam, G., San- keti, P., et al. Openvla: An open-source vision-language- action model. arXiv preprint arXiv:2406.09246. 2024b.
Lamb, A., Islam, R., Efroni, Y., Didolkar, A., Misra, D., Fos- ter, D., Molu, L., Chari, R., Krishnamurthy, A., and Lang- ford, Guaranteed discovery of control-ente -entooe-n la- tent states with multi-step inverse models. arXiv preprint 2oXiv:2207.08229, 2022.
Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., and Dormann, N. baseline::: Reliable rein- forcement learning implementations. Jourmal of Machine Learning Research, 22(268):1-8, 2021. URL http: html.
//jmlr wwtppppppgp/p/pppooootppteeereet
Levine, A., Stone, P., and Zhang, A. Multistep inverse is not all you need, arXiv preprint arXiv:2403. 1940, 2024.
Scannell, A., Kujampää,, K., Zhao, Y., Nakhaei. M., Solin, A., and Pajarinen, iqrl-implicitly quantized representa- tions for sample-efficient reinforcement learning. arXiv preprint arXiv:2406.02696 2024.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.016433 2020.
Schmeckpeper, K., Rybkin, O., Daniilidis, K., Levine, S., and Finn, C. Reinforcement learning with videos: Combining offline observations with interaction. arXiv preprint arXiv: 1.06507 2020.
Liu, Q., Zhou, Q., Yang, R., and Wang, J. Robust represen- tation learning by clustering with bisimulation metrics for visual reinforcement learning with distractions. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 37, pp. 9843-8851, 2023a.
Schmidt, D. and Jiang, M. Learning to act without actions. arXiv preprint 1rXiv:2312.10812 2023.
Liu, Y., Huang, B., Zhu, Z., Tian, H., Gong, M., Yu, Y., and Zhang, K. Learning world models with identifiable factorization, Advances in Neural Information Processing Systems, 36:31831 06:31831-31888 831-41664 2023b.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and
Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv: 177.02770. 2017.

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P. Data-efficient reinforcement learn- ing with self-predictive representations. arXiv preprint arXiv:2007.05929, 2020.
Ye, S., Jang, J., Jeon, B., Joo, S., Yang, J., Peng, B., Mandlekar, A., Tan, R., Chao, Y.-W., Lin, B. Y., al. Latent action pretraining from videos. arXiv preprint arXiv:2410. 11758, 2024.
Shah, H., Tamuly, K., Raghunathan, A., Jain, P., and Netra- palli, P. The pitfalls of simplicity bias neural networks. Advances in Neural Information Processing Systems, 33: 9573-9585, 2020.
Zhang, A., McAllister, R., Calandra, R., Gal, Y., and Levine, S. Learning invariant representations for rein- forcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.
Stone, A., Ramirez, O., Konolige, K., and Jonschkowski, R. The distracting control suite-a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722, 2021.
Zhang, Q., Peng, Z., and Zhou, B. Learning drive by watching youtube videos: Action-conditioned contrastive policy pretraining. In European Conference Computer Vision, pp. 111-128. Springer, 2022a.
Struckmeier, O. and Kyrki, V. Preventing mode col- lapse when imitating latent policies from observations, 2023. URL https: http//ppreeeteeeeeeeeeeeeeeeeeee net/ forum? id-ME9 EQ00gMeo..
Zhang, W. GX-Chen, A., Sobal, V., LeCun, Y., and Car- ion, N. Light-weight probing of unsupervised repre- sentations for reinforcement learning. arXiv preprint arXiv:2208. /2345, 2022b.
Tomar, M., Mishra, U. A., Zhang, A., and Taylor, M. E. Learning representations for pixel-based control: What matters and why? arXiv preprint arXiv:2111.07775, 2021.
Zhao, Y., Zhao, W. Boney, R., Kannala, J., and Pajarinen, J. Simplified temporal consistency reinforcement learning, In International Conference on Machine Learning, pp.
PMLR, 2023.
Torabi, F., Warnell, G., and Stone, P. Behavioral cloning from observation. arXiv preprint arXiv:1805.01991 2018.
Zheng, Q., Henaff, M., Amos, B., and Grover, A. Semi- supervised offline reinforcement learning with action- free trajectories. In International conference on machine learning, pp. 22339-42366 PMLR, 2023.
Torabi, F. Warnell, G., and Stone, P. Recent advances in imitation learning from observation. arXiv preprint 19XXvvI905.13355 2019.
Zhou, Q., Wang, J., Liu, Q., Kuang, Y., Zhou, W., and Li, Learning robust representation for reinforcement learn- ing with distractions by reward sequence prediction. In Evans, R. J. and Shpitser, I. (eds.), Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial In- telligence, volume 216 of Proceedings of Machine Learn- ing Research, pp- 2551-2562. PMLR, 31 Jul-04 Aug 2023. URL https:ttp///c httpttsss//ppcccediingnnooo /ttp://eoocceeeneeee/ee www.mmrrr.......... press/ v216 6 ht//2hou2333 a html.
Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances neural information processing systems, 30, 2017.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wan, S., Wang, Y., Shao, M., Chen, R., and Zhan, D.- C. Semail: eliminating distractors in visual imitation via separated models. In International Conference Machine Learning, pp. 054426335443 PMLR, 2023.
Zhu, C., Simchowitz, M., Gadipudi, S., and Gupta, A. Repo: Resilient model-based reinforcement learning by regu- larizing posterior predictability.
Oh, A., Naumann, T., Globerson, A.. Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information P Processing Sys- tems, volume 36, pp. 32445-32467. Curran Associates, Inc., 2023.
Wang, T., Du, S. S., Torralba, A., Isola, P., Zhang, A., and Tian, Y. Denoised mdps: Learning world models better than the world itself. arXiv preprint arXiv:2206.1552 15477, 2022.
Wang, Y., Wan, S., Gan, L., Feng, S., and Zhan, D.-C. Ad3: Implicit action the key for world models to dis- tinguish the diverse visual distractors. arXiv preprint 2009999 2024.
Yamada, J., Pertsch, K., Gunjal, A., and Lim, Task-inducee representation learning. arXiv preprint arXiv:2204. /127, 2022.
12

---

Latent Action Learning Requires Supervision in the Presence of Distractors
A. Additional Figures
26
66
21
24
LACM Prose from lstert actions
LAGH Proze from dbs emb
L40Ms superviionn frobe from ibbnt actions supee rision Probe from cbs emb
Nomallzed Return
LAG LSGM
34
16
Action
23
12
28
1.
1.
28
Latent action dim
Lotent action dim
trajoctorics
(a) Probe without supervision
(b) Probe with supervision
(c) LAOM with FSQ ablation
Figure 12. We provide additional ablations on the walker environment with three random seeds. (a)-(b) We took the observation embedding from the LAOM visual encoder and trained linear probe to predict real actions, similar probing from latent actions. (a) As can be seen, for LAOM the probe from observation embedding is better for smaller latent action dimensionality. This can be explained by fact that the information bottleneck induces the IDM to mainly encode noise in latent actions, as it can better explain the dynamics (deterministic distractors in the background), while observation embedding mostly preserves the information. At higher latent action dimensions, they are expected to equalize, as latent actions without bottleneck can encode the full dynamies, including noise and real actions. This is exactly the effect described in Section 4 which motivated us add supervision. (h) However, we different effect with LAOM-supervision, where the probe from the embedding observation is generally worse than from the latent actions, as with supervision
ground the latent actions focus on features relevant for control with small dimensions, filtering the noise. (c) We re-evaluate the effect of the quantization during LAM training, given all other LAOM improvements from Section and measuring actual performance in the environment instead of probing. As be seen, quantization indeed harmful for performance.
LAOM+supervision (ours)
LAOM (ours)
LAPO (original)
IDM
BC
cheetah-run
walker-run
hopper-hop
humanoid-walk
E 0.8
0.9
0.8
Nomaalzed Datuur
0.0
Normalized Return
0.4
0.3
0.6
7
0.5
0.5
0.4
0.2
0.4
0.4
0.3
0.1
0.2
0.1
0.1
0.0
0.0
8
16
37
64
128
2
37 Labeled trajectories
8
16
64
128
16 32 Labeled trajectories
8
64
128
16
37
64
12a
Labeled trajectories
Label trajectories
Figure 13. Main results without distractors, analogously to our main result in Figure 8. As can be seen, supervision help even without distractors, although all methods work good in this setting. Notably, IDM strong haseline.
LAOM (ours)
LAOM (aurs)
LAPO (ariginal)
BC
cheetah-run
walker-run
hopper-hop
humanoid-walk
Normalizzo Retum
Nomallize Return
0.3
Normalizeo loze Retur
0.2
Normalized memallized Retur
0.1
0.3
0.2
0.0
0.0
8
16
32
64
128
8
16
32
64
128
8
16
32
64
128
16
32
64
128
Labeled trajectories
Labeled trajectories
Labeled trajectories
Labeled ldbeleed trajectories
Figure 14.
pre- pre-triinin experiment results for each environment, For details see Figure 14.
13

---

Latent Action Learning Requires Supervision in the Presence of Distractors
LuOM+superviiion (ours)
0.7
LAOM (ours)
Normalized Return
0.6
LAPO (original)
IDM
0.5
BC
0.4
0.3
0.2
0.1
2
32
64
128
Labeled trajectories
Figure 15. Figure summarizing the results from Figure 13, analogously main result in Figure As be supervision help even without distractors, although all methods work good in this setting.
B. Additional Related Work
Learning with distractors. Distractors in various forms are commonly used in many sub-fields of reinforcement learning, such as: visual model-based learning, model-free learning, and representation learning.
In model-based learning, researchers explore ways to efficiently train world models that do not waste their capacity to model task-irrelevant details, either via decomposing world models to predict relevant and irrelevant parts separately (Fu al., 2021; Wang et al., 2022: Wan et al., 2023; Wang al., 2024) or by avoiding reconstructing observations (Okada Taniguchi, 2021; Deng al., 2022; Zhu et al., 2023; Liu al., 2023b; Burchi & Timofte, 2024). In our work, have similar need to not model action irrelevant details, as this will result in latent actions that describe changes in exogenous noise, not changes cased by ground truth actions. Thus, we use the commonly occurring latent temporal consistency loss (Schwarzer et al., 2020; Hansen et al., 2022; Zhao et al., 2023).
In model-free learning, researchers explore various techniques to improve generalization to new distractors and domain shifts (Hansen Wang. 2021; Hansen al., 2021; Bertoin al., 2022; Huang et al., 2022b; Batra & Sukhatme, 2024; Almuzairee al., 2024), which often revolves around the of augmentations (Ma al., 2022). In our work also augmentations, specifically subset of ones proposed by Almuzairee et al. (2024), stabilize LAM training with latent temporal consistency loss (Schwarzer et al., 2020).
In representation learning, researchers search for ways to obtain minimal representations that contain only task- (Yamada et al., 2022), reward- (Zhou et al., 2023) or control-related information (Zhang et al., 2020; Lamb et al., 2022; Liu et al., 2023a; Ni et al., 2024; Levine et al., 2024), as this can greatly increase sample efficiency and generalization (Kim et al., 2024a). In our work, inspired by Lamb et al. (2022), we incorporate the multi-step IDM into LAM and show that it can help learn better latent actions in the presence of exogenous noise. Moreover, when small number of ground truth actions is available for pre-training (see Figure 1), our model them conceptually reduces one proposed by Levine al. (2024), for which it has been theoretically shown that it can recover control-endogenous minimal state. This may explain why incorporating labels during LAM pre-training, rather than during final fine-tuning, brings so much benefit, since discovering true actions is trivial given a minimal state. We however, found a contradicting evidence, as Figure 11 shows that our proposed methods do not learn minimal state practice.
Overall, although we were inspired by existing approuches, they have not previously been used improve latent action learning, especially in combination, which, as show (see Figure 6) is essential for good performance in the presence of distractors.
C. Data Collection
We used environments from the Distracting Control Suite (DCS), wrapped with Shimmy wrappers for compatibility with the Gymnasium API. For cheetah-run, walker-run and hopper-hop we used PPO (Schulman al., 2017), adapted from
14

---

Latent Action Learning Requires Supervision in the Presence of Distractors
the CleanRL (Huang al., 2022a) library. For humanoid-walk, used SAC (Haarnoja et al., 2018) from the stable- baselines3 (Raffin et al., 2021) library, PPO from CleanRL. was not able solve it at the expert level. We used default hyperparameters and trained on 1M transitions in each environment, except for numanoid-walke where we trained on 100k transitions. Importantly, for speed, all experts were trained with proprioceptive states and no distractors, we later rendered proprioceptive states to 64px images with or without distractors during data collection. For each environment, we collected 5k trajectories, with an additional 50 trajectories for evaluation with novel distractor videos (from the evaluation set the DCS). As each trajectory consists of 1000 steps, the datasets contain 5M transitions. We include ground truth actions and states for debugging purposes. The datasets will be released together with the main code repository.
Table Datasets statistics.
Dataset
Average Return
Size (GB)
cheetah-run
837.70
57.7
walker-run
739.79
57.8
hopper-hop
306.63
57.6
numanoid-walee
617.22
58.9
D. Implementation Details
All experiments were run on H100 GPUs, in single-gpu mode and PyTorch 6 precision with AMP. For the visual encoder, we used ResNets from the open-source LAPO (Schmidt & Jiang, 2023) codebase, which also borrowed from baselines originally provided as part of the ProeGen 2020 competition. For the action decoder, we used a two-layer MLP with 256 hidden dimensions and ReLU activations.
In contrast to the commonly used cosine similarity, used MSE for temporal consistency loss. We also found that projection heads degraded performance, we did not use them. We use slightly non-standard MLP for latent IDM and FDM: we compose it from multiple MLP blocks inspired by Transformer architecture (Vaswani, 2017) and condition on latent action and observation representation on all layers instead of just the first. We have found that this greatly improves prediction, especially for latent actions. We also use ReLU6 activations instead of GELU, as it naturally bounds the activations, which helps with stability during training, similar to target networks in RL (Bhatt et al., 2019). Without supervision, we use the EMA target encoder. With supervision, we find that a simple stop-grad is sufficient to prevent any signs of collapse, a finding also reported by Schwarzer et al. (2020).
For all experiments open-source the code at https: / / thub. http//ddnllaaaaaooooooo
the cosine learning late schedule with warmup. For hyperparameters Appendix F. We
Table 2. Methods training time summed from all stages (including online evaluation) for each method.
Method
Training Time
LAPO
7h 38m
LAOM
6h 43m
-supervinn
7h 6m
BC
lh 10m
IDM
~ 5h 30m
E. Evaluation Details
We outline the evaluation procedures used in our experiments for each method. First, we review the general setup. For each environment, we have large dataset without action labels, with and without distractors. To decode the learned latent actions to ground truth for evaluation, we allow a small amount of action labeled data, in line with previous work (Schmidt & Jiang, 2023; Ye et al., 2024). We sample once from the existing dataset, revealing true actions, to ensure that all methods are on equal conditions. We use identical backbones where possible, and try our best to make all methods equal in the number of
15

---

Latent Action Learning Requires Supervision in the Presence of Distractors
trainable weights. For hyperparameters, see Appendix F. We report the scores achieved by BC trained datasets with all actions revealed in Table 4. We use these for normalization in all experiments.
BC. We trained BC from scratch to predict ground-trui truth actions on available labels, i.e. on 2 or 128 trajectories.
IDM. We used two-staged pipeline. First, we trained IDM to predict actions on available labels, i.e. on 2 trajectories. Then, we trained BC on full unlabeled dataset, providing labels via pre-trained IDM. We report BC final return.
LAPO and LAOM. We used three-stage pipline. First, we pre-train latent actions on full unlabeled datasets. Then, we trained BC, providing latent action labels via pre-trained LAM. Finally, we trained action decoder on small amount of labels, while freezing the rest of the policy weights.
LAOM+supervision. Almost like LAOM, with the difference being that we exactly aligned stages in terms of action labels used. While in LAOM we can pre-train it once and then re-use for later stages regardless of the number of action labels, in LAOM+supervision we trained separate LAM for each budget of labels. Thus, for LAOM+superiiiion trained with supervision from trajectories of labels, on final stage the decoder was trained only on the same trajectories. We repeat this process for all cases, from 2 to 128 trajectories.
Table Evaluation returns of BC trained on full datasets with ground-truth actions revealed. We use them for normalization.
Dataset
With distractors
Without distractors
cheetah-run
823
840
walker-run
749
735
hopper-hop
253
300
humanoid-we
428
601
Table 4. Total parameters for each method according to the hyperparameters used in Appendix F.
Dataset
Total Parameters
LAPO
214447849
LAOM
192307 136
LAOM+supervision
192479189
BC (on all stages)
107541504
IDM
192258965
16

---

Latent Action Learning Requires Supervision in the Presence of Distractors
F. Hyperparameters
Table LAPO hyperparameters. We use the same hyperparameters for all experiments and explicidy mention any exceptions. Names are exactly follow the configuration files used in code.
Stage
Parameter
Value
grad nord
None
batch_size
512
num epochs
10
frame-stack
encoder_deep
False
weight_decay
None
Latent actions learning
pncoder-scale
learning rraee
0.0001
warmup-epose -epcchs future obs offset
10
latent_action_din
8192
encoder num res blocks
dropout
0.0
use_aug
False
batch_size
512
num_epochs
10
frame.stack
Latent behavior cloning
encoder_des deep
False
weight_decay
None
encoder scale
32
learning_nad rrate
0.0001
warmup epochs
incoder_num_res_bie
use_aug
False
batch_size
512
hidden dim
256
beight-decay eval episodes
None
Latent actions decoding
25
learning_rat
0.0003
total updates
2500
warmup_epochs
0.0
17

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Table LAOM hyperparameters. We use the same hyperparametens for all experiments and explicitly mention any exceptions, Names are exactly follow the configuration files used in code.
Stage
Parameter
Value
use aug
True
grad_norm batch_sin
None
512
num_epochs
10
target tau frame_stack act_ heeddeeddddmm encoder_deep
0.001
1024
False
obs head dim Beiightddcaay encoder-scale
1024
None
Latent actions learning
learning_ran .rate warmup_epoke _epochs encoder dropout act_ .dropot encoder_nore wnommonoommmmmmmmmmm....mm oout
0.0001
0.0
0.0
False
bos_head_droppe
0.0
future obs offset aaaticnndiin target.update_evv eeeery
10
8192
encoder num res blocks
dropout
0.0
use_aug
False
batch size
512
num.epochs
10
frame-stack
Latent behavior cloning
encoder_deep
False
weight decay
None
dicoder_scalo
32
learning_rat rate
0.0001
warmup epochs
0.0
encoder.num.res.bleck
use_aug batch size hidden_dim weight evisoooss learning rate total updates
False
512
256
None
Latent actions decoding
25
0.0003
2500
warmup cermup-epochs
18

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Table LAOM-superviion hyperparameters. We same hyperparameters all experiments explicitly mention any exceptions. Names are exactly follow the configuration files used in code.
Stage
Parameter
Value
use_aug
True
grad norm
None
batch_size
512
num_epochs
10
target_tau frame stack
0.001
3
act act_head_din encoder deeep
1024
False
obs_head _dim
1024
seight.decayy encoder
0.0
Latent actions learning
learning-rate
0.0001
warmup_epocht
3
mncoder_dropout act head dropout encoder_none _orrmmou obs hbdd drroout future_obs, citure_obb_offfeee oof_effset labeled loss coef latent cecntttction__ _dim size
0.0
0.0
False
0.0
10
0.01 (0.001, cheetah-run)
8192
128
target update every incoder.num.res.bleee.
dropout use aug
0.0
False
batch_size num_epochs frame_stack
512
10
Latent behavior cloning
encoder deep
False
Beight_decaa
None
encoder_scite learning rate
32
0.0001
warmup-epocht
encoder num res blocks
aug
False
batch_size hidden_dim weight_decay eval. eplleepssdess
512
256
Latent actions decoding
learning rate total_updates warmup epochs
0.0003
2500
19

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Table 8. IDM hyperparameters. We use the same hyperparameters for all experiments and explicitly mention any exceptions. Names exactly follow the configuration files used in code.
Stage
Parameter
Value
use_aug
False
grad_nore norm
None
batch_size
512
frame stack
3
act_head_ding
1024
encoder_deep
False
weight decay
None
IDM learning
encoder scale
12
learning_ rate
0.0001
total-updater
10000
warmup epochs
3
.dropout act head dropout
0.0
0.0
future_obs, _bs_offsettttttooooooooeessss.e_
encoder num res. blocks
2
dropout
0.0
use aug
False
batch size
512
num_epochs frame-stack encoder deep
10
3
Behavior cloning on IDM actions
False
deight
None
encoder scale
32
eval_epiae rea_eepsodess
25
learning rate
0.0001
warmup.epor
ancoder.num_res.bee
2
Table 9, BC as baseline hyperparameters. We use the same hyperparameters for all experiments and explicitly mention any exceptions. Names are exactly follow configuration files used in code.
Parameter
Value
dropout
0.0
use_aug
false
batch-size
512
frame stack encoder.deep
3
false
weight_
encoder _scale
32
eval episodes learning rate
25
0.0001
total updates carmup_epochhe
10000
0
cooldown ratio encoder num res blocks
20

---

Latent Action Learning Requires Supervision in the Presence of Distractors
Table 10. BC for normalization hyperparameters. We use the same hyperparameters for all experiments and explicitly mention any exceptions, Names are exactly follow configuration used code.
Parameter
Value
dropout
0.0
use_aug
false
batch size
512
frame_stack
3
encoder deep
false
weight_decay
encoder scale
32
eval evil_epis
25
learning rate num eem_es
0.0001
10
warmup-epoccs
0
cooldown_ratee
0
sonider_num_ree_blooe
2
21


Reviews:
[{"review_id": "NhL5p3ukCu", "reviewer": "Reviewer_HSFB", "summary": "- The paper focuses on LAMs, which aim to infer control actions from unlabelled videos\n- Here the authors note a benefit of reusing action labels from later in the pipeline to help focus (through supervision) latents on control actions\n- This is most effective in the presence of 'distractors', ie non-control action changes\n- Empirical results are show on distracting control suite\n- A broader investigation of various model design choices and extensions are conducted", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful and constructive feedback. We address the main questions below.\n\n> Missing this, which follows a pipeline similar to LAPA: IGOR…\n> \n\nWe will include the citation, thank you for your suggestion.\n\n> … naturally a wider bottleneck will allow more information about both distractors and actions into the latent …. But is having extra distractor information not harmful?\n> \n\nAs we explained in more detail in Section 4, the removal of quantization and the increase in dimensionality is a necessity. Without these changes, latent actions will have no information about real actions at all, and as a consequence will be useless for further pre-training, as demonstrated by the LAPO performance. The best we can hope for in the general case (without supervision) is to encode all dynamics, including noise, but most importantly real actions, into latent actions. That's what the changes in LAOM do. LAOM also does not guarantee that latent actions will be minimal. However, now that we are sure that the real actions are contained in the latent ones, we can hope that with a little supervision we can extract them in a generalizable way. We believe that the results in Figures 1, 7 and 8c clearly show that this is indeed the case.\n\n> I'm surprised with this capacity the model doesn't just pipe the entire next observation through the latent\n> \n\nThis is an interesting question that we did not explore in depth as we felt it was beyond the scope of the study. However, we did not observe any evidence of shortcut learning. It is important to note that we used encoders that were not very small for a given task. Given that we are working with 64x64 images, they have enough capacity to predict the next state in pixel space very accurately. As we show in the appendix, the encoders take up about 200M parameters in total.\n\n> I'd be interested to dive deeper into the differences between LAM and IDM\n> \n\nWe believe that IDM will perform better than LAM in the limit (e.g. see Figure 9 in [GR00T N1](https://arxiv.org/abs/2503.14734)), but it's also limited to a single action space. However, when the number of labels is very limited, LAM will perform better due to better generalization. Overall, we believe that LAM+supervision combines the best of both worlds with better use of existing action labels.\n\n> Could you detail the architecture for IDM FDM and pre-trained latent policy? Also BC baseline. Are these CNNs? In appendix I only saw 'encoder num res blocks'\n> \n\nWe provide some details of the architectures in Appendix D. We use the same visual encoder architecture for all methods (IDM, FMD, BC), which is a simple ResNet borrowed from the open source LAPO code. For FDM we use an identical architecture, swapping conv downsampling with transposed conv upsampling. LAPO, similar to the original code, uses separate encoders in IDM and FDM, while LAOM shares one encoder between them. For LAOM latent FDM we use multiple MLP blocks, inspired by MLP from transformer architecture. To process successive observations, we concatenate images across channels. BC uses ResNet + small action head. The action decoder is a two-layer MLP with a hidden dim of 256.\n\n> Augmentation details\n> \n\nWe use several types of augmentations: shift, rotate, change perspective, and combinations such as shift-rotate, rotate-perspective, etc. (as we note in the paper, they are taken from Almuzairee et al., 2024). We sample augmentations for each sample in a batch, but share them across the dimension of the frame stack.\n\n> Would love to see a zoomed in plot of fig 8 b as can't tell how the MSE changes for the LAOM+supervision.\n> \n\nSorry, here's an enlarged [Figure](https://ibb.co/5X6B5KdH). We will add it to the appendix. Overall, the MSE does not change that much here.\n\n> For fig 10, what representation is used for the IDM?\n> \n\nThe IDM can be schematically described as $a = h(f(s_t), f(s_{t+1}))$, where $f$ is a ResNet encoder and $h$ is an action head consisting of MLP. For visualization we used observation embedding after the ResNet encoder, that is $f(s_t)$.\n\n> In fig 8 c, if it seems like the models keep improving with latent dim -- why stop at 8192?\n> \n\nMostly for practical reasons, since increasing the dimensionality of latent actions increases the overall computational requirements for pre-training. For example, we had to significantly increase the size of the BC in order for it to learn latent actions of dimension 8192 accurately enough (real actions have dimensionality of about 4-16). In fact, to perform with similar quality on real actions, the BC might have been about 4-8 times smaller, because predicting 4 numbers is much easier (LAPA report similar results). To ensure fairness we used larger BC size in all experiments. Our main goal was to show a trend, so we did not see the need to go further.", "claims_and_evidence": "Yes, see strengths weaknesses.", "methods_and_evaluation_criteria": "Yes, see strengths weaknesses.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "See strengths weaknesses.", "relation_to_broader_scientific_literature": "See strengths weaknesses.", "essential_references_not_discussed": "Missing this, which follows a pipeline similar to LAPA.\n- IGOR: Image-GOal Representations are the Atomic Control Units for Foundation Models in Embodied AI", "other_strengths_and_weaknesses": "Strengths\n- While the problem setting departs from the standard three-stage LAM setup, in my opinion it remains a realistic use-case, where the same set of labels are being used twice in a smart way\n- The method is simple\n- Overall the method works well, the paper in general directly addresses an issue many LAM users may be interested in\n- Good related work description\n- Thorough investigation of various aspects of LAMs going beyond the main message provides breadth to the paper (eg fig 8, 9, 10, 12)\n- Reasonable eval metrics, normalizing by BC with all labels was a nice touch\n- Reasonable baselines, nice to see improvements over the IDM baseline\n- Interesting intuition about differences between LAM and IDM approaches (generalization)\n\nWeaknesses\n- One major difference with previous work is the decision to open up the bottleneck -- both removing the quantization and moving to a large 8192 dims. This is originally justified by fig 4 and 5, however these only measure MSE -- naturally a wider bottleneck will allow more information about both distractors and actions into the latent, so measuring MSE on actions will reduce. But is having extra distractor information not harmful? Fig 8c partially suggests not.\n- In general I'm surprised with this capacity the model doesn't just pipe the entire next observation through the latent -- one possible reason I thought is that the architectures are not powerful enough to unpack all details through the bottleneck, see q about architectures.\n- Overall this is not a critical point as it seems to be necessary to get the method to work anyway.\n- Some of the architectural choices -- eg latent reconstruction and augmentation -- are becoming pretty standard in other LAM works, so there is limited novelty there, though the paper doesn't overclaim on this aspect anway.\n- I'd be interested to dive deeper into the differences between LAM and IDM -- fig 8a suggests the gap closes at some point. But this is probably better left to future work.\n\nQuestions\n- Could you detail the architecture for IDM FDM and pre-trained latent policy? Also BC baseline. Are these CNNs? In appendix I only saw 'encoder num res blocks'\n- Augmentation details -- is there any structure in how you do the random augs? Like do inputs in a sample receive the same crop or?\n- Would love to see a zoomed in plot of fig 8 b as can't tell how the MSE changes for the LAOM+supervision.\n- For fig 10, what representation is used for the IDM? The final predicted action?\n- In fig 8 c, if it seems like the models keep improving with latent dim -- why stop at 8192?", "other_comments_or_suggestions": "NA", "questions_for_authors": "See strengths weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "puwLJ2rmJS", "reviewer": "Reviewer_sdKR", "summary": "This paper focuses on the setting of learning latent actions in the presence of background distractions. The authors investigate improving upon prior latent action pretraining work with recent advances in dynamics and latent action modeling. It shows that multi-step inverse dynamics, large latent action dimension without quantization, forward dynamics in latent space, and augmentations help improve latent action quality when distractors are present, and that supervised training with a small amount of ground truth actions can help with latent action quality and close the gap to behavioral cloning performance under the background distraction setting.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for their time, constructive feedback, and suggestions for additional experiments, which we found very valuable. We have tried to answer the questions below.\n\n> Clarification: is the BC policy also only trained on up to 128 trajectories?\n> \n\nYes, the BC baseline we show in Figures 1, 7, 9 uses the same architecture as the BC in LAM methods, but is simply trained from scratch only on trajectories with available ground truth action labels, from 2 to 128.\n\nWe also use a separate BC for normalization on all figures. We pre-train it on full datasets with all action labels revealed to get the maximum possible score with ground truth actions. With such normalization, we can quantify how much performance we have recovered compared to having access to a fully action-labelled dataset.\n\n> Clarification: what is the IDM trained on? Is it trained on up to 128 trajectories of ground truth actions, then used to relabel the full dataset for downstream behavioral cloning?\n> \n\nYes, this is an accurate description of the overall pipeline.\n\n> Baseline: how large is the observation embedding dimension? In Figure 8(b), what is the action probe MSE if you probe directly from the observation embedding?\n> \n\nThanks for the suggestion! Due to time constraints, we ran these experiments with three random seeds, but only in the walker environment. Given the previous evidence, we are confident that the results will hold in the remaining environments and will include additional results in the camera-ready version of the paper. We took the observation embedding from the LAOM encoder and trained the linear probe to predict real actions, similar to probing from latent actions. We visualize the results in the following figures [[Figure 1](https://ibb.co/wFRGMhDV), [Figure 2](https://ibb.co/gL6L0zNf)].\n\nAs can be seen, for LAOM it is indeed the case that probe from observation embedding is better for smaller latent action dimensionality. This can be explained by the fact that the information bottleneck induces the IDM to mainly encode noise in latent actions, as it can better explain the dynamics (deterministic distractors in the background), while observation embedding mostly preserves the information. At higher latent action dimensions, they are expected to equalize, as latent actions without bottleneck can encode the full dynamics, including noise and real actions. This is exactly the effect we described in Section 4, which motivated us to add supervision.\n\nHowever, we see a different picture with LAOM+supervision, where the probe from the embedding observation is generally worse than from the latent actions, because with supervision we can ground the latent actions to focus on features relevant for control even with small dimensions, filtering out the noise.\n\n> Question: why is it the case that LAOM + supervision achieves a low action probe MSE regardless of latent action dimension, and LAOM without supervision action probe MSE depends on the latent action dimension?\n> \n\nWe believe that in the absence of supervision, as we discuss in Section 4, the information bottleneck is detrimental, as it incentivises the IDM to encode into latent actions a minimum amount of information that is maximally predictive of the next observation. In the case of distractors, this will mostly be noise, as it is easier to predict deterministic videos in the background than actual actions (which also explain much more variation in the overall dynamics). By increasing the latent action dimension, we remove the bottleneck and allow LAOM to encode the full dynamics in actions, including actions but also noise.\n\nOn the other hand, LAOM+supervision grounds the latent action space to be predictive of actual actions, which can be much smaller because it does not need to explain noise (actual actions are only ~4-16 dimensions).", "claims_and_evidence": "The paper claims that background distractions requires several changes to current latent action inference methods to improve latent action quality and supports it with experimental results.", "methods_and_evaluation_criteria": "The evaluations overall makes sense. The authors test on extracting latent actions with background distractions, and evaluates the latent action quality with linear probe and downstream policy evaluation.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Overall the experiments make sense. However, in this paper, there seems to be no bottlenecks (neither VQ nor dimensional bottleneck) in the latent action. One concern is whether the latent action, as evaluated by downstream policy performance and action probe MSE, is not only a high-dimensional projection of the observation embedding. It would be convincing to see an additional baseline where downstream policy training / action probe MSE is evaluated directly on the observation embedding.", "relation_to_broader_scientific_literature": "The paper develops on recent advances in latent action pretraining such as LAPO, DynaMo, LAPA, etc., which investigates pretraining latent action policies and visual representations from video data. There is also a rich line of work investigating policy robustness under visual distractions. Extracting latent actions from videos is a longstanding important topic for learning robotic policies, as well as dealing with background distractors.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: the paper investigates extracting latent actions from video data with background distractors, an important setting to consider for generalizing to Internet-scale videos. It shows that with a few modifications, prior work in latent action pretraining such as LAPO can be improved both in the distractor-free and background distraction settings.\nWeaknesses: see experimental designs and questions.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- Clarification: is the BC policy also only trained on up to 128 trajectories? \n- Clarification: what is the IDM trained on? Is it trained on up to 128 trajectories of ground truth actions, then used to relabel the full dataset for downstream behavioral cloning?\n- Baseline: how large is the observation embedding dimension? In Figure 8(b), what is the action probe MSE if you probe directly from the observation embedding?\n- Question: why is it the case that LAOM + supervision achieves a low action probe MSE regardless of latent action dimension, and LAOM without supervision action probe MSE depends on the latent action dimension?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rwcDfm2uwL", "reviewer": "Reviewer_o8og", "summary": "Latent actions prove to be useful for efficient policy pretraining from unlabeled videos. This paper aims to enhance the quality of latent actions by removing the original information bottleneck, leveraging multi-step future observations, and predicting future states in the latent space. The authors also suggest that adding a small amount of action supervision can significantly mitigate the effects of distractors. The proposed methods are validated through linear probing accuracy and normalized returns in downstream tasks.\n\n## *Update after rebuttal*\n\nI remain unpersuaded by the authors' response. Specifically, while I agree that the modifications bring positive gains in the scenarios created by the authors, I still have the following major concerns:\n\n- Due to the prediction loss and the extremely large latent dimension, the use of action supervision does not effectively produce a latent action representation. As revealed by Figure 10 and the responses to reviewers HSFB and sdKR, the latent capacity is so redundant that it passes almost all observation information through the latent. As a result, the policy is more like a video prediction policy rather than a latent action policy. However, the proposed method is not compared or discussed with previous video prediction policies [1,2].\n- The conclusions may not be applicable to real-world applications (the authors have also mentioned this in Appendix A). The distractors are created by copying and pasting the background directly, which are quite different from real-world patterns. Additionally, the agents in the studied scenarios are always centered and share a simple appearance, which differ from the settings studied by previous latent action models [3,4]. While it is common to utilize simplified environments for research, the scenarios created by the authors are too unique to verify their applicability to true decision-making scenarios.\n- Typical decision-making scenarios, such as Open-X-Embodiment and ProcGen, do not require using the proposed method (the authors confirmed in their response to my rebuttal comment). This also greatly compromises the applicablity of the proposed methods to real applications.\n- It would be beneficial to incorporate convincing visualizations like LAPO [3] to demonstrate the distribution of latent actions.\n\n---\n\n[1] Learning Universal Policies via Text-Guided Video Generation\n\n[2] Learning to Act from Actionless Videos through Dense Correspondences\n\n[3] Learning to Act without Actions\n\n[4] Latent Action Pretraining from Videos", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and feedback. We have tried to address the concerns below.\n\n> I have some concerns about how the paper assesses the quality of latent actions. The objective of extracting latent actions is to fully encode action information while minimizing background noise, and the paper utilizes linear probing error to judge the quality of latent actions. However, a lower error does not necessarily equate to higher quality latent actions. Instead, it only indicates that the encoded latent space contains more action information, regardless of the ratio of useful action information.\n> \n\nThis is indeed a valid concern, which we already discuss in detail in the paper (see Section 4, second column, from lines 269-270). As the reviewer correctly points out, linear probing does not tell us that the resulting latent actions are minimal (in the sense that they contain only information relevant to actions without noise), but only allows us to detect the amount of information about real actions in the latent ones.\n\nIt should be noted that the fact that latent actions contain information about real actions is a **necessary condition** for their usefulness. If the latent actions do not contain any information (as expressed by poor probe loss), this automatically means that they are useless for further BC pretraining.  Therefore, before worrying about minimality, it is important to make sure that the latent actions contain real actions at all, which is exactly what we do with our LAOM modifications.\n\nWith our experiments in Section 4, we show that in the presence of distractors, naive LAPO (especially with quantization) does not produce latent actions that contain sufficient information about real actions, and thus cannot be used for efficient pre-training. LAOM improves this by a factor of eight, ensuring that latent actions contain real actions, which directly translates into a twofold improvement in return. Does LAOM guarantee latent action minimality? It does not, and we state this explicitly in the paper (first column from line 320). On the contrary, without supervision, LAOM still performs poorly (Figure 6), although we can now hope that supervision will allow us to extract action information from latents, which would be impossible with LAPO (which is clearly illustrated by our main experiment in Figure 1). Demonstrating this was one of our main goals.\n\n> W1) **The reliability of the metric.** … I believe that linear probing error is not a reliable metric for assessing the quality of latent actions, as it only evaluates how much action information are contained in the latent.\n> \n\nWe re-emphasise that linear probes were used only to show that LAPO does not work in the presence of distractors, and that even LAOM does not guarantee minimality. Our main contribution and claim is about the need for supervision, and we explicitly demonstrate this in the second part of the paper with experiments in the real environment and with real return. We believe that such clear improvements in return (see Figure 1) clearly indicate the higher quality of latent actions of LAOM vs. LAPO, and LAOM+supervision vs. LAOM.\n\n> W2) **The effectiveness of the method.** \n\nAs we stated above (and in responses to other reviewers), linear probing is not our final (and only) metric. It was only used to show that LAPO does not learn good latent actions, and that even after the modifications considered in LAOM - performance did not increase. So the fact that LAOM performs poorly on Figure 6 is not a problem with our approach. On the contrary, this is exactly what we wanted to show in order to highlight the general issues with latent action learning in the presence of distractors. And this is what is greatly improved by the addition of supervision (see Figure 1, 7).\n\n> W3) **The scope of the proposed setting.**\n\nWe have to respectfully disagree with the reviewer's assessment. We're not changing the setting in any way, as LAPO, LAPA and other methods still require real actions, we're just suggesting that they be used differently. On the contrary, it is precisely LAPO and related work that considers a simplified setting that lacks the distractors common to real videos, which in turn is not aligned with the goal of pre-training on Internet-scale data. The main purpose of our work was to show that vanilla LAPO will not scale to Internet videos due to distractors.\n\n> Q1) **Procgen results without distractors.**\n\nOur aim was not to provide a state of the art method, but to study the properties of latent action learning methods in the presence of distractors, which is absent in the existing literature. LAOM is more of a set of suggestions for practitioners than a stand-alone method. Therefore, we think that the experiments on ProcGen are out of scope, as they do not allow us to study any additional properties of LAPO in the presence of distractors.", "claims_and_evidence": "Yes, there are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "I have some concerns about how the paper assesses the quality of latent actions. The objective of extracting latent actions is to fully encode action information while minimizing background noise, and the paper utilizes linear probing error to judge the quality of latent actions. However, a lower error does not necessarily equate to higher quality latent actions. Instead, it only indicates that the encoded latent space contains more action information, regardless of the ratio of useful action information.", "theoretical_claims": "I have checked the theoretical claims in this paper.", "experimental_designs_or_analyses": "The paper provides extensive experiments with sufficient details.", "relation_to_broader_scientific_literature": "The paper aims to improve the effectiveness of action pretraining, which could benefit the development of embodied AI.", "essential_references_not_discussed": "As far as I know, all closely related works are cited appropriately.", "other_strengths_and_weaknesses": "W1) **The reliability of the metric.** As I mentioned earlier, the quality of latent actions is influenced not only by the amount of action information encoded but also by the ratio of useful information to noise. Therefore, I believe that linear probing error is not a reliable metric for assessing the quality of latent actions, as it only evaluates how much action information are contained in the latent.\n\nW2) **The effectiveness of the method.** Guided by the inappropriate metric, the proposed LAOM removes the information bottlebeck and significantly increase the latent action dimension to 8192. However, according to the downstream performance (blue lines) in Fig.6, these modifications yield little improvement and can even lead to worse performance with 2 and 4 labeled trajectories.\n\nW3) **The scope of the proposed setting.** While introducing action labels mitigates the negative effects of distractors, it may be somewhat unfair compared to the motivation of LAPO (Schmidt & Jiang, 2023). LAPO and its related works exclude action labels during the pretraining stage to ensure their algorithms only rely on videos. This assumption has huge potential to exploit Internet-scale and cross-embodiment data, even without a consistent action format. However, the proposed setting is not aligned with this goal.", "other_comments_or_suggestions": "Please see the questions below.", "questions_for_authors": "Q1) **Procgen results without distractors.** Aside from the environments with distractors, is it possible to compare the proposed techniques with LAPO on the original Procgen benchmark to demonstrate their effectiveness? I believe this would help readers fully understand the properties of the proposed method.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kzCVhAjoCs", "reviewer": "Reviewer_nu2R", "summary": "This paper presents an empirical study of latent action learning in the presence of distractors. They found that latent action learning struggles with distractors, and propose several changes in architecture to improve latent action learning. Notably, they found supervision with a small amount of action labels could significantly improve latent action learning with distractors.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort. We address the questions below.\n\n> I am curious about the claim on \"Quantization hinders latent action learning\". it seems that this is only verified by linear probing. However, as the author mentioned, linear probing has a major limitation - it can only tell us whether real actions are contained in latent actions or not. Removing quantization, in some sense, is similar to increasing the dimensionality. As a result, I believe it would be better if the author could study this claim with experiments on the BC stage.\n> \n\nYou are absolutely right that removing quantization has a similar effect to increasing the dimensionality of latent actions, as both changes loosen the information bottleneck, which we believe can be harmful in the presence of distractors (for details see the last response to reviewer **sdKR**).\n\nHowever, we disagree with the conclusion about linear probes. As you rightly point out, probes can only tell us whether real actions are contained in latent actions or not. If probe loss is low, we cannot be sure that latent actions are minimal, only that they contain real actions. To test their true quality, we need to pre-train BC and fine-tune it in a real environment. Thus, low probe loss is a necessary (but not sufficient) condition for good latent actions. This means that if we get a high probe loss, it definitely means that the latent actions do not contain real actions and are therefore useless for subsequent fine-tuning. This is the reason why we did not consider FSQ further in the later stages, as its addition only worsens the probe loss and therefore definitely worsens the final result.\n\nHowever, to be certain, we performed an experiment with LAOM+FSQ. Due to time constraints, we only used walker environment and 3 random seeds (see [Figure](https://ibb.co/VY0djs7T)). As can be seen, FSQ does indeed worsen the downstream performance after fine-tuning. We will include this Figure (and for the remaining environments) in the Appendix.\n\n> the claim on line 126~127 is not very accurate. The NN architecture in LAPO is not used in motoGPT, Dynamo, and LAPA. I believe they use different architectures.\n> \n\nYou are right, we did not express ourselves clearly. In our case, the details of the architecture itself are not very important (even LAOM itself does not work well without supervision). What is important is that mathematically all these architectures do exactly the same thing as LAPO, and so inherit the same limitations. We will correct this statement in the new version of the paper.\n\n> The paper fails to cite the … \"IGOR: Image-GOal Representations …\"\n> \n\nThank you for the suggestion, this is indeed a highly relevant paper. We will include the citation.\n\n> I'm curious about the results shown in Figure 10. How is the separate decoder trained? Is the distractor constant, or does it display temporal correlations within each trajectory? If that's the case, then even when using only true actions, one might expect that an FDM decoder could learn to predict the distractor.\n> \n\nFor the decoder, we used an observation embedding after the ResNet encoder for each method. We trained it to reconstruct the observation and did not pass the gradients through the embedding to avoid changing the main training loop. The distractors are dynamic and change as the episode unfolds (video plays in the background, agent colour changes, camera shakes).\n\nYou are right that with real actions, FDM from the original LAPO will learn to predict distractors. However, this is a problem of prediction in the pixel space. A method similar to LAOM, which predicts the next observation in latent space, with ground truth actions will provably recover the control endogenous minimal state, filtering out the distractors (see Preliminaries and [Multistep Inverse Is Not All You Need](https://arxiv.org/abs/2403.11940) paper). Why doesn't this happen with LAOM+supervision? Probably because the number of ground truth actions is extremely small in our case.", "claims_and_evidence": "I am curious about the claim on \"Quantization hinders latent action learning\". it seems that this is only verified by linear probing. However, as the author mentioned, linear probing has a major limitation - it can only tell us whether real actions are contained in latent actions or not. Removing quantization, in some sense, is similar to increasing the dimensionality. As a result, I believe it would be better if the author could study this claim with experiments on the BC stage.", "methods_and_evaluation_criteria": "Yes. It follows the experiment setting proposed in LAPO.", "theoretical_claims": "no theoretical claim", "experimental_designs_or_analyses": "yes. seems sound to me", "relation_to_broader_scientific_literature": "the claim on line 126~127 is not very accurate. The NN architecture in LAPO is not used in motoGPT, Dynamo, and LAPA. I believe they use different architectures.", "essential_references_not_discussed": "The paper fails to cite the following work on latent action learning that also introduce some improvements, such as using random cropping on data—an approach similar to the \"adding augmentation\" method discussed in the submission.\n\n\"IGOR: Image-GOal Representations are the Atomic Control Units for Foundation Models in Embodied AI\"", "other_strengths_and_weaknesses": "Overall, I appreciate the paper as an empirical study on latent action learning with distractors. The problem setting is realistic, and the experimental work is solid.\n\nOne limitation is that the study focuses solely on LAPO without comparing it to other latent action learning methods such as Genie, LAPA, IGOR, and MotoGPT, which utilize varying neural network architectures for the latent action model. However, this appears to be an inherent challenge in latent action research—each work introduces similar yet distinct latent action models tailored for different applications. As a result, direct comparisons between existing approaches are limited, making it difficult to determine which neural network architecture represents the state-of-the-art.", "other_comments_or_suggestions": "No.", "questions_for_authors": "I'm curious about the results shown in Figure 10. How is the separate decoder trained? Is the distractor constant, or does it display temporal correlations within each trajectory? If that's the case, then even when using only true actions, one might expect that an FDM decoder could learn to predict the distractor. Additionally, I'm not surprised that the LAM model is able to predict distractors, since it is trained by reconstructing the next frame, which naturally extends to predicting distractors as well.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper received generally positive reviews for its empirical investigation into latent action learning in the presence of distractors. Reviewers appreciated the realistic problem setting, clear motivation, and thorough experimental analysis. The work demonstrates that small amounts of action supervision and design changes—such as increasing latent dimensionality and incorporating forward dynamics—can significantly improve performance in distractor-heavy environments.

While the method shows solid empirical gains, one reviewer expressed strong reservations regarding the scalability and real-world applicability of the setting. Concerns include the lack of strong comparisons to prior video prediction models, limited generalizability beyond synthetic distractions, and the potential misalignment with the broader goals of actionless learning. Additionally, some questioned the informativeness of the chosen metrics and the lack of evaluation on more diverse benchmarks like Procgen.

Despite these limitations, the paper offers valuable insights into a challenging and under-explored aspect of latent action learning. The AC finds the contributions relevant and timely and recommends acceptance, while encouraging the authors to better contextualize their setting and discuss broader applicability in the final version.
Decision: Accept (poster)