Thanks for your review. We hope our extra work in providing new experiments that our results do carry over to practical LAMs (see the results in the response to the first reviewer) may warrant an increase in score.

[W1] It is not clear how results shown here carry over to practical LAMs. I think that is anyways the main point and motivation of this work. Text contains references to practical LAM papers in certain places where obtained results are similar to those in previous studies. This is good, but more effort would be needed here to make results practically useful. Appendix A is also too weak in this respect. Real difference is that dynamics model is not linear in practical LAM. About actions being Gaussian is a design choice by authors.

We agree that many real-world dynamics are not linear. This choice has been made to allow for tractable analysis and be compatible with linear LAM architecture (i.e., linear LAM can capture linear dynamics, and nonlinear LAM can capture more complex dynamics). Linear LAM design has two benefits: 1) Providing theoretical justification for tricks currently found to work well in practice (such as data augmentation and auxiliary action prediction). 2) Produces new insights into details affecting the learning of LAM (namely we noted the importance of data collection policy, which has not been mentioned in prior literature).

To evaluate how linear LAM captures the real-world LAM, we conduct further experiments as follows to validate that the main conclusion we obtained from linear LAM also applies to the setting where the dynamics is not linear and the action is not Gaussian. Please see the new experiment results in the response to the first reviewer (Reviewer VEZh).

[W2] How does linear LAM compare empirically to some practical LAM? Considering that paper already has quite a bit of simulation results it would not be a huge leap to add such a study to appendices.

The experiment results presented in the response to [W1] indicate that our main arguments still hold for practical LAMs.

[W3] Some results are just mentioned in the text, but no complete rationale is given in appendices, such as the difference between Eq. (5) and Eq. (6). Can Eq. (6) be interpreted as mutual information?

Eq. (6) approximates Eq. (5) by exploiting the connection between mutual information and the MSE loss between the target variable and the least squares estimate (LSE) of the target variable.

[Q1] In line 198, authors invert matrix B, but matrix B is not a square matrix. So you actually use pseudoinverse instead? -Section 4.1 Case 3, you only state that similar analysis of Case 1 can be conducted and what it results to. But I do not see reference to an appendix wheres such an analysis can be found.

Yes, we use pseudoinverse in Line 198. For Section 4.1 Case 3, we can define $\bar{X} = [X|Y] \in \mathbb{R}^{d_o \times (d_a + d_b)}$ (column concatenation) and $\bar{a} := [a^T|b^T]^T$ (row concatenation). Consequently, the results in Case 1 about $X$ and $a$ can be applied directly to $\bar{X}$ and $\bar{a}$ in Case 3. We will add the deduction to the revised appendix.

[Q2] Section 4.1 contains three cases, where Case 1 has no noise, Case 2 has stochastic noise and Case 3 has a systematic noise. Wouldn't it make sense to have fourth case, which has both stochastic and systematic noise?

This is a very good question. Stochastic noise increases the reconstruction loss but does not change the semantics of the learned latent. Systematic noise, when intense enough, can change the semantics of the learned latent. It is a bit hard for us to obtain a clear analytical solution when combining them together, but intuitively they will change both the loss and the latent semantics.

[Q3] Simulation results seem to produced with one seed, as there is stochasticity in the experimental procedure it is natural to ask whether different seeds to lead to different numbers? I am guessing that trends and conclusions would still hold.

We find that seed does not influence the results too much. To show this, we run vanilla linear LAM with $d_a = d_z = 8$ with 10 seeds and find that the standard deviation for linear LAM objective (LLO) is only 0.0173 (compared with the range for LLO [-2, +2]).

[F1] Line 226 "Gaussion" -> "Gaussian" -Grammatically sentences do not always parse. A bit more work in this respect is needed.

We will further improve the paper based on the reviewer’s suggestion.

Many thanks for your detailed review. We have aimed to address your concerns, including a new set of experiments bridging the gap between linear and practical LAM (see the results in the response to the first reviewer). We hope you might view this as sufficient to update the paper’s score.

[W1] Within the broader scope of LAM usage (mostly large-scale video data), it is unclear whether these claims remain impactful. No thorough investigation into this has been provided. The afore-mentioned weakness of the paper remains largely unaddressed in the paper. It would be great to see more justification for how this work may translate to more complex domains.

