MIMEx: Intrinsic Rewards from Masked Input Modeling

Summary

The paper introduces a novel method for Exploration in RL, called MIMEx (Masked Input Modeling for Exploration). Previous approaches for exploration investigated intrinsic rewards, usually computed as a measure of a state or transition’s “novelty”, adding them to extrinsic rewards (the actual task’s rewards) to enhance the algorithm’s exploration and thus performance, especially for tasks with sparse rewards. The present paper argues that those previous intrinsic reward methods can be viewed under a unifying lens as approaches that use pseudo-likelihood estimation to estimate novelty. A general algorithm is introduced called MIMEx that computes pseudo-likelihood over entire trajectories, as the prediction loss of a masked sequence autoencoder. This loss corresponds to the intrinsic reward for an entire trajectory, and can be simply added to task rewards when running any RL algorithm. The paper then demonstrates the effectiveness of the approach with evaluations, baseline comparisons and ablation studies on a “PixMC-Sparse” benchmark suite.

Strengths

The paper is very clearly and naturally presented, with great care put into arguing for a coherent story. The work is original, arguing for a straightforward but interesting generalization of existing concepts in the RL literature. The natural blending of theoretical justification for the work (pseudo-likelihood as novelty), application of recently popular techniques (masked sequence modeling) and experimental results make the paper highly significant. The generalization of the concept of intrinsic rewards to entire trajectories is also extremely interesting. The experimental results are thorough, including evaluation on different benchmark settings, comparisons with other approaches, and extensive ablation studies. In particular, the ablation study showing how scaling of the autoencoder transformer affects results seems very interesting, and should warrant further investigation.

Weaknesses

It seems that adding a whole masked autoencoder to compute intrinsic rewards may incur a hefty computational cost. What is the overall runtime / resource utilization of MIMEx compared to other baselines? Even if MIMEx is more computationally demanding, it can still be a good choice due to sample efficiency, especially for offline RL tasks. In any case, the paper would greatly benefit from such an analysis.

Questions

A few points: * **[a]** Can the authors address the point raised in “Weaknesses”, related to computational cost? * **[b]** In Section 5.3, “Trajectory-level exploration”, specifically in Figure 4: Why does a sequence length of 6 reduce performance? From the authors’ presentation, it may seem that increasing trajectory length should in principle always guarantee a benefit. Why is there instead an optimum at lengths shorter than 6? * **[c]** Minor point: For all plot figures, font sizes are too small.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

I believe the authors appropriately addressed the limitations of their work.

Summary

This paper proposes to use masked autoencoding (similar to MAE) in RL and use the loss as an intrinsic reward for exploration in sparse reward domains. Their method, MIMEx, does masked reconstruction of latent observation o_t, based on the previous T observations, and assigns the loss as intrinsic reward r_t. They propose that count-based or prediction error-based intrinsic rewards can be viewed in the same masked autoencoding framework, as specific inputs and masks (such as ICM and RND). They show experiments in DM Control Suite and a harder version of PixMC denoted as PixMC-Sparse, which sparsifies the shaped reward. They show that MIMEx outperforms ICM and RND baselines.

Strengths

Experiment results suggest it is stronger than ICM and RND, and the method itself is relatively straightforward. The paper does clearly demonstrate evidence that this is a viable method for exploration.

Weaknesses

The formulation of masked autoencoding and trying to express other methods like RND and ICM in terms of it is not a very strong argument. One of the key parts of RND is the usage of a fixed random network, which is not properly captured by a varying mask; it is captured by allowing an arbitrary transformation of the inputs. For ICM, the representation learning part from inverse control is also not captured by the formulation, but also delegated to a transformation of the inputs. Representation learning is often a key part of an exploration algorithm, which MIMEx does not directly capture. By allowing an arbitrary transformation of the inputs, the actual masking portion of MIMEx becomes less important, since masking can equally be delegated as a different transformation of the inputs, resulting in only needing a trivial mask for MIMEx. I think there should be less emphasis on this part for the paper, as the experiments also do not investigate this idea further. An example of further investigation would be to try to emulate RND/ICM or other exploration approaches through specific masks (not just random masks) and showing whether specific masks may be better or worse than uniform random masking. I believe your reference [14] Byol-explore, does incorporate sequence-level information in its prediction error for intrinsic reward, which is not mentioned in related work. ---- After Author Rebuttals ---- After reading other reviews and the author rebuttals, I think the authors have addressed the addressable parts of my concerns, as well as have stated that they will clarify their main claim to be more nuanced. Thus I'm inclined to slightly raise my score.

