Paper: $\lambda$-AC: Effective decision-aware reinforcement learning with latent models (ICLR-2024)

Reviewer comment (Writing):
While the paper is well-written overall, there are some technical details and notations that can make key formulas harder to parse (see the minor comments below). For a paper with relatively heavy notation, as it needs to consider ground-truth models, estimated models, transitions collected from interactions, and transitions from learned models, such mistakes, while understandable, can make some key concepts hard to grasp. As someone with little prior knowledge on MuZero and IterVAML, I found the mathematical details behind both losses hard to understand.

Author response:
We thank you for your thorough review and your detailed suggestions! We acknowledge that there are several typos that make key concepts harder to grasp, and we have used your feedback to update the manuscript. As a general rule of thumb, we aim to use superscript (i) to denote an i-th sample, while we aim to use superscript i without brackets to denote power operations such as chained autoregressive model applications. The places you pointed out are indeed inconsistencies in this usage.

Reviewer comment (Presentation):
The legends and labels in the figures are slightly confusing. The paper's text seems to use $\lambda$-MuZero and $\lambda$-IterVAML to refer to the variants of the two losses under the $\lambda$-AC framework, and MuZero and IterVAML to refer to the "vanilla" versions of the algorithms. However, I cannot find $\lambda$-MuZero nor $\lambda$-IterVAML in figures 4 - 7, which seems to suggest that the experiments are done on the "vanilla" algorithms themselves. Is this the correct interpretation?

Author response:
> “The legends and labels in the figures are slightly confusing.” 

This is an oversight on our part, we only compare the lambda variants in this paper. We will change the labels to reflect this.

Reviewer comment (Theory):
I am slightly concerned by the correctness of eq (1) and eq (2) as written here. Particularly, eq (1) takes expectation over some $\mu ∈ \Delta(S × A)$ (even though here the paper lets $\mu$ be a distribution over the state space only), whereas in eq (2) the samples follow a distribution where the initial state $x\_{i\_1}$ is drawn from some $\mu ∈ \Delta(S)$, and actions $a$ and subsequent states are drawn from the ground-truth transition kernel $p$ (which generates $\mathcal{D}$). Without any assumption on $\mu$, I am not sure if (2) provides an unbiased estimate of (1). Note that the cited work (Farahmand, 2018) also seems to take expectation over some joint distribution over state and action.

Author response:
> “I am slightly concerned by the correctness of eq (1) and eq (2) as written here.” 

The expectation is indeed meant to be taken over the state-action space, we have fixed this omission.

Reviewer comment (Theory):
Shouldn't the sample-based versions of MuZero and IterVAML depend on some policy $\pi$ as well? Aren't the actions $a\_{i\_j}$ are collected by some particular policy $\pi$?

Author response:
> “Shouldn't the sample-based versions of MuZero and IterVAML depend on some policy as well?” 

Both are written for an off-policy case, so actions can be taken from any number of policies. We clarify this in the updated draft. The off-policyness is not as relevant for the model loss as it would be for a value function loss, therefore previous work has not investigated this in detail and we follow this convention.

Reviewer comment (Theory):
I am not familiar with either MuZero or IterVAML, so my assumption may be misguided. However, shouldn't it be expected that MuZero will always have some bias? From equation (3), my understanding is that the loss function uses some deterministic mapping to estimate the transition. As such, the assumption is inherently not compatible with stochastic transition. Is this intuition oversimplifying the problem? (Of course this is not to say the theoretical results are trivial or not interesting. Rigorous proofs cannot be replaced by hand-waving "analysis".)

Author response:
> “However, shouldn't it be expected that MuZero will always have some bias?”

This is a very good question: in part, this is the (somewhat surprising) result of our first proposal: if value awareness is the goal, a deterministic model can be sufficient and no stochastic transition model is needed. Intuitively the reason is that instead of aiming to predict a next state distribution, we are only interested in the expectation of the value function, which is a simple scalar. This scalar is all that is needed for value iteration or TD learning. The proof highlights that this scalar can be predicted with a deterministic function.

In addition, proof 2 highlights that even with a stochastic model, the MuZero value function learning algorithm will be biased (potentially quite badly). Note that our proof places no assumptions on the implementation of the model. This is the second highly surprising result, as merely fixing the model architecture is insufficient to fully adapt MuZero and descendant algorithms to stochastic environments. Our tabular experiments in the appendix highlight this fact: we use stochastic models there (just with constrained capacity), and the MuZero algorithm is still incapable of learning the correct value function.

You are completely correct to point out that MuZero (as well as almost any other algorithm in the field) has some bias due to the choice of model class, irrespective of whether stochastic or deterministic models are used. However, the bias we show is separate from expected bias-variance tradeoffs and is irreducible with better models or more samples. The problem we highlight exists _in addition_ to the regular approximation error that any model class incurs when predicting a complex function, which makes it unique from the problems that arise from simply picking the wrong function class for the problem.

