Latent Action Learning Requires Supervision in the Presence of Distractors

We thank the reviewer for their time and effort. We address the questions below.

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.

You 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).

However, 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.

However, 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). 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.

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.

You 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.

The paper fails to cite the … "IGOR: Image-GOal Representations …"

Thank you for the suggestion, this is indeed a highly relevant paper. We will include the citation.

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.

For 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).

You 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 paper). Why doesn't this happen with LAOM+supervision? Probably because the number of ground truth actions is extremely small in our case.

We thank the reviewer for their time and feedback. We have tried to address the concerns below.

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.

This 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.

It 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.

With 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.

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.

We 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.

W2) The effectiveness of the method. 

As 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).

W3) The scope of the proposed setting.

We 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.

Q1) Procgen results without distractors.

Our 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.

We thank the reviewer for their thoughtful and constructive feedback. We address the main questions below.

Missing this, which follows a pipeline similar to LAPA: IGOR…

We will include the citation, thank you for your suggestion.

… naturally a wider bottleneck will allow more information about both distractors and actions into the latent …. But is having extra distractor information not harmful?

As 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.

I'm surprised with this capacity the model doesn't just pipe the entire next observation through the latent

This 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.

I'd be interested to dive deeper into the differences between LAM and IDM

We believe that IDM will perform better than LAM in the limit (e.g. see Figure 9 in GR00T N1), 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.

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'

We 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.

Augmentation details

We 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.

Would love to see a zoomed in plot of fig 8 b as can't tell how the MSE changes for the LAOM+supervision.

Sorry, here's an enlarged Figure. We will add it to the appendix. Overall, the MSE does not change that much here.

For fig 10, what representation is used for the IDM?

The 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)$.

In fig 8 c, if it seems like the models keep improving with latent dim -- why stop at 8192?

Mostly 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.

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.

Clarification: is the BC policy also only trained on up to 128 trajectories?

Yes, 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.

We 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.

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?

Yes, this is an accurate description of the overall pipeline.

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?

Thanks 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, Figure 2].

As 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.

However, 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.

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?

We 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.

On 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).