Can a MISL Fly? Analysis and Ingredients for Mutual Information Skill Learning

We thank the reviewer for the responses and suggestions for improving the paper. The reviewer brought up two main questions about the paper: the interpretation and the effect of the coefficient $\xi$, and the importance of the inner product parameterization for the critic. We have attempted to address these questions by ablating different values of $\xi$ and explaining intuitions of the inner product parameterization respectively. We have also revised the paper to include those interpretations. Together with the discussion below, does this fully address the reviewer’s questions?

Where did the "fixed coefficient ξ = 5" come from/does it have any interpretation?

We introduce the coefficient $\xi$ to scale the negative term $\text{LB}\_{-}^{\beta}(\phi)$ in the contrastive lower bound and found that a proper choice of $\xi$ improves the performance of CSF empirically. This coefficient has an effect similar to the tradeoff coefficient used in information bottleneck optimization ($\beta$ in [1]). We also note that when setting $\xi \geq 1$, the new representation objective $\text{LB}\_{+}^{\beta}(\phi) + \xi \text{LB}\_{-}^{\beta}(\phi)$ is still a (scaled) contrastive lower bound on the mutual information $I^{\beta}(S, S’; Z)$ (see Appendix A.6 for further discussions).

How sensitive is CSR to this hyperparameter?

To answer this question, we ran ablation experiments to study the effect of the scaling coefficient $\xi$ on the representation objective of CSF empirically. We conducted experiments on the $`Ant`$ again, selected a set of different values of $\xi \in \\{ 0.5, 1, 2, 5, 10\\}$, and report means and standard deviations of coverage of $(x, y)$ coordinates over 5 random seeds. Results in Fig. 10 (Appendix C.3) suggest that while the state coverage saturates after increasing $\xi$ to $2$, our choice of $\xi = 5$ is optimal for CSF.

… an ablation study shows that this is a crucial choice. Can the authors provide a further comment on the importance of the temporal distance metric?

Yes. We used inner product parameterization as our distance metric for the critic and this is inherited from METRA. Unlike METRA, our method does not explicitly enforce a $L^2$ temporal distance constraint and instead optimizes the contrastive lower bound on $I^{\beta}(S’, S; Z)$. Our experiments in Fig. 3 (Right) show that the inner product parameterization is important for the good performance of CSF. There are two explanations for these observations. First, prior work [2, 3, 4] has mentioned that mutual information $I(S; Z)$ is invariant to bijective mappings on $S$ (See Fig. 2 of [3] or Fig. 2a of [4]), e.g., translation and scaling, indicating that maximizing the mutual information between changes of states and skills ($I(S, S’; Z)$) might encourage better state space coverage. Second, the inner product parameterization allows us to analytically compute the second-order Taylor approximation in Proposition 2, drawing the connection between the METRA representation objective and contrastive learning. We have revised the paper (Appendix A.9 blue texts) to include these insights.

it is unclear how well these algorithms scale beyond the relatively simple MuJoCo benchmarks

We have already started to investigate applying CSF to MiniHack [5], a simplified version of the immensely complex NetHack [6]. Specifically, we used the MiniHack-Corridor-R5-v0 environment which has a discrete action space and requires the agent to navigate between 5 different rooms using various doors (that require opening) and corridors. We have also implemented a version of METRA for these tasks. Our preliminary results suggest that both METRA and CSF perform similarly after the first 1M env steps, and both explore better than random. Specifically, they both manage to open doors and explore some of the corridors, sometimes even reaching the goal location. We're currently debugging some stability issues for both methods that arise later in training. We'll be sure to include the final results in the final paper.

[1] Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. In International Conference on Learning Representations, 2017.

[2] Park, Seohong, et al. "Lipschitz-constrained unsupervised skill discovery." International Conference on Learning Representations. 2022.