Questions

MIMEx seems to be encoding and decoding latents (as opposed to pixels), so are you allowing gradients from MIMEx to also flow back through to those latents? I.e. are you stopping gradients from the input/target sides? If not stopping gradients, then is there a danger that the representation will collapse to a constant? If yes, then are you relying on the RL to train the latent representations? In your Figure 1, MIMEx only receives latent observations as input/target. Why are actions not included? Many forward model predictions rely on state and action information for prediction, including ICM. What was the reasoning behind this? If actions were included, then MIMEx would be able to more closely emulate many multi-step prediction methods such as [14] Byol-explore or SPR (Data-Efficient Reinforcement Learning with Self-Predictive Representations https://arxiv.org/abs/2007.05929) by picking a mask the only masks out observations but keeps around actions.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

This paper addresses some of the limitations, but another potential limitation for this method is partial observability, especially for long horizon problems. Only using the previous T observations as input might not be a rich enough context for predicting the next observation, so there would need to be another way of adding more past information. Another potential limitation is stochasticity. Since MIMEx is predicting observations, it will suffer from high loss if the observations are stochastic, similar to many other prediction-error-based approaches that try to directly predict the observation. RND avoids the issue by same-step prediction, while ICM relies on the representation learning of inverse control to filter our dynamics-irrelevant noise. With MIMEx, it seems like even adding white noise to observations could end up with high losses everywhere.

Summary

The paper introduces a novel approach to exploration in reinforcement learning (RL) called Masked Input Modeling for Exploration (MIMEx). MIMEx uses a masked autoencoding objective on variable-length input sequences to derive intrinsic rewards for exploration. The paper claims that MIMEx improves exploration efficiency in sparse-reward tasks and show that MIMEx outperforms claimed competitive baselines on tasks from the PixMC-Sparse suite and the DeepMind Control Suite except for in a few cases, such as certain cases in which the reward is not sparse.

Strengths

**Originality** The paper introduces a fresh approach to exploration in reinforcement learning, applying existing concepts in a new way. This application of masked autoencoders and pseudo-likelihood estimation to the problem of exploration in sparse reward environments is at least innovative. **Quality** The paper is well-structured and provides a comprehensive examination of the proposed MIMEx method. The paper creates a new benchmark dataset from an existing one with non-trivial modifications, and explains these. The paper has conducted extensive experiments and provided a detailed ablation study, which adds to the quality of the paper. The paper does an adequate job at discussing it's limitations and potential future directions for its work. **Clarity** The paper is very well-written and mostly clear in its presentation. The paper has done an excellent job of explaining the MIMEx method, its implementation details, and the experimental setup. The explanation of dense rewards for the PixMC environment, the construction of the PixMC-Sparse, and the use of figures and tables to illustrate the results is superb. **Significance** MIMEx shows superior performance to certain benchmark exploration algorithms, __random action noise__, __intrinsic curiosity module__, and __random network distillation__, some of which have been shown to perform well at the time of their publication on long-horizon sparse reward tasks, such as Montezuma's Revenge. The model allows for easy adjustment of the time horizons it considers, though tuning this hyperparameter may be difficult. MIMEx is agnostic to standard RL algorithm and model choice, allowing it to be used on several problems. --- Overall, this paper is easy to read and combines ideas into a novel exploration method that could be useful. However, it's unclear to me the generalizability of most of the results and the increased compute, memory, and wall-clock expenses wrt the horizon needed for the MIMEx module. Thank you for providing code with your submission for both sets of benchmarks! **I have changed my score from a 5 to a 6 as a result of the authors rebuttal, in which they show evidence of additional work that addresses the key concern I had as well as additional results clarify ambiguities I had. Thank you, Authors. I'm happy that my rebuttal helped you improve your work and that you considered and addressed my concerns by performing additional work to improve your work and sharing these results clearly in your rebuttal.**

Weaknesses

The paper's main idea, while innovative in its application, is built upon existing concepts in the field, such as masked autocoders and prediction error based exploration. It uses standard latent embeddings of observations and then reconstruction of masked embeddings to compute the prediction error, which is work that has been done for years, a recent example of which is [1], work that the paper transparently references (this transparency is a strength in clarity and overall). The paper could have better highlighted the unique aspects of their approach and how it diverges from or improves upon existing methods in reinforcement learning exploration. Though PixMC-Sparse gets a discrete bonus in some environments for reaching a certain intermediate state, the paper could've evaluated MIMEx on easily accessible or created environments by easily modifying existing ones that focus on more diverse types of sparse rewards, such as discrete instead of continuous (wrt to goal distance) rewards. - The paper's quality could be improved with a more diverse set of experimental tasks and environments. The current selection of environments for which hyperparameter and ablation studies were done do not generalize to much outside of their specific domains, so they don't demonstrate the generalizability of the MIMEx method. However, I do understand that these experiments can be computationally expensive, but I would recommend performing additional ablations on environments that differ significantly from Kuka_____ in order to maximize confidence in the generalization of the results of these ablations. - In lines 118-119, the paper states that "MIMEx can be added...as a lightweight module..." but there is no mention, let alone evaluation wrt horizon $T$ of the increased memory footprint or wall-clock time needed to calculate the intrinsic reward. - Later exploration methods [2, 3] that are catered toward sparsity or perform better than the baselines presented are not used as baselines. But, to the paper's defense, [3] was published just Sep 2022. However, the issue is that the most recent baseline used is RND, which was published in 2018, and there are many better ones that exist to evaluate MIMEx against. - The paper could better justify the usage of each of the exploration baselines used for readers that may not be familiar with RND's performance on Montezuma's revenge though it does do this somewhat in the Related Works section. - Again, while the paper shows promising results on specific tasks, its impact could be limited if the method's effectiveness doesn't extend to a broader range of tasks and environments. The paper could increase its relevance and significance by demonstrating MIMEx's effectiveness in more diverse scenarios and discussing potential applications beyond the current scope. [1] Xiao, T., Radosavovic, I., Darrell, T., & Malik, J. (2022). Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173. [2] Zhang, T., Rashidinejad, P., Jiao, J., Tian, Y., Gonzalez, J. E., & Russell, S. (2021). Made: Exploration via maximizing deviation from explored regions. Advances in Neural Information Processing Systems, 34, 9663-9680. [3] Eberhard, O., Hollenstein, J., Pinneri, C., & Martius, G. (2022, September). Pink noise is all you need: Colored noise exploration in deep reinforcement learning. In The Eleventh International Conference on Learning Representations.

Questions

1. Should the equality in the display mode equation in line 84 be an approximation since $|X|$ is not necessarily infinite? 2. Referring to lines 204-206, what are the internals of the transformer blocks? Could you clarify because not all transformer blocks are in the same order or contain the same properties as [4]. 3. Why choose the three chosen exploration baselines, random action noise, intrinsic curiosity module, and random network distillation? 4. I'm confused by the inclusion of the sample curriculum in Table 3. Is this only an informative point, or is this something you evaluated? 5. In lines 233-234, the paper states "Trajectory-level exploration To our knowledge, MIMEx is the first framework that successfully incorporates sequence-level intrinsic reward to solve hard exploration tasks." Doesn't [5] do this implicitly as the discount factor in [5] is varied? Again, this is a relatively recent publication. 6. Minor suggestion, to make the results of Figure 2 more tangible and to improve readibility, I suggest providing a visual and explanation of the KukaPick task alongside Figure 2, even though it is in the Appendix. 7. Do you have ideas on extensions of Online masked prediction to offline RL and any advantages or disadvantages using MIMEx with this paradigm? [4] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30. [5] Ramesh, A., Kirsch, L., van Steenkiste, S., & Schmidhuber, J. (2022). Exploring through random curiosity with general value functions. Advances in Neural Information Processing Systems, 35, 18733-18748.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The main limitation is the limited understanding of the generalizability of the results to problem domains, even sparse reward ones, that are different than PixMC-Sparse, which most are. The paper does address that MIMEx may hurt performance when rewards aren't sparse, but it could provide more limitations on absence of additional memory, compute expense, and wall-clock time needed to add the MIMEx intrinsic reward generation model, which has multiple, multiple-layer transformers and more.

Summary

This work proposed a general framework for deriving intrinsic rewards called Masked Input Modeling for Exploration (MIMEx). This method starts from the observation that existing intrinsic reward approaches are special cases of conditional prediction, where the estimation of novelty can be seen as pseudo-likelihood estimation with different mask distributions. From this perspective, MIMEx derives intrinsic reward based on masked prediction on input sequences, which naturally lends itself to controlling the difficulty of the underline conditional prediction task. Empirically results on eight tasks from PixMC-Sparse and six tasks from DeepMind Control Suite demonstrate that MIMEx outperforms other baselines regarding sample efficiency.

Strengths

1. This work proposed a general framework for deriving intrinsic rewards, which can be applied to various hard-exploration tasks. 2. The most interesting part of MIMEx is that it enables extremely flexible control over the difficulty of conditional prediction for deriving intrinsic rewards. Section 5.4 provides comprehensive studies to understand how varying mask distribution affects the performance of MIMEx. 3. Extensive ablation studies have been provided to understand why MIMEx works better than other approaches. The comparisons are provided among diverse directions, including trajectory-level exploration, variance reduction, and model scalability.

Weaknesses

1. The smoothness of writing can be improved, e.g., section 3.1 is not closely related to other parts of MIMEx. The motivation of using masked prediction loss as the intrinsic reward is not clear enough. 2. Line 81 cited a retracted paper.

Questions

1. What motivates using masked prediction loss as the intrinsic reward? 2. When reading the introduction, I have the following questions: how masked language modeling is connected to pseudo-likelihood? Why approaches that estimate novelty can be viewed as modeling different conditional prediction problems or masked prediction problems with different mask distributions?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

1. One limitation of MIMEx is that it can potentially hurt performance on easier tasks that do not require exploration, as in the case of exploration bonus approaches. 2. Stronger bias from other sources, like the expert demonstrations, is needed to improve the exploration ability of MIMEx further. ------------ After rebuttal --------------- I would like to thank the authors for addressing part of my concerns. I agree to increase my score slightly.

Reply to Authors

Thank you for your rebuttal. The wall clock time and memory usages are indeed impressive. I don't have any additional comments, and maintain my score as previously stated.

Masking Distributions and Actions

Thanks for the response to my questions! I appreciate the response and additional ablation on the masking distribution, which does give more insight into its importance. However I do want to emphasize my point that without actions, MIMEx is unable to properly capture many intrinsic reward methods. The main reason I am re-emphasizing this point is because this paper is making a very strong claim for the expressiveness of the MIMEx framework, which is only partially true. For example in Section 3, line 99: "Inspired by this perspective, we propose Masked Input Modeling for Exploration (MIMEx), a unifying framework for intrinsic reward methods...". The term "unifying" is used multiple times throughout the paper and abstract. MIMEx generalizes 1-step to multi-step with very flexible mask distributions, which is its strength. However without actions, MIMEx fails to model any 1-step dynamics model type of intrinsic reward, including ICM. In your ablation of [l2_fixed_50 - seq length 2, fixed mask, 50%], if you are trying to do 1-step next state prediction, having the action vs. not having the action are fundamentally different problems, and MIMEx can only model one version. Thus I think this claim of MIMEx being a "unifying" framework is too strong, and it would be better to make a more nuanced claim such as "generalized framework for state-based multi-step intrinsic reward".