Reviewer comment (Novelty):
The paper is weak in significance because it mainly focuses on analyzing two existing value-aware losses and obtains a few insights that only apply to the two specific losses.

Author response:
We focus on these two losses as they are the main losses considered in the literature on value-aware/value equivalent model learning. Other approaches do exist, and a wider comparison is indeed interesting, but we decided that the “core” approaches deserved a closer analysis. We would also like to point out that MuZero is not a niche approach, but a landmark approach in model-based RL with almost 2000 citations in the field. Its loss has been used in several follow up works and can be considered a “backbone” component of the field, comparable to Dreamer. We therefore think that our results are highly relevant also for the many other algorithms building on MuZero.

Reviewer comment (Theory):
Regarding soundness, the claim that “$\lambda$-IterVAML leads to an unbiased solution in the infinite sample limit, [conditions]…” is not well justified. Proposition 1 only shows the existence of such an unbiased solution, and it is not immediately apparent that $\lambda$-IterVAML *leads to* it. If this is an implied result, clarification is needed.

Author response:
The proof indeed only shows existence and we point this out below proposition 1 in the main text. Showing convergence is a significantly harder problem which we decided not to tackle in this paper. Even for well established algorithms like MLE, we cannot obtain general convergence results with deep models. We still think the inclusion of this proof is relevant to the community because it shows that the goal of learning value-aware models, even with deterministic backbones, is not impossible. This result has been a core argument used in the MuZero line of work, starting from Value Prediction Networks, but no proof of the conjecture existed in the literature so far. The necessity of invoking complex measure theoretic topics highlights that this is a highly non-trivial result.

In our opinion, the necessity of the stabilizing loss points to the fact that the learning problem is hard and convergence is not directly implied by existence here. Further analysis is indeed necessary, yet out of scope for this work.

Reviewer comment (Evaluation):
In the caption of Figure 6, the explanation that the performance decrease of IterVAML is due to the lack of real reward signal in the value function loss is not supported by evidence.

Author response:
We agree that this is a conjecture that is not fully proven by the evidence and we should have been more careful in the caption. We will clarify the claim to be in accordance with the main text. To be clear: this is our working hypothesis as to why the difference between the different approaches was so drastic, even though we expected them to behave very similarly in this experiment.

Reviewer comment (Theory):
Question — Are there results similar to Proposition 1 when $\mathcal{X}$ is a discrete space? If not, how likely could there be such a result?

Author response:
This is an excellent question, and the answer is no, not immediately. Intuitively the problem is that the average value over a set of discrete points does not have to correspond to the value of any single one of them, so no deterministic mapping is possible. The main technical problem stems from the fact that no meaningful continuous topology can be defined on a discrete space. However the introduction of a continuous representation space “fixes” this. If the discrete space is represented as a set of points embedded in a continuous space (i.e. by encoding each state as a one-hot vector and allowing the model to predict continuous vectors) in which the value function is also continuous, that space can be used to obtain our proposition. We will add this note to the appendix with more formal language.

Reviewer comment (Writing):
This paper assumes too much background on the reader. It uses jargon without clearly and sufficiently introducing them, for example, latent model, decision-aware learning framework, IterVAML, MuZero, BYOL loss, and so on. Most importantly, it is very hard to figure out what is the contribution of this paper. Both the introduction and the conclusion did not clearly point the main contribution out.

Author response:
> This paper assumes too much background on the reader. It uses jargon without clearly and sufficiently introducing them

We are happy to improve our writing. We present background descriptions in Section 2 and would love to hear what parts are not sufficiently well explained? We have added additional text explaining the difference between observation-space and latent space models. If additional clarification is needed, we are happy two write an expanded background section in the appendix.

> Both the introduction and the conclusion did not clearly point the main contribution out.

Regarding the contribution, the final 4 paragraphs of the introduction as well as the conclusion list our main insights. We have updated the draft to put more clear emphasis on these:

_In this work, we ask three major question: *(a)* what design decisions explain the performance difference between IterVAML and Muzero, *(b)* what are the significant theoretical or practical differences between the two, and *(c)* in what scenarios do decision-aware losses help in model-based reinforcement learning?_

_[...]_

_In summary, our contributions are *a)* showing that IterVAML is a stable loss when a latent model is used, *b)* proving and verifying that the MuZero value loss is biased in stochastic environments, and *c)* showing how these algorithms provide benefit over decision-unaware baseline losses._

Reviewer comment (Writing):
The readability of this paper could be greatly improved by deleting unnecessary words and sentences. More tables should be introduced in place of large paragraphs of words.

Author response:
Dear reviewer, thank you for your review. We will post a more detailed answer soon addressing your feedback. To properly engage with the points raised by your review, we would quickly like to ask for some clarifications. Which sentences did you feel did not contribute to your understanding of the paper, and what tables or figures would be helpful to replace these? We look forward to hearing from you.