Thank you for this, which we agree was an important point not fully addressed in the initial paper. The gap between linear LAM and LAM usage in practice can be summarized as 1) linear to nonlinear model architecture, 2) continuous to sometimes vector quantized bottleneck; and 3) linear to complex world dynamics in video data (cf. Appendix A). To justify how our conclusions can be translated to complex domains, we conducted an additional set of experiments. Please see the new experiment results in the response to the first reviewer (Reviewer VEZh).

[W2] The writing could be substantially improved. It is often difficult to understand what is being said in the paper.

Apologies for not having had time to thoroughly edit the paper before submission. We commit to doing a full edit in the next version.

[W3] Figures could be improved for better readability. The y-axis should be labelled.

We agree. Currently the y-axis is only labeled in the title, we will revise.

[W4] At this point, it is unclear to me whether the analysis presented in the paper may be impactful for the community. My major concern is the vast oversimplification of the underlying model. Although the motivation for this oversimplification is clear (the model becomes analysable), it nevertheless generates a large gap between the results presented in the paper and the actionable significance of these results on the existing LAM methods. Primarily, can we trust these results to be applicable to complex, high-dimensional domains? To be clear, I do not necessarily think that they are not, rather it seems that the paper does little to address this concern in general. As such, it would be great to hear whether the authors believe these concerns could be addressed with more experiments in more complex settings, without taking away from the main claims of the paper.

Thank you, hopefully our experiments, shared above in W1 response, using nonlinear model architecture, vector quantized bottleneck, and image-based input, will help alleviate this concern.

[W5] The authors should provide more details / justifications as to how their analysis may be applicable. What are the limitations? Where, if at all, do the assumptions break down?

Apart from the new experiment results, the reason why linear LAM would catch the essence of LAM in practice is because it captures the key designs including IDM, FDM, bottleneck and reconstruction loss. The assumption that the dynamics is linear is made to match the model capacity of linear LAM, and we believe this can scale to the scenario where both are nonlinear. The limitation of linear LAM is that it cannot shed light on the LAM research on deep learning techniques (e.g., cross-attention vs self-attention in IDM), empirical data selection (e.g., the best proportion for each real dataset), or optimization (e.g., the optimizer or the learning rate schedule). We argue that the value of our paper is not we can obtain results based on assumptions that may not hold. Instead, our value is that the conclusions drawn from our analysis can be applied to the practice (e.g., using more exploratory data to train LAM) or explain the observations in practice (e.g., the data augmentation technique).

[Q1] I would appreciate if the authors could improve the writing in the paper; this would be essential for making it more accessible to the readers. (e.g. it is unclear to me what 'practical LAM' means when it is first introduced).

See W2 response. We will add a practical LAM definition.

[Q2] L147: It seems that minimising mutual information with an observation should be argued more. Isn't it possible to imagine the case where certain actions are only taken in specific states?

Minimizing mutual information with observation is still a proper objective, even if there are specific states where only certain actions can be taken. In this case, it is true that the mutual information between latent and state cannot be zero since we can infer the state from the latent. However, minimizing this mutual information still helps LAM to discover action representations to be as independent with states as possible so that it can be shared across states. Note that the ultimate objective for LAM is to learn shared action representation with consistent semantics (e.g., the latent could be “moving the arm left if not blocked”).

[Q3] L230: The formulation "of other agents" is a bit confusing. Wouldn't this better be phrased as predictable changes that are not being caused by agent's actions?

We agree this term could be changed to improve clarity. We propose to use use “predictable noise” in our revision – please let us know if this satisfies your concern.

[Q4] L279: "Therefore, the analysis in this section suggest us to use the data collected by the policies with higher randomness whenever possible." Although the theoretical component is somewhat insightful, this point seems rather trivial at first glance. If the policy is deterministic, a world model would naturally be capable of implicitly modelling the behavior of the agent?

Our main observation is that, when the action is easily predicted from the current state, there is less incentive for the latent to encode the action, with the FDM (or a world model as the reviewer mentioned) predicting the next state directly from the current state. Hence, we should prefer a higher entropy policy for the purpose of training a LAM. Like all good insights, as you note, this is intuitive, but we only made this discovery through our analysis, and do not believe related literature has made this observation.