[3] Park, Seohong, Oleh Rybkin, and Sergey Levine. "METRA: Scalable Unsupervised RL with Metric-Aware Abstraction." The Twelfth International Conference on Learning Representations.

[4] Park, Seohong, et al. "Controllability-Aware Unsupervised Skill Discovery." International Conference on Machine Learning. PMLR, 2023.

[5] Samvelyan, Mikayel, et al. "Minihack the planet: A sandbox for open-ended reinforcement learning research." arXiv preprint arXiv:2109.13202 (2021).

[6] Küttler, Heinrich, et al. "The nethack learning environment." Advances in Neural Information Processing Systems 33 (2020): 7671-7684.

We thank the reviewer for the responses and suggestions for improving the paper. It seems like the reviewer’s main question is about the contribution of the paper. Our paper mainly focuses on (1) providing an alternative explanation of METRA theoretically and empirically within the well-studied mutual information skill learning framework, and (2) proposing a simpler algorithm to match the SOTA performance. These contributions are well within scope for ICLR, and ICLR has a history of recognizing papers that make similar contributions [1, 2]. Together with the discussion below, does this address the reviewer’s concerns?

I also find the argument that the method removes 5 hyper-parameters compared to METRA to be bad faith … Could you clarify which hyper-parameter actually require tuning and how many hyper-parameters you introduce?

Thanks for mentioning this; we agree that counting hyperparameters is subjective. For now, we have removed the exact quantification from the main paper and instead said “simpler”; we have also added a new discussion in Appendix B.1 to discuss the hyperparameters that we have removed (and the one hyperparameter $\xi$ that we’ve added) in more detail. We would love to get the reviewer's feedback on this, and would happily consider alternative ways of presenting this comparison.

The statement line 278 does not hold for all cited methods (like DIAYN), as the "omitted" is sometimes a useless constant.

Thanks for mentioning this. We have revised this sentence to clarify (blue texts). The purpose of this sentence is to contrast that these prior methods use the same function form of the lower bound (Eq. 2 \\& 3) on different variants of mutual information to learn both representations and policies, while METRA uses different functional form of objectives for the representation and the actor. We elaborate which mutual information each prior method tries to optimize in Appendix A.5 (blue texts).

The final model is very close to previous work and do not present substantial improvements on the different environment compared to METRA.

Our main contributions are not to surpass METRA in terms of performance, but rather the things listed at the top of this response (i.e. explaining METRA within MISL, and proposing a simpler algorithm that matches SOTA).

Nevertheless, to get further insight into the behavior of CSF and METRA, we have already started to investigate applying CSF to MiniHack [3], a simplified version of the immensely complex NetHack [4]. We used the MiniHack-Corridor-R5-v0 environment which has a discrete action space and requires the agent to navigate between 5 different rooms using various doors (that require opening) and corridors. We have also implemented a version of METRA for these tasks. Our preliminary results suggest that both METRA and CSF perform similarly after the first 1M env steps, and both explore better than random. We're currently debugging some stability issues for both methods that arise later in training. We'll be sure to include the final results in the final paper.

[1] Higgins, Irina, et al. "beta-vae: Learning basic visual concepts with a constrained variational framework." International conference on learning representations. 2017.

[2] Arora, Sanjeev, et al. "A simple but tough-to-beat baseline for sentence embeddings." International conference on learning representations. 2017.

[3] Samvelyan, Mikayel, et al. "Minihack the planet: A sandbox for open-ended reinforcement learning research." arXiv preprint arXiv:2109.13202 (2021).

[4] Küttler, Heinrich, et al. "The nethack learning environment." Advances in Neural Information Processing Systems 33 (2020): 7671-7684.

Thanks again for your review and further suggestions. We are running ablation experiments studying the sensitivity of METRA with respect to $\lambda$, $\epsilon$, and the learning rate of the dual variable optimizer, helping to understand the simplicity of CSF versus METRA. We will be sure to include discussions about these in the final version.

We thank the reviewer for the response, and for the suggestions for improving the paper. To address the concerns about experiments, we have run additional experiments to ablate different distance metrics / kernel functions for the critic. Together with the discussion for other questions below, does this fully address the reviewer’s concerns about the paper?

Are there any theoretical insights supporting the assertion that parameterization is key for contrastive successor features?

Yes. Our experiments show that the inner product parameterization of the critic $f(s, s’, z) = (\phi(s’) - \phi(s))^{\top} z$ is important for CSF (Fig. 3 (Right)). There are two explanations for these observations.

First, the inner product parameterization tries to push together the difference of the representation of transitions $\phi(s’) - \phi(s)$ and the corresponding skill $z$, instead of focusing on representations of individual states $\phi(s)$ or $\phi(s’)$. This intuition is inline with the observation that mutual information $I(S; Z)$ is invariant to bijective mappings on $S$ (See Fig. 2 of [1]), e.g., translation and scaling, indicating that maximizing the mutual information between change of states and skills ($I(S, S’; Z)$) encourages better state space coverage.

Second, the inner product parameterization allows us to analytically compute the second-order Taylor approximation in Proposition 2, drawing the connection between the METRA representation objective and contrastive learning. We have revised the paper to include these theoretical insights (see Appendix A.9 blue texts).

Did you experiment with different kernel functions?

As suggested by the reviewer, we ran additional experiments to study the effect of different distance metrics / kernel functions for the critic in Fig. 3 (Right). We selected $L^1$ distance (Laplacian kernel) $-\\| \phi(s') - \phi(s) - z \\|_1$ and squared $L^2$ distance (Gaussian kernel) -\frac{1}{2} \left \\| \phi(s') - \phi(s) - z \right \\|_2^2 in addition to the inner product (vMF kernel) and the monolithic MLP parameterizations of the critic to conduct ablation studies. For the sake of consistency, we chose to use the $`Ant`$ task and report means and standard deviation of coverage of unique (x, y) coordinates over 5 random seeds. Results in Fig. 3 (Right) suggest that the inner product critic is the best for CSF among different choices of parameterizations.

It would be advantageous to include demonstrations or other forms of visualization for the skills learned, as I did not find this in the appendix code.

Thanks for the suggestion. We visualized skills learned by CSF and METRA (anonymous videos in https://anonymous.4open.science/r/csf-3BF4/README.md under “Videos of learned policies” section) and have revised the paper to include visualizations of $(x, y)$ trajectories of these skills in Appendix C.4 (blue texts).

A significant question arises especially in the Quadruped experiments … why does the empirical performance (or at least the rate of convergence) not exceed that of METRA?

Our hypothesis for this observation is that METRA might overfit to some of the tasks. The reason for this is that METRA uses dual gradient descent to optimize the representation objective, resulting in 2 more hyperparameters than CSF (see Appendix B.1). While using more hyperparameters benefits the flexibility of METRA [2], it also means that the algorithm tends to overfit especially when the effective dimensions of a task are low, e.g., $(x, y)$ in Quadruped.

We note that CSF only slightly underperforms METRA (within one standard deviation confidence interval) in terms of both the average state coverage and the rate of convergence. As supported by our theory, the main practical benefit of CSF is not that it should achieve better results on every task but rather that it is easier to implement (see algorithm summary in Sec. 5.2 and Appendix B.1).

We thank the reviewer for the responses and helpful suggestions for improving the paper. The reviewer mainly raised two questions about the paper: evaluation on diverse benchmarks and conditions under which the theoretical analysis holds. We have attempted to address the question about evaluations by starting to run experiments on MiniHack benchmarks [1], and we answer the question about conditions and assumptions for theoretical analysis below. Together with the discussion below, does this address the reviewer’s concerns?

… it would also be nice to show where (if anywhere) one fails when the other succeeds. Perhaps partially observed MPDs, more interactive objects or discrete actions spaces would be key in identifying where exactly both methods stand.

We have included the full results of the state space coverage in Appendix B.3 and Fig. 5 shows that CSF matches the prior SOTA performance on 4 / 6 tasks. Specifically, CSF outperforms METRA on the $`Robobin (Pixels)`$ while slightly underperforming METRA on the $`Quadruped (Pixels)`$.

In addition, as suggested by the reviewer, we have started to investigate applying CSF to MiniHack [1], a simplified version of the immensely complex NetHack [2]. Specifically, we used the MiniHack-Corridor-R5-v0 environment which has a discrete action space and requires the agent to navigate between 5 different rooms using various doors (that require opening) and corridors. We have also implemented a version of METRA for these tasks. Our very preliminary results suggest that both METRA and CSF perform similarly after the first 1M gradient steps, and both explore better than random. We're currently debugging some stability issues for both methods that arise later in training. We'll be sure to include the final results in the final paper.

Can the authors elaborate on the assumptions made in the theoretical analysis, and how they might affect the generalizability of the results?

Our theoretical analysis mainly builds on four conditions (assumptions) which are inline with empirical findings in the didactic experiments. Below, we discuss these assumptions in detail.

First, when casting the mutual information maximization problem into a min-max optimization problem, we assume the realizability of the variational function class $Q = \\{q(z \mid s, s’)\\}$, i.e., $Q$ is expressive enough to cover the ground truth posterior $p^{\pi}(z \mid s, s’)$. Many prior MISL methods [3, 4, 5] adopt the min-max optimization procedure and implicitly apply this assumption.

Second, our analysis focuses on the case where $p(z)$ is distributed as a uniform distribution on the $d$-dimensional unit hypersphere, as is common in much (but not all) prior work [3, 6]. Using this prior distribution helps us to unify analysis for both the representation objective and the policy objective in METRA.

Third, in Proposition 2, we claim that the constraints used in the METRA representation objective correspond to a second order Taylor approximation of the negative term in the contrastive lower bound ($\text{LB}_{-}^{\beta}(\phi)$). However, since METRA uses dual gradient descent to adaptively update the dual variable $\lambda$, the tightness of the Taylor approximation depends on the value of $\lambda$. We implicitly assume that the Taylor approximation holds and the METRA representation objective has the same effect as contrastive learning. In practice, our experiments in Fig. 8 and Appendix C.1 verify that this assumption is reasonable. Note that we also apply this Taylor approximation in the proof of Proposition 3 (step (c)).

Finally, we apply Proposition 1 in the proof of Proposition 3 (step (d)), which implicitly assumes that the optimization problem in Eq. 7 is solved without errors. We also empirically verify this condition in the didactic experiments Fig. 2a and Sec. 6.1. We have revised the paper accordingly (Sec. 3 and Appendix A.4 blue texts) to clarify these conditions.

[1] Samvelyan, Mikayel, et al. "Minihack the planet: A sandbox for open-ended reinforcement learning research." arXiv preprint arXiv:2109.13202 (2021).

[2] Küttler, Heinrich, et al. "The nethack learning environment." Advances in Neural Information Processing Systems 33 (2020): 7671-7684.

[3] Eysenbach, Benjamin, et al. "Diversity is All You Need: Learning Skills without a Reward Function." International Conference on Learning Representations.

[4] Hansen, Steven, et al. "Fast Task Inference with Variational Intrinsic Successor Features." International Conference on Learning Representations.

[5] Gregor, Karol, Danilo Jimenez Rezende, and Daan Wierstra. "Variational intrinsic control." arXiv preprint arXiv:1611.07507 (2016).

[6] Park, Seohong, Oleh Rybkin, and Sergey Levine. "METRA: Scalable Unsupervised RL with Metric-Aware Abstraction." The Twelfth International Conference on Learning Representations.

Reviewer 8yQh: if you can, please reply the rebuttal.