Reviewer comment (Writing):
The author should directly articulate their research goal at the beginning of the research paper. Currently, readers cannot understand the research goal until the first full pass of the paper.

Author response:
> The author should directly articulate their research goal at the beginning of the research paper.

We present the main contributions in the introduction, with a summary in the last paragraph. We have updated the writing accordingly (note the quote in the previous paragraph). We are happy to include suggestions of things that could be added to make the goal clearer?

Reviewer comment (Novelty):
In terms of contribution, a comparison among three different loss functions in three environments may not be significant enough to offer strong insights. And the novelty of this work is limited because it is a direct completeness extension and evaluation of previous works cited in section 2.1.

Author response:
> A comparison among three different loss functions in three environments may not be significant enough to offer strong insights.

We will answer this together with the following question on why we focus on IterVAML and MuZero, and do not compare against Dreamer: We purposefully chose VAML and MuZero because they are the main approaches in the family of value equivalent learning. Our paper is not meant (and not introduced) as a survey on model-based RL, but as a comparison of value-aware/value-equivalent approaches. BYOL has been used in several papers as an additional stabilizing loss, so it is a natural ablation. It could also be replaced by Dreamer, but Dreamer or other models have not been adapted in the framework of decision-aware learning before. Our insights are specifically targeted at the two major approaches to the value equivalent model family. 

> And the novelty of this work is limited because it is a direct completeness extension and evaluation of previous works cited in section 2.1

We want to gently push back against the claim that there is limited novelty in a thorough evaluation. Our work is in line with similar papers at top conferences (such as “Hamrick et al. On the role of planning in model-based deep reinforcement learning” published at ICLR 2021), which aim to understand the problems of existing methods instead of proposing additional ones. Our proofs are new contributions to the literature and touch upon previously unknown *and completely unrecognized* issues in a landmark algorithm in the field in addition to highlighting a possible solution. We do not merely extend MuZero, we show that it is an incorrect algorithm in a large number of important scenarios.

We are working on providing additional experiments in the appendix to broaden our basis of comparison.

Reviewer comment (Experiments):
Why do you choose IterVAML, MuZero, and BYOL loss functions as benchmarks to compare? Are they broad enough to give a representative comparison of model-based RL methods?

Author response:
> A comparison among three different loss functions in three environments may not be significant enough to offer strong insights.

We will answer this together with the following question on why we focus on IterVAML and MuZero, and do not compare against Dreamer: We purposefully chose VAML and MuZero because they are the main approaches in the family of value equivalent learning. Our paper is not meant (and not introduced) as a survey on model-based RL, but as a comparison of value-aware/value-equivalent approaches. BYOL has been used in several papers as an additional stabilizing loss, so it is a natural ablation. It could also be replaced by Dreamer, but Dreamer or other models have not been adapted in the framework of decision-aware learning before. Our insights are specifically targeted at the two major approaches to the value equivalent model family.

Reviewer comment (Writing):
The framework $\lambda$-AC seems to be vague. Towards the end, when the authors discuss about using model for policy learning or not, both of these will fall under this framework as per the definition: "and an actor-critic algorithm to obtain policy".

Author response:
> “The framework -AC seems to be vague.”

We use the term lambda-AC as a way to highlight similarities, not differences, between the presented algorithms. It is correct that both variants of the algorithms presented fall into the same framework, several other extensions do as well. We point this out in detail in Section B in the appendix.

We acknowledge that the term can lead readers to believe that $\lambda$-AC is intended as a novel algorithm. We have updated the writing accordingly to clarify our usage of the term.

Reviewer comment (Experiments):
If presenting an evaluation paper, why not compare more model-based methods like Dreamer and using discrete settings as well.

Author response:
> If presenting an evaluation paper, why not compare more model-based methods like Dreamer and using discrete settings as well.
We purposefully do not compare to Dreamer or many other model-based RL methods, because they do not fall into the family of value equivalent learning and therefore insights are not naturally transferable to these other cases. Our paper is not meant (and not introduced) as an evaluation or survey on model-based RL, but as an investigation of value-aware/value-equivalent approaches, of which we focus on the two most prominent examples.

Reviewer comment (Writing):
A small preliminary on model based value gradients like SVG should be presented.

Author response:
The second paragraph of 2.2 addresses SVG, however, the term is not mentioned explicitly there, only the eponymous paper is. We corrected that oversight in the updated draft.

Reviewer comment (Experiments):
When it is established in the first experiment, that the auxiliary loss does not add the MuZero, why is it still used as it is adding additional bias. What will happen if I use the MuZero directly (without auxiliary loss) in Section 4.3?

Author response:
We are working on expanding our experimental evaluation with a test of using the MuZero loss without stabilization in other environments. Preliminary results show that, due to the similarities with the IterVAML loss, we do not obtain good performance this way. Note that experiment number 1 tests the removal of the latent encoder, not the removal of the stabilizing loss.