[Q5] As mentioned in the Strenghts & Weaknesses section, I believe the authors should more clearly address the apparent gap between the presented results and the settings in which LAMs would most likely be applied (complex, high-dimensional settings; large-scale video data). If addressed convincingly, or if the authors can argue their case in the rebuttal, I would be happy to increase my score.

In response to the reviewer’s concern about the gap, we have provided additional experiment results in the response to the first reviewer (Reviewer VEZh)..

[Q6] IGOR paper claims: "In practice, we found that the video quality has a big impact on the model performance. We exclude low-quality videos characterized by excessive shakiness or rapid camera movement, and apply stabilization techniques to the remaining videos." I may be mistaken, but this slightly contradicts the conclusions of your paper? Could this be initial evidence that your claims do not translate to more complex settings?

No. Filtering data is an orthogonal technique from data’s perspective. Combining with the information that IGOR also uses data augmentation (advocated by our paper), we conclude that camera shakiness could be an issue and both the data augmentation (to help LAM robust to gentle shakiness) and data filtering (to avoid LAM suffer from excessive shakiness) are needed.

Thank you for taking the time to review our paper. We hope our responses, and a new set of experiments addressing the reviewers' concern about the gap from linear LAM to real (see the results in the response to the first reviewer) might justify an uplift in score.

[W1] The paper is mathematically heavy.

Thanks for this comment, since we would like to make the paper accessible to many readers. If the reviewer could mention any specific areas they thought were tricky, we’d be happy to add more intuitive explanations.

[W2] The analysis is based on simplified linear model with assumptions (Line 170)

We have studied the effect of this assumption carefully. Only the analysis in Section 4.1 is based on the assumption in Line 170 (noise and action are uncorrelated with observation). This assumption leads to perfect learnability of LAM. We later relax this assumption in Section 4.2 to show that stochastic data collection policy matters in reducing the correlation between action and observation. Also note that the noise correlated with the observation can be regarded as the signal.

[W3] The paper lacks empirical evaluations on real-world datasets

Thanks for this comment. At the request of the review, to provide evidence on whether our conclusions still hold on real-world datasets which are mainly image-based, we conducted a new set of experiments aimed at bridging the gap between linear LAM and real-world LAM. Please see the new experiment results in the response to the first reviewer (Reviewer VEZh).

[Q1] Is the analysis applicable when observations consist of multiple frames? Intuitively, capturing actions based on only two adjacent frames seems limiting.

Our analysis still applies when we simply redefine the vector observation $\bf{o}$ to contain the information about the history, corresponding to the setting in Genie and LAPO. For example, the learned latent can still capture $\bf{q} + \bf{\epsilon}$ in the idea case (cf. Proposition 4.1). The only difference is that the semantics of $\bf{q} + \bf{\epsilon}$ may be different. Intuitively, the latent may capture the “acceleration” instead of “velocity” when an object moves across multiple frames.

Moreover, even Genie and LAPO still aim at capturing movement between two frames (with additional input about the history). We find this may not be very limiting. Consider another option that extracts a latent to encode the change across multiple frames instead of extracting a latent between each two frames. However, we find that this formulation is not common in practice, maybe because such multi-frame movement can be easily represented as the concatenation of compact representations of two-frame movements or just the latent of two skip frames.

[Q2] All analysis relies on the additive transition function in Line 139. Is this assumption valid in practice?

This linear assumption is a simplification required for tractability. Since our linear LAM architecture is not as complex as practical LAMs, it can only handle additive transition dynamics. We believe that when LAM is more capable (as in practice), it can handle more complex dynamics. Our setting (linear LAM + linear dynamics) is a good approximation of the real case (nonlinear LAM + complex dynamics). Our new set of experiments give evidence that our insights hold in settings when dynamics are not linear.

[Q3] I understand in the ideal case when z contains no information about the noise and observation, the LLO value is 2 in Line162, but why this is the optimal value of LLO.

For a random variable $x$, $||\hat{x} – x||^2 / \text{Var}(x)$ is no smaller than 0 due to the non-negativeness of norm. It equals 0 when $\hat{x}$ is a perfect prediction of $x$. Another notable condition is that $\hat{x}$ is trained to predict $x$ to minimize the MSE loss appearing in the numerator. Since the worst prediction for $\hat{x}$ would be predicting the mean, this term will be no larger than 1. We will make it clearer in our revision.

[Q4] Some typos. Line 195, should be A=(1-alpha)o ; Line 323, missing lambda in the objective.

Thank you, we will fix the typo in Line 323. After checking we believe Line 195 is correct as stands, it should be $A=(1-\alpha) I$ where $I$ is the identity matrix, and thus $A \bf{o} = (1-\alpha) \bf{o}$ which can be used to cancel out the information $\alpha \bf{o}$ mixed in the latent.

We thank the reviewer for their valuable feedback. We are grateful that the reviewer acknowledged our paper as “one of the first works to provide a theoretical analysis of the learnability of LAM” with “thorough mathematical derivation” and “intuitive explanations”. We address the concerns as below, and hope these might be worthy of an upgrade in score.

[Simplified-to-real gap] Discuss the gap between simplified model and real-world architecture. How well does the simplified linear LAM model capture the behavior of real-world latent action models? Could you discuss the gap between this abstraction and practical architectures, and its impact on the applicability of your theoretical results?

Thank you for this important point – in order to better understand how linear LAM captures the real-world LAM, we have conducted further experiments to validate that the main conclusions we obtained from linear LAM also applies to the setting that bridges the three gaps between linear and practical LAM: 1) linear to nonlinear model architecture, 2) continuous to sometimes vector quantized bottleneck; and 3) linear to complex world dynamics (also see Appendix A). We present the new experiment results in the end.

[Paper writing] There are minor typos and the referencing issue.

Apologies for allowing these to pass in the first version. We have fixed these issues in the next version.

---

New Experiment Results

We provide a set of experiments using practical LAM to understand how our insights from linear LAM analysis transfer to real settings. Specifically, we use image as the input, non-linear CNNs in IDM and FDM, and vector quantization. We show that the main conclusions in our paper still hold in this complex setting.

Dataset

We designed a 4x4 grid-world style synthetic dataset. The top half of the observation contains a square (intensity=1.0) that can be controlled with 5 actions (up, down, left, right, and stay still). The bottom half of the observation contains random Bernoulli noise (with prob 0.5). An intensity parameter controls the noise magnitude (none=0.0, low=1.0, high=2.0).

Policy

By default, we use a uniform policy, where each action is equally probable. For one experiment, we also use a correlated policy, where state and action are correlated. With 95% prob, the action moves the square on a fixed snaking pattern through the grid, and with 5% chance a random action is selected.

Model

For the IDM, we use a small CNN to process $o$ and $o’$, followed by a VQ bottleneck with codebook size of 5 outputting the latent. Finally, for the FDM, a separate UNet takes the latent and previous observation $o$ to output the predicted $\hat{o}’$. When predicting actions, codes are preassigned to actions, and latents are trained to minimize L2 distance to their true action code. For data augmentation, we shift the 4x4 image left/right for one grid with periodic padding. Models were trained for 16k updates with the Adam optimizer.

Default setting

Unless specified, we use low stochastic noise, no action prediction, no data augmentation, and 5 codebook vectors.

Evaluation

Not being linear, it is not possible to measure the mutual information between latent and quantities of interest exactly. Instead, since by design the observation separates the controllable (top half) from stochastic (bottom half) changes, we measure the reconstruction loss on the relevant portion of the observation to assess how what information the latent has captured following training. A lower controllable loss means the latent contains more information about the action, and lower stochastic loss means more noise has been captured. Similar to LLO, this is normalized by the variance of the signal. Note there is unfortunately no straightforward way to measure the information about the observation $o$ in the latent for this non-linear case.

Results

Firstly, we train vanilla LAM on different noise levels. Results show that as noise is increased from none to high, the action information in the latent decreases (increasing controllable loss). At the same time, the information about the stochastic noise increases (decreasing stochastic loss). This is consistent with our linear LAM theory in Section 4.1, that latents encode whatever leads to most variance.

Secondly, we see that switching from a uniform to a correlated policy reduces the information about actions in the latents, while increasing the information about stochastic noise. This is predicted by our linear LAM theory in Section 4.2, that deterministic data collection policy can lead to degenerated performance.

Thirdly, we see that using data augmentation can improve LAM learning. This is consistent with the conclusion in Section 4.3, that data augmentation can help LAM learning.

Finally, we show that incorporating 1% of actions labels into the training process improves over vanilla LAM by increasing information about actions, and reducing information about noise. This is consistent with the conclusion in Section 4.4, that predicting true action can improve LAM learning.