Paper: How the Level Sampling Process impacts Zero-Shot Generalisation in Deep Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
It would be also great if the authors could further provide more analysis or even theoretical analysis of the hypothesis.

Author response:
We thank the reviewer for his feedback and positive comments on our paper. We hope they will appreciate even more the new and improved section 4! In particular, we point to sections 4.1 and 4.2 which provides a more rigorous analysis of our hypothesis.

We believe the reviewer's points have been addressed in our post titled titled "Summary of changes in updated version" which gives a comprehensive overview of the changes made in the new revision. We have fixed the notation issues pointed out by the reviewer, and we hope the structure added to section 3 makes it now easier to follow.

Reviewer comment (Evaluation):
The paper overpromises many things, for example from the abstract:  
- "As a first step, we measure the mutual information (MI) between the agent’s internal representation and the set of training levels, which we find to be well-correlated to instance overfitting": there is no empirical measurement of MI or its correlation with overfitting in the paper. All evidence is circumstantial and the interpretation of Figure 2, in this reviewer's opinion, wrong.  
- "adaptive sampling strategies prioritising levels based on their value loss are more effective at maintaining lower MI", "We then turn our attention to unsupervised environment design (UED) methods, which [..] minimise MI more effectively than methods sampling from a fixed set.", "SSED generates levels using a variational autoencoder, effectively reducing MI": again, MI is never measured and the claim that the evaluated methods minimize MI is pure conjecture.

Author response:
1. `there is no empirical measurement of MI or its correlation with overfitting in the paper.` and `MI is never measured`.
In section 4.3 we include an additional result measuring the correlation of our estimate for $MI(L,b)$ with the generalisation gap and with the value loss, with additional analysis conducted in appendix B.1

Reviewer comment (Writing):
The paper is confusingly written and introduces many terms without defining them. For example, $\text{MI}(i,\pi)$ is never defined in the main paper. $S^V_i$ is defined with $\hat V_t$ and $V_t$, which have never been defined. On Page 5 the authors state $V_i^\pi(s) = V^\pi(s) + v_i^\pi(s)$, but $V^\pi(s)$ has never been defined (only for a general POMDP, the CMDP only knows context-specific values called $V_{i_x}^\pi$). Sometimes the used PPO baseline uses non-recurrent intermediate representations (p.5), another time the same implementation uses an LSTM-based architecture (which is recurrent, p.7). The reviewer has the impression that many of these details could be understood upon reading the cited papers, but as is, the paper does a poor job explaining them.

Author response:
2. We thank the reviewer for pointing out these missing definitions, which are now properly introduced in the new revision. Regarding recurrent and non-recurrent representations, we now clarify how our analysis is equally applicable to recurrent and non-recurrent architectures in section 4.1.

Reviewer comment (Theory):
The connection between generalization, MI and scoring the value loss ($S^V$) is indirect, and not rigorous enough. The connection between MI and generalization is not very clear: in Figure 1 both following the black and the green cells leads to the goal in (a) and (b), but only one of the two generalizes to (c). Both have thus the same MI, but different generalization. While the statement "An agent learning level-specific policies implies high MI between its internal representation and the level identities" (p.1-2) makes generally sense, the conclusion that agents with high MI "will not transfer zero-shot to new levels" is less clear. Moreover, it is not apparent why the negation *agents with low MI generalize well*, which seems to be the basis for this paper, should be true. So the only clear connection seems to be the bound in (eq.2). The reviewer also missed a connection why "Sampling levels associated with highly negative classifier cross-entropies therefore results in less mutual information" (p.5). Finally, the idea that $V_i^\pi(s) = V^\pi(s) + v_i^\pi(s)$ separates into two errors and the second is somehow connected to MI (without any clarification how exactly) is flawed, as the errors of $V^\pi$ and $v_i^\pi$ could also cancel each other out, yielding low error, but high MI.

Author response:
3. 
a) We hope our restructuring of section 4 helps in understanding the connection between $S^V$, $MI(L,b)$ and the generalisation gap. We agree with the reviewer that agents with high MI will not always generalise poorly, as MI characterises an *upper bound* on the generalisation gap, while the gap itself may be small if that bound is not tight. However, our experiments results reported in section 4.3 and appendix B.1 and B.2 indicate that, in practice, high MI does imply a significant generalisation gap.

b) To clarify, we are not making the claim that "agents with low MI generalise well". We claim that "agents with low MI have a small generalisation gap". This does not always imply strong generalisation. For example, achieving train and test scores of zero yield a zero generalisation gap, however the agent cannot be described as generalising well. To an extent, this example is an extreme version of what happens under the $S=S^{MI}$ sampling strategy, where the generalisation gap is reduced by adversely impacting the performance on the train set. We hope the updated section 4, and in particular section 4.3 and the updated Figure 2 are now clearer in explaining this failure mode induced by certain mutual information minimisation strategies.

c) Section 4.2 now focuses exclusively on motivating the $V_i^\pi(s) = V^\pi(s) + v_i^\pi(s)$ decomposition and its relationship to MI. the unbiased value estimator lemma from Bertran et al. (2020), included as Lemma 3.1 in the new revision, states that $\mathbb{E}_L[V^\pi_i]$ is an *unbiased* estimator for the CMDP value function (redefined as $V^\pi_C$ in the new version for added clarity) under any set of levels $L$. As $\hat{V}$ is trained to predict $V^\pi_i$ under $L$, at convergence it is an unbiased estimate of $V^\pi_C$. Finally, the bias-variance trade-off tells us $\hat{V}$ will tend to have low bias but high variance as it starts overfitting. In fact, we may informally view $v^\pi_i$ as a level-specific "noise function". In order to accurately predict this "noise function" $v^\pi_i$, $\hat{V}$ needs to be learn to identify level identities, as different levels have different $v^\pi_i$. Hence $MI(L,b)$ becoming high as the value loss goes to zero.

Reviewer comment (Theory):
The approximation in Lemma 4.1 only works if the levels are constantly sampled with $p(i)$, but your method aims to change this distribution. How does this work when your scoring selects the levels that make up the batch $B$?

Author response:
1. Prioritised sampling effects the model by changing the training data. However the resulting trained model must be evaluated over the context distribution associated to the CMDP and MI also must be measured over this distribution. In procgen, all train seeds correspond to in-context levels and can be considered equally likely in the CMDP contextual distribution. Therefore the empirical estimation for that distribution is the distribution giving all training levels equal weights. We emphasise this distinction in section 4.1 of the new revision.

Reviewer comment (Theory):
Your goal is to produce policies that have low MI on the training set, but approximated MI is not used to change the policy, only the sampling of levels. Is there a reason why you do not simply reward actions with lower MI?

Author response:
2. In this work, we chose to focus specifically on the level sampling process as we view it as a form of data regularisation that is understudied when compared to the number of existing contributions tackling ZSG through modifications of the reward function or the agent architecture, or by augmenting the observations. However, as our method is orthogonal to such approaches there may be interesting synergies to explore in future work.