MIMEx: Intrinsic Rewards from Masked Input Modeling
Authors: Toru Lin
Allan Jabri
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
exploration
intrinsic reward
intrinsic motivation
masked autoencoder

Abstract: Exploring in environments with high-dimensional observations is hard. One promising approach for exploration is to use intrinsic rewards, which often boils down to estimating "novelty" of states, transitions, or trajectories with deep networks. Prior works have shown that conditional prediction objectives such as masked autoencoding can be seen as stochastic estimation of pseudo-likelihood. We show how this perspective naturally leads to a unified view on existing intrinsic reward approaches: they are special cases of conditional prediction, where the estimation of novelty can be seen as pseudo-likelihood estimation with different mask distributions. From this view, we propose a general framework for deriving intrinsic rewards -- Masked Input Modeling for Exploration (MIMEx) -- where the mask distribution can be flexibly tuned to control the difficulty of the underlying conditional prediction task. We demonstrate that MIMEx can achieve superior results when compared against competitive baselines on a suite of challenging sparse-reward visuomotor tasks.

MIMEX: Intrinsic Rewards from Masked Input Modeling
Toru Lin University of California, Berkeley toru@berkeley.co edu edu
Allan Jabri University of California, Berkeley babrieberkeley. edu
Abstract
Exploring environments with nigh-dimensional observations hard. One promis- ing approach for exploration is to use intrinsic rewards, which often boils down to estimating "novelty" of states, transitions, or trajectories with deep networks, Prior works have shown that conditional prediction objectives such as masked autoencod- ing can be seen as stochastic estimation of geeudollikelhood perspective naturally leads to a unified view on existing intrinsic reward approaches: they are special cases of conditional prediction, where estimation novelty can be seen as this view, we propose a general framework for deriving intrinsic rewards- Masked Input Modeling for Exploration (MIMEx) where the mask distribution be flexibly tuned control the difficulty the underlying conditional prediction task. We demonstrate that MIMEX can achieve superior results when compared against competitive baselines on a suite of challenging sparse-reward visuomotor tasks.
We show how this
estimation with different mask distributions. From
Introduction
In supervised learning (SL), it is often assumed that data is given, i.e. the data collection process has been handled by humans. In reinforcement learning (RL), however, agents need to collect their own data and update their policy based on the collected data. The problem of how data is collected is known the exploration problem. Agents face the eeploraationexxpoiinnee balance between maximizing its cumulative rewards (exploitation) and improving its knowledge of the environment by collecting new data (exploration). Exploration is therefore of central importance in the study of RL; it is the only way through which an agent can learn decision-making from scratch.
trade-off to
Simple exploration heuristics like using e-greedy policy [22 (33 and adding Gaussian action noise 21 30, 31 are often too inefficient to solve challenging tasks with sparse reward. Meanwhile, designing dense reward functions is intractable scale to many tasks across environments, due the amount of domain knowledge and engineering effort needed. Training agents with intrinsic rewards, ie. providing exploration bonuses that agents jointly optimize with extrinsic rewards, is promising way to inject bias into exploration. This exploration bonus often derived from a notion of "novelty", such that agent receives higher bonus in more "novel" states. In classical count-based methods (35). for example, the exploration bonus is derived from state visitation counts directly. To scale to high dimensional states, recent works often derive the exploration bonus from either (1) pseudo-counts [|244] 39] (i.c. approximate state visitation counts) or (2) prediction error from a proxy model. In the latter case, the proxy model responsible for representing the relative novelty of states, for example with respect to learned forward models 25
34 or randomly initialized networks (8)
While count-based (including pseudo-counts) methods and prediction-errot based methods are typ- ically viewed as distinct approaches for intrinsic reward [ these methods often boil down similar formulations in practice. That is, they rely on the same underlying idea of estimating
Code available httttp://ggcuuu..c////rroo//////
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

novelty by approximating likelihood, while different in the quantities they model: states, transitions, trajectories, etc. Recent works have demonstrated how conditional prediction objectives such as masked language modeling are connected to pseudo- likelihood, Under this view, approaches that estimate novelty can be viewed as modeling different conditional prediction problems, or masked prediction problems with different mask distributions.
[29 43 one way of approximating
From this perspective, we propose Masked Input Modeling for Exploration (MIMEX), a generalized framework for intrinsic reward methods. On high level, MIMEX allows for flexible trajectory-Ievel exploration via generalized conditional prediction. To evaluate on eard-exploration problems with high-dimensiona observations and dynamics, we develop a benchmark suite of eight challenging robotic manipulation tasks that involve realistic visuomotor control with sparse rewards. We show that MIMEx achieves superior results when compared with common baselines, and investigate why it works better than other approaches through extensive ablation studies.
Preliminaries
2.1 From RL Exploration to Intrinsic Reward
Among approaches for RL exploration, one the most successful to jointly optimize an intrinsic reward in addition to the environment reward (a.k.a. the extrinsic reward). At each time step, the agent is trained with reward r, = rf + B ri. where rf is the environment reward, rf is the intrinsic reward, and / is parameter control the scale of exploration bonus. The intrinsic reward approach is flexible and general, as it simply requires modifying the reward function. This also means that intrinsic reward function can be learned and computed in manner compartmentalized from the policy, allowing one to leverage models and objectives that scale differently from policy optimization.
The key principle underlying intrinsic reward approaches is optimism in the face of uncertainty (7]. which has been empirically observed to be effective for encouraging active exploration [37]. In other words, the intrinsic reward is supposed to encourage agents visit states that are less frequently visited, or pursue actions that lead to maximum reduction in uncertainty about the dynamics; note that the latter is equivalent visiting transitions that are less frequently visited.
2.2 Intrinsic Reward and Conditional Prediction
Works in intrinsic reward literature tend to draw on different concepts such as novelty, surprise, or curiosity, but their practical formulation often follows the same principle: to encourage the agent to visit novel states or state transitions, is desirable for r to be higher in less frequently visited states state transitions than in frequently visited ones. Deriving intrinsic reward therefore requires knowledge on state density or state transition density.
prediction- methods S 24), state density is modeled directly; in iitiiction-errr-bbaee methods [18 25 forward model is learned and prediction error is used as proxy of state transition density. Under this view, we can unify existing intrinsic reward methods into one framework, where only the underlying conditional prediction problem differs.
Recent intrinsic reward methods can be categorized into either pseudo-counts-based error-based [2]. In pseudo-counts-base
2.3 Pseudo-Likelihood from Masked Prediction Error
have shown how the masked autoencoding objective in models like BERT 9 be viewed stochastic maximum beeudo-likelhhodd estimation on training set of sequences of random variables X = (#1,..., ). Under this objective, the model approximates the underlying joint distribution among variables, by modeling conditional distributions Xx, resulting from masking at position t. Despite being an approximation, the aseudo-likelihowe has been shown to be a useful proxy, e.g. as a scoring function for sampling from language model 29
Prior works 29 43
More concretely, define pseudo log-likelihood as:
X
PLL(O;
1
MΣΣ log logp(zz|Xx)). XED 1=1
D
where D is a training set. To more efficiently optimize for the conditional probability of each variable in sequence given all other variables, they propose to stochastically estimate under masking
'While 43 was later retracted due to an error, our argument does not depend on the part with error. We provide additional clarification on the validity of our claim in Appendix A.]

---

distribution &y -u([1....... (xI):
[
K
1
1 K
logp(z,(Xx.) =Beu [logp(x.(Xx,))
2
log plpp [X (is)
X
While this is defined for categorical variables, for continuous variables we can simply consider reconstruction eтror on masked prediction as regressing to unit Normal distribution of variables such as observations, features, and trajectories. From this perspective, we can view different cases of conditional prediction as different forms of maximum pseudo-likelihood estimation and implement them as masked autoencoders.
Method
Masked
Prediction Objective
Pseudo-counts [4
subset of state ST current-step feature FT
|&T [mask( (T) ) with mask state dims
RND
2r f (mask(ærr))
with ø7 ((ST) under a random network e [[ ar] [[[mas(*TT+ 27, ar) (r])
ICM
next-step feature 27+1
225
with xr = ((87) under inverse dynamics encoder Ø
Table I: Examples of existing intrinsic reward methods viewed as masked prediction problems with different mask distributions.
3 Sequence-Level Masked Autoencoders
The observations outlined in Section in particular the link between masked prediction and pseudo- likelihood, allow us to draw connections between different classes intrinsic reward methods. Prediction-error-basedl methods consider the problem of conditionally predicting 8 from (st, ax) by masking 8141 from ($t+1; 8t; 04) tuples, which can be viewed as approximating the pseudo- likelihood of the next state under a Gaussian prior. Pseudo-counts- based methods approximate the likelihood of st, which can achieved by marginalizing the prediction error across maskings of st- As summarized in Table prediction with different mask distributions.
view these approaches estimating pseudo- likel--i via masked
Inspired by this perspective, we propose Masked Input Modeling for Exploration (MIMEx), generalized framework for intrinsic reward methods. Simply put, MIMEX derives intrinsic reward based masked prediction on input sequences with arbitrary length. Such framework naturally lends itself to greater control over the difficulty of the underlying conditional prediction problem; by controlling how much information we mask, we can flexibly control the signal-to-noise ratio of the prediction problem and thus control the variance of the prediction error. Morcover, it provides the setup to fill space in literature: trajectory-Iever exploration. Existing approaches framed as conditional prediction often consider one-step future prediction problems, which can saturate early useful exploratory signal & 16 25). While longer time-horizon prediction problems capture more complex behavior, they can suffer from high variance (77) MIMEx allows for tempering the difficulty of prediction by varying input sequence length; by setting up conditional prediction problems on trajectories, we can obtain intrinsic rewards that consider transition dynamics across longer time horizons and extract richer exploration signals. We can also easily tune the difficulty of the prediction problem, by varying both the input length and the amount of conditioning context given a fixed input length.
The fact that MIMEX sets up a masked prediction learning problem comes with several additional benefits. First, masked autocncoding relies on less domain knowledge compared to methods like contrastive learning, and has proven success across many different input modalities ( 10 17
32 43
leverage standard architectures such those used in masked language modeling [ and masked image modeling [7]. for which the scalability and stability have been tested. MIMEx can be added to any standard RL algorithm as a programmatically simple module that encourages exploration.
Moreover,
We describe the implementation details of MIMEx in the remainder of this section. An overview of our framework is shown in Figure !
3.1 Sequence Autoencoders for Exploration
Given a sequence of trajectory information stored in the form of tuples (ox, ak, r for h = 1,.., we now describe how to derive intrinsic reward rf for time step t.
3

---

1-2
2-2
MIMEX (T=3)
1
Z
2
Encodor
MIMEX
7
Enoade
1++1
Decoda
2++1
c++
H+1
Mask
4++2
2442
r+4
2++2
4+2
[+2
Figure 1: Overview. The overall schematic of MIMEx. in the case of exploration sequence length == (Left) Given a trajectory high high-dimensional observations {ok} J. the base RL encoder first encodes them into latent embeddings } MIMEx then takes in { and derives the corresponding intrinsic rewards {ri } for &
2, 1, I.... 2. (Right) For each timestep MIMEX performs online masked prediction on the sequence of embeddings up to T steps before k. padding missing timesteps with zeros. The masked prediction loss L2 is used the raw intrinsic reward, i.e. Lg rk
Input sequence generation We define parameter 7 that represents the length of sequence we use derive the exploration bonus. From the selected time step we extract the sequence of inputs up to T steps before f, including time step f. When >= T,, we take sequence (O--T+1.- 0r-1, 0z); when T, we pad observations at time steps before the initial step with zeros such that the input sequence length is always T. We use a Transformer-basee Og latent embedding 2k. obtaining an input sequence (2. 2t).
encoder to convert each input observation
Online masked prediction Given the generated input sequences, train masked autoencoder model online using the sequences as prediction target. Specifically, we treat observation embedding 2g each time step & as discrete token, randomly mask out X% of the tokens each sequence using a learned mask token, and pass the masked sequence into a maansforme-bbsse decoder to predict the reconstructed sequence. The random mask generated with a uniform distribution. The more "novel" an input sequence is, the higher its masked prediction error is expected to be. MIMEx compatible with wide range of high- dimenssiinii inputs across different modalities, thanks the scalability of its architecture and the stability of its domain domain-agnostic objective.
3.2 Masked Prediction Loss as Intrinsic Reward
We use the scalar prediction loss generated during the online masked prediction as an intrinsic reward, i.c. r Lt. The intrinsic reward is used in the same way described in Section 2 i.c. r. = rf 3 r where 3 is a parameter to control the scale of exploration bonus. Intuitively, this encourages the agent explore trajectory sequences that have been less frequently encountered.
3.3 Backbone RL Algorithm
Our exploration framework is agnostic to the backbone RL algorithm since it only requires taking trajectory sequence input. In Section 5 empirically demonstrate how MIMEx can work with either on-policy or off-policy RL algorithms, using PPO 31 and DDPG 21 33 as examples.
4
PixMC-Sparse Benchmark
4.1
Motivation
Exploration can be challenging due sparsity of rewards, or high-diiennioonality of observation and action spaces paired with complex dynamics. Existing benchmarks for exploration algorithms
span a spectrum between two ways controlling the sparsity of reward: (1) long task horizon; (2)
-igh-dimensional observations (e.g. pixels). Either side of the spectrum can contribute to increased sample complexity, along the time dimension the space dimension. While existing works on deep
RL exploration often evaluate on popular benchmarks such as the Arcade Learning Environment [6]
and DeepMind Control Suite 40), these benchmarks often involve densely engineered rewards or deep exploration in environments with simple dynamics (e.g. Montezuma's Revenge)
3 12 38

---

Identifying missing space of exploration benchmarks that involve gigher-dimensional observation inputs and more complex dynamics, we construct a challenging task suite named PixMC-Sparse (Pixel Motor Control with Sparse Reward). The suite includes eight continuous control tasks, with realistic dynamics and egocentric camera views of robot as observation. PixMC-Sparse built on PixMC [4] as an extension to the original suite.
4.2
From PixMC to PixMC-Sparse
The original PixMC task suite does not necessarily pose challenging exploration problems due to densely shaped rewards (see Appendix A.2 for details of the original PixMC reward functions). Therefore, we modified the tasks to reflect more realistic hard-exploration challenges. To summarize, we sparsify the reward signals by removing certain continuous distance-based rewards from each original task. We show details of how each task's ' reward function modified in Table 2 and include additional annotated visualization in Appendix A.3
4.3 Evaluation Protocol
The difficulty level of each hard-exploration task presented in PixMC-Sparse can be tuned by (I) increasing the environment's 's time space resolution, (2) minimizing reward shaping. Either approach reduces the extent to which the benchmark is saturated, and the latter is itself a practical and realistic motivation for studying exploration in RL. In Table show sample curriculum to minimize reward shaping for Cabinet, Pick, and Move tasks, starting from the fully dense reward functions (included in Appendix A.2
Task
Removed Dense Reward Terms
Task
Example Steps Spunify Rewards
FrankaReach
distance
Remove "disince to goal" term
KukaReack
distance to goal
KatuClttine!
Kamove "Bager Thunb disiace lundle tenms
FrankaCahing Cabinet
distance to handle
Remove "distance to gool" term
KakaCabinet
finger
liadle, thumb distance handle
FronkaPhik
1. Kanaae "isiaane to stjea" form
FrankaPick
distance object
KakaPPick
Ramone "linger Tilunl disiance object tents
2. Remove "distance to goal" term
KudaPPick
linger distance to object, thumb distance to object
2. Remove "distance to pool tem
FrankaMove
distance to object
FrankaMow
"Ilstanc ohjez""
KukaMove
inger distance object, thumb distance object
Remove
to gool"
KolaMore
Nomore "linger tiuuab disiance abyect" tenms
Table 2: We define PixMC- Sparse tasks by removing
RRmmvve "distance to pool" term
reward terms from the reward functions of PixMC Table 3: A sample curriculum to gradually tasks. In the table, show details of how each task reward function is modified.
increasing exploration difficulty for each task in PixMCCSpprrs
In Figure 2 we present case study of evaluating exploration algorithms on increasingly sparse KukaPick tasks. Task success is defined by agent picking up an object on table using Kuka arm
Figure 2: Evaluating exploration algorithms with increasingly sparse reward. We show an example on PixMC-Sparse's KukaPick task. When using fully dense reward, the base RL agent with random action noise exploration can already achieve success rates of 80% consistently (dense). With one dense reward term removed, the base RL agent completely fails to solve the task (sparse). With the addition of our exploration algorithm, MIMEx, some success observed and the success rates again start to provide useful in- formation for evaluation (sparse+MIMEx)) With more dense reward terms removed, performance of MIMEx agent further away from saturation (spurser+ We can continue to develop and evaluate more powerful exploration algorithms by using increasingly sparse reward.
KukaPick
80
Succeaa ftato 1N0
60
Steps (M)
1+7
dence
sparse
Sparse-MIHEE
sparser MINEX
5 Experiments
In this section, we show that MIMEx achieves superior results and study why works better than other approaches. We start by evaluating MIMEx against three baseline methods, on eight
More details the tasks can be found in Appendix A.3
5

---

tasks from FiXM--Spars (Section 4 and six tasks from DeepMind Control suite [40] Then, we present ablation studies of MIMEx on PixMC-Sparse to understand specific factors that contribute to MIMEx's s performance and obtain insights for its general usage. We report all results with 95% confidence intervals and over random seeds.
5.1 Implementation Details
We describe details of MIMEx implementation, independent of the base RL algorithms. On a high level, implementation of the MIMEX model is similar to that of MAE [7]. The model consists of an encode & mask module and a decode module.
Given an input sequence x with shape (B. I,, D) where B is the batch size, I. is the sequence length, and D is the feature dimension, the encode & mask module first projects the last dimension of x to the encoder embedding size through a fully-connected layer; the corresponding positional embeddings are then added the time dimension (length L) of x. We denote this intermediate latent vector as
2. We apply random masking with uniform distribution to the latent vector z, following the same protocol used in MAE [7. The masked latent vector Zm is then passed through a Transformer-based encoder.
The decode module first inserts mask tokens to all masked positions of 2m, then passes 2m to Transformer- based decoder. Finally, another fully-connected layer projects the last dimension of Am back the input feature dimension D. We calculate the reconstruction loss on only the masked positions, and use the loss as intrinsic reward.
For all environments, we use batch size of 512, the Adam [19 optimizer, and a MIMEx learning rate of 0.0001. We use a mask ratio of 70% for all tasks; 3 (exploration weight) of 0.05 for Reach tasks and 0.05 for Cabinet, Pick, Move tasks. For the encoder, we use embedding dimension of 128, with 4 Transformer blocks and heads; for the decoder, we use an embedding dimension of 64, with Transformer block and 2 heads. We run each experiment
NVIDIA A100 GPU.
5.2 Comparison with Baselines
We compare MIMEx to three RL exploration baselines: random action noise (noise), intrinsic curioisty module (25 (iem) and random network distillation [8] (rnd). We choose these baselines for their generality in formulation and robustness in performance. We show that MIMEx can be flexibly incorporated into both on-policy algorithms and off-policy algorithms: for tasks in PixMC-Sparse. we implement MVP 44 with PPO (31 (an on-policy algorithm) being the core RL algorithm; for
tasks in DeepMind Control suite (DMC), we implement DrQv2 [45) with DDPG [3] (an off-policy algorithm) being the core RL algorithm. Note that MVP and DrQv2 are the respective state-of-the-an algorithm on each environment. Results are shown in Figure S
For Pix.C-Sppase we find that MIMEx outperforms all baselines on all tasks, in terms of both final task success rate and sample efficiency. Importantly, the performance gap between MIMEx and baseline methods becomes more pronounced on the harder Pick and Move tasks, where exploration also becomes more challenging.
For DMC, we find that MIMEX achieves higher average performance than all baselines on 2 out of 6 tasks (cartpole wwteeewwwwsssseeee -parse. quadruped _wakk), while having similar average performance as the noise and mn baselines on the other 4 tasks pacrobotswiingwe finger turn hard). We also observe that MIMEx exhibits lower variance compared similarly performant baselines on all tasks. As discussed in Section 4 we find it difficult to benchmark RL exploration methods on DMC even when using high-dimensional pixel inputs, and recommend using harder exploration task suite such as PixMC-Sparse to better evaluate state-of-the-a exploration algorithms.
funger-tum_easyy quadruped rumm
5.3 How does MIMEX improve performance?
We have conducted extensive ablation studies make relevant design choices and understand why MIMEx works better than other approaches. For brevity, here we only present factors that result in significant improvement in task performance. Due to computational constraints, we only show results on 2 tasks from PixMC-Sp an easier KukaReach task and harder KukaPick task. We summarize our findings below.
Trajectory-level exploration To our knowledge, MIMEx the first framework that successfully incorporates sequence-level intrinsic reward to solve hard exploration tasks. Intuitively, doing so

---

60
80
40
-0
Steps
Steps (M)
KukaCabinet
FrankaCabinet
100
80
80
20
0.0
0.
1.0
1.5
2.0
2.5
300
0.0
0.
1.0
1.5
2.0
Steps
167
Steps
1e7
100
60
40
40
KukaMove
FrankaMove
40
-00
20
Steps
3e7
Steps (M)
167
md
quodruped
800
800
200
200
02
2.4
05
1.4
02
Stran
Steon (11
Stepl
cartocle_sttpp
quadruped
1230
Edicde Rivert
12
0.4
1.4
0.
1.2
Stepn
DUP' noise
iom
md
Figure 3: Comparison with baselines. We compare method (ours) against 3 exploration baselines: random action noise (noise), intrinsic curiosity module (icm), and random network distillation (rnd). The top figure shows results 8 tasks from the PixMC-Sparse suite, and the bottom figure shows results on tasks from the DeepMind Control Suite.

---

encourages agent to search for new trajectories rather than one transitions, which may lead to more complex exploration behavior. The empirical results from varying the exploration sequence length used in MIMEx, as shown in Figure suggest that using sequences with length greater than 2 (i.e. transitions that are longer than one step) indeed improves exploration efficiency.
Variance reduction High variance is classical curse to RL. algorithms [36] In deriving intrinsic reward using MIMEx, a source of high variance the random masking step during the mask- autoencode process. We hypothesize that reducing variance in conditional prediction can lead to large performance gain, and empirically test the effect of doing so. In particular, we find that a simple variance reduction technique applying random mask multiple times and taking the average prediction error over different maskings can greatly improve agent performance. We show our results in Figure 4
Model scalability Recent advances in generative models have showcased the importance of scale, which we hypothesize also applies to the case of trajectory modeling for deriving intrinsic reward. We show our empirical results that verified this hypothesis in Figure[4] Performance gains are achieved by simply increasing the size of MIMEX's base model (specially, doubling the size of Transformer-basee decoder), especially in the harder KukaPick environment.
5.4 How does varying mask distribution affect performance of MIMEx?
As discussed in Section 3 one characteristic of MIMEx is that enables extremely flexible control over the difficulty of conditional predic- tion for deriving intrinsic rewards. In theory, this flexibility allows MIMEx to encompass the full spectrum of pseudo-likelihood estimation via masked prediction. We hypothesize that this would enable MIMEX to be both general and powerful: the same implementation of MIMEx can achieve superior performance compared to baselines in wide range of settings. We are therefore particularly interested in studying how varying mask distribution affects the per- formance of MIMEX. Below, we show ablation studies on two attributes of MIMEx that can be easily adjusted to change the mask distribution: mask ratio and mask dimension.
445 475
17%
14
0600
Mask ratio One straightforward way to vary the mask distribution and tune the difficulty of conditional prediction is to simply adjust the mask ratio when applying random mask. Intuitively, when the mask ratio is lower, the amount of information given to the model is larger, and the conditional prediction problem becomes easier. On the two PixMC-Sparse tasks we used, we find that results from using certain fixed mask ratios exhibit high variance across tasks, for example the results from us- ing mask ratios 20% and 50% in Fi We hypothesize that, because the difficulty of ex- ploration varies from task to task, there exists different "optimal" mask ratio parameter for each task. In theory, MIMEx can be tuned to achieve "optimal" performance by choosing a specific mask ratio for each task. In practice, we use fixed mask ratio of 70% for our de- fault implementation since we find it achieve a good trade-off between hyperparameter tuning and robustness of performance.
rms
Knu Nouth
0h0c
Figure 4: (top) Varying exploration sequence length. In refers to exploration with sequence length of n for n = 2,3, 4,5,6. (middle) Varying number of times of masking for intrinsic reward calculation. nm/ is masking once, and nms5 is masking 5 times followed by tak- ing the average of all prediction errors.
(bottom) Scaling MIMEx decoder. Compared with the smaller model, the larger model uses a decoder that doubles in depth, number of output heads, and embedding dimensions.
8

---

Mask type Another degree of freedom of MIMEx lies in the flexibility of choosing which input dimension to apply random masks to. Fully random masking on the last dimension of input (i.e. feature dimension) will enable the most fine-grained control over the underlying conditional prediction problem. We hypothesize that this greater degree of freedom might allow MIMEx to achieve better performance on certain tasks, but might also require more substantial effort to find the optimal parameter. In Figure 5) we compare the performance of (1) masking only on the time dimension with mask ratio of 70% and (2) masking the feature dimension with mask ratios of 80%, 90%, and 95% on our chosen tasks. Empirically, we observe similar results between these two types of masking with the chosen mask ratios, opted to use the more computationally efficient time-only masking in our default implementation.
350
6.2
8.6 1.0
-50
15
0.4
31n 300
1.
-20
70% 20%
50%
70%
50%)c)
90%)63
30%
Figure 5: (left) Varying mask ratio. Results from masking 20%, 50%, and 70% during the random masking step of MIMEx. (right) Varying mask type and ratio. Results from (1) masking only on the time dimension with mask ratio of 70% (70%) and (2) masking on the feature dimension with mask ratios of 80%, 90%, and 95% (80%(a), 90%(()) 95%(a)).
Related Works
Undirected exploration [41 approaches add noise to individual actions (c.g. random action noise) [T] action policy (c.g. e-greedy), or even network weights [3]. These methods are often surprisingly
strong baselines for hard-exploration problems (38).
Classical count-based methods provide another effective way to perform exploration [35]. To bring
bount-based exploration to high-dimensiona state spaces, 4 and (24] propose pseudo-count as
useful surrogate of novelty, using CTS (5) and PixelCNN [42 as the density model respectively. (39)
is another example of count-based exploration, but uses hash code instead of pseudo-count for state
visitations.
( takes a different approach by approximating novelty with prediction error on random features. The intuition is the same - prediction error should be higher for less visited states. The closest to our work are perhaps works that generate prediction errors by learning a forward dynamics model. For example,
25 34 can be seen as special cases of our framework in that they only perform one-step prediction. [16] is similar to our work in spirit, but learns a multi-step dynamics model with recurrent networks and generates multi-step prediction error with B YOL-like [5]] prediction objective. Relatedly, (28) derives intrinsic rewards through predicting temporally extended general value functions, though focusing on partial observable tabular environment. [88 26 use reduction of uncertainty under the learned model as intrinsic reward; the former uses Bayesian neural network and the latter uses ensemble disagreement. In practice, these methods perform worse because is usually hard to get good estimate of the reduction in uncertainty with all the stochasticity inherent in modern neural networks and optimization tools.
Limitations and Future Work
Inspired by unified perspective on intrinsic reward approaches, we present Masked Input Modeling for Exploration (MIMEx), a general framework for deriving intrinsic rewards. By using a masked autoencoding objective variable-length input sequences, MIMEx enables flexible control over the difficulty of its underlying pseudo-likell likelihood estimation. By leveraging a scalable architecture and objective function, MIMEx can be easily scaled up and stabilized to solve challenging sparse-reward tasks, achieving superior performance compared to several competitive baselines.
One limitation of MIMEx that is that it can potentially hurt performance on easier tasks that do not require exploration, as in the case of exploration bonus approaches in general [88. Moreover, while MIMEx can improve sample efficiency on hard-exploration tasks like those in PixMC-Sparse,

---

to tractably solve much harder problems like real-world sensorimotor learning, we might need a combination of MIMEx and stronger bias from other sources for example expert demonstrations. Improving exploration from either angle will be important steps for future work.
While we did not investigate into how different representation learning methods affect the perfor- mance of MIMEx in this work, adding representation learning techniques is possible in the MIMEx framework and is an interesting direction for future work. Looking forward, we hope to study how MIMEX can be leveraged to transfer knowledge from pretrained representations to the exploration problem, particularly given its scalability and flexibility.
Acknowledgements
We thank Ilija Radosavovic, Fangchen Liu, Qiyang Li, Jitendra Malik, and Alexei Efros for helpful discussions and support. TL is funded by an NSF Graduate Research Fellowship.
References
[1] S. Amin, M. Gomrokchi, H. Satija, H. van Hoof, and D. Precup. A survey of exploration methods in reinforcement learning. arXiv preprint arXiv22099000155 2021.
[2] A. Aubret, L. Matignon, and Hassas. A survey on intrinsic motivation in reinforcement learning. arXiv
preprint orXiv:9908 06976, 2019.
[3] A. Badia, B. Piot, S. Kapturowski, P. Sprechmann, A. Vitvitskyi, Z. D. Guo, and C. Blundell. Agent57: Outperforming the atari human benchmark. International Conference on Machine Learning, pages
507-517. PMLR, 2020.
[4] M. Bellemare, S. Srinivasan, G. Ostrovski, T. Schaul, D. Saxton, and R. Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
[5] M. Bellemare, J. Veness, and E. Talvitie. Skip context tree switching. In E. P. Xing and T. Jebara, editors, Proceedings of the 31 International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1458-1466, Bejing, China, 22-24 Jun 2014. PMLR
[6] M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[7] R. I. Brafman and M. Tennenholtz. R-max-a general polynomial time algorithm for near-optimal reinforce-
ment learning. Journal of Machine Learning Research, 310ct):213-231, 2002.
[8] Y. Burda, H. Edwards, A. Storkey, and O. Klimov. Exploration by random network distillation. arXiv
preprint arXiv: 1284,, 2018.
[9] J. Devlin, M.-W. Chang. K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[10] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenbom, X. Zhai, T. Unterthiner, M. Dehghani, M. Min- derer, G. Heigold, S. Gelly, et al. An image is worth 16x words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
[11] O. Eberhard, J. Hollenstein, C. Pinneri, and G. Martius. Pink noise is all you need: Colored noise exploration in deep reinforcement learning. In The Eleventh International Conference on Learning Representations, 2022.
[12] A. Ecoffet, J. Huizinga, Lehman, K. O. Stanley, and Clune. First return, then explore. Nature, 590(7847::80-5886 2021.
[13] M. Fortunato, M. G. Azar, B. Piot, J. Menick, I. Oshand, A. Graves, V. Mnih, R. Munos, D. Hassabis, O. Pietquin. et al. Noisy networks for exploration. arXiv preprint 1rXiv:1706.1033 2017.
[14] K. Goyal, C. Dyer, and T. Berg- Kirkpatrick. Exposing the implicit energy networks behind masked
language models via metropolis- hastings. arXiv preprint arXiv:2106. 02736. 2021.
[15] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to sulf-supervised learning.
Advances in neural information processing systems, 33:21271
2020.
[16] Z. Guo, S. Thakoor, M. Pislar, B. Avila Pires, F. Altché, C. Tallec, A. Saade, D. Calandriello, J.-B. Grill,
Y. Tang, et al. Byol-explore: Exploration by bootstrapped prediction. Advances in neural information processing systems,
2022.
[17] K. He, X. Chen, S. Xie, Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 16000-16009, 2022.
[18] R. Houthooft. X. Chen, Y. Duan, J. Schulman, F. De Turck, and P. Abbeel. Vime: Variational information
maximizing exploration. Advances in neural information processing systems, 29, 2016.
[19] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.699
2014.
[20] P. Ladosz, L. Weng, M. Kim, and H. Oh. Exploration in deep reinforcement learning: A survey, Information Fusion, 2022.
10

---

[21] T. P. Lillicrap. J. J. Hunt, A. Pritzel. N. Heess, T. Erez, Y. Tassa, D. Silver, and D. Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: 02971, 2015.
[22] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing
atari with deep reinforcement learning. arXiv preprint arXiv:1312.5002. 2013.
[23] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller,
A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep reinforcement learning. nature,
51897540):529-33,, 2015.
[24] G. Ostrovski. M. G. Bellemare, A. Oord, and R. Munos. Count-based exploration with neural density models. In International conference on machine learning, pages 2721-2730. PMLR, 2017.
[25] D. Pathak. P. Agrawal, A. A. Efros, and T. Darrell. Curiosity- durioe exploration by self-supervised prediction. International conference on muchine learning. pages 2778-2757 PMLR, 2017.
[26] D. Pathak, D. Candhi, and A. Gupta. Self-supervised exploration via disagreement. In International conference on machine learning, pages 5062-5071. PMLR, 2019.
[27] A. Ramesh, I.. Kirsch, S. van Steenkiste, and J. Schmidhuber. Exploring through random curiosity with general value functions, 2022.
[28] A. Ramesh, I.. Kirsch, S. van Steenkiste, and J. Schmidhuber. Exploring through random curiosity with general value functions. Advances in Neural Information Processing Systems, 35:18733-18748, 2022.
[29] J. Salazar, D. Liang, T. Q. Nguyen, and K. Kirchhoff. Masked language model scoring. arXiv preprint
urXiv.1910.16633 14659, 2019.
[30] J. Schulman, S. Levine, P. Abheel, M. Jordan, and P. Moritz. Trust region policy optimization. In
International conference on machine learning, pages 1889-1897. PMI.R 2015.
[31] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms.
arXiv preprint arXiv: 170717007444444 2017.
[32] Y. Seo, D. Hafner, H. Liu, F. Liu, S. James, K. Lee, and P. Abbeel. Masked world models for visual control.
arXiv preprint arXiv:2206, 14244, 2022.
[33] D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, and M. Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pages 387-395. PMLR, 2014.
[34] B. C. Stadie, S. Levine, and P. Abheel. Incentivizing exploration in reinforcement learning with deep predictive models. arXiv preprint 150%00050700084 2015.
[35] A. L. Strehl and M. L. Littman. An analysis of model-based interval estimation for markov decision processes. Journal of Computer and System Sciences, 1309)1330-1133 2008
[36] R. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[37] R. S. Sutton, A. G. Barto, al. Introduction to reinforcement learning. 1998.
[38] A. A. Taīga, W. Fedus, M. C. Machado, A. Courville, and M. G. Bellemare. Benchmarking bonus- exploration methods on the arcade learning environment. arXiv preprint arXiv: 02388, 2019.
[39] H. Tang, R. Houthooft, D. Foote, Stooke, O. Xi Chen, Duan, J. Schulman, F, DeTurck, and P. Abbeel. # exploration: A study of count-basee exploration for deep reinforcement learning. Advances in neural information processing systems, 30, 2017.
[40] Y. Tassa, Y. Doron, A. Muldal, Erez, Y. Li, D. d. L. Casas, D. Budden, A. Abdolmaleki, J. Merel,
A. Lefrancq, et al. Deepmind control suite. arXiv preprint 1801, 1801,00690 2018.
[41] S. B. Thrun. Efficient exploration in reinforeement learning. Technical report, Citeseer, 1992.
[42] A. Van den Oord, N. Kalchbrenner, L. Espeholt, O. Vinyals, A. Graves, et al. Conditional image generation
with pixelenn decoders. Advances in neural information processing systems, 29, 2016.
[43] A. Wang and K. Cho. Bert has a mouth, and it must speak: Bert as a markov random field language model.
arXiv preprint
(ttp:::99200009990 2019.
[44] T. Xiao, I. Radosavovic. T. Darrell, and J. Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.0611 2022.
[45] D. Yarats, R. Fergus. A. Lazaric. and L. Pinto. Mastering visual continuous control: Improved data-
augmented reinforcement leamning, arXiv preprint arXiv:2107.09645, 2021.

---

A Appendix
A.1 Clarification reference to 43
As we make a reference to 443 a work that was later retracted due to an error - we would like to clarify that we are aware of the error and our claims remain valid despite the enror in cited work.
In 43), the authors mistakenly viewed BERT as a Markov Random Field (MRF). While the original goal of 43 was to derive procedure for sampling from masked language models (MLMs) by viewing them as MRFs, the work also inspired the use of MLM prediction error as way for scoring sequences when decoding from language model. The way we propose to use MLMs is similar to the latter, i.e. as a proxy metric for scoring sequence-level predictions of trajectory information. In other words, we do not formally treat the MLM in MIMEx as an MRF, and we do not attempt to obtain conditional distributions from which one generates samples. We also note that a correct energy-basees view was later proposed in [14), which does not change the argument that put forth either. While sampling from an energy-based based model is expensive, we only seek to obtain a useful stochastic estimate of the energy function for the purpose of scoring.
A.2 Detailed reward design of the original PixMC tasks
We present the detailed reward terms of each task in the original PixMC. The total environment reward of each task is the sum of all reward terms, each multiplied with a tunable scale parameter. "Bonus" type of reward terms are discrete while other types of reward terms are continuous.
FrankaReach: distance to goal (from parallel gripper); goal bonus when distance to goal is smaller than threshold value; action penalty.
KukaReach: distance to goal (from humanoid hand); goal bonus when distance goal is smaller than a threshold value; action penalty.
FrankaCabinee distance to handle (of cabinet); handle bonus when mesh of parallel gripper intersects mesh of handle; open bonus when cabinet open; distance to goal (from parallel gripper); open pose bonus when parallel gripper is within certain pose distribution; goal bonus when distance to goal is smaller than a threshold value; action penalty.
KukaCabinet: distance to handle (of cabinet); handle bonus when mesh of humanoid hand intersects mesh of handle; open bonus when cabinet is open; distance to goal (from humanoid hand); open pose bonus when humanoid hand is within certain pose distribution; goal bonus when distance to goal is smaller than a threshold value; action penalty.
FrankaPick: distance to object (from parallel gripper); lift bonus when object is lifted above the table; distance goal (from parallel gripper); goal bonus when distance to goal is smaller than threshold value; action penalty.
distance to object (from humanoid hand); lift bonus when object is lifted above the table; distance to goal (from humanoid hand); goal bonus when distance to goal is smaller than a threshold value; action penalty.
KukaPick:
FrankaMove: distance to object (from parallel gripper); lift bonus when object is lifted above the table; distance to goal (from parallel gripper); goal bonus when distance to goal is smaller than threshold value; action penalty.
KukaMove: distance to object (from humanoid hand); lift bonus when object is lifted above the table; distance goal (from humanoid hand); goal bonus when distance to goal is smaller than threshold value; action penalty.
Due to the dense reward terms, these tasks are not sufficiently challenging for benchmarking state-of- the-art exploration algorithms.
A.3 Detailed comparison between PixMC and PixMC-Sparse tasks
We visualize partial trajectory of all PixMC-Sparse tasks in Figure 6 Below, we provide more details on how agents are rewarded differently in PixMC-Sparse compared to in PixMC.
For Reach tasks, agent needs to move its end effector (parallel gripper in the case of Franka; humanoid hand in the case of Kuka) such that the end effector reaches a specific goal location. In PixMC, agent receives variable- variable-valued valued reward at every time step; the reward value is continuous and inversely
12

---

Figure 6: Visualization of tasks. (Top-left) KukaReach. (Top-right) FrankaCabinet. (Bottom-left) KukaPick. (Bottom-right) FrankaMove.
proportional to end effector's distance goal location. In PixMC-Sparse, agent only receives beed-discrete-valued reward when its end effector intersects with a small spherical space (visualized as the green sphere) around the goal location.
For Cabinet tasks, agent needs to grasp the cabinet handle using its end effector, open the drawer, and keep pulling until it reaches a goal location. In PixMC, agent constantly receives variable-valued reward inversely proportional to end effector's distance to cabinet handle. In PiXMC-Sparse, agent only receives fixed-valued reward when the mesh of its end effector intersects with the mesh of cabinet handle (visualized as the green sphere) around the goal location.
For Pick tasks, agent needs to grasp cube-shaped object using its end effector and lift up the object until it reaches goal height. PixMC, agent constantly receives a variable-valus reward inversely proportional to end effector's distance object, and lift bonus when it starts successfully lift up the object (i.e. when object is above the table). In SixMC-Sparse, agent does not receive any distance-to-object wew-ttoooobobbsssssseeeee reward, i.e. reward until it starts to lift up the object.
For Move tasks, agent needs grasp a cube-shaped object using its end effector and lift up the object until it reaches goal location. In PixMC, agent constantly receives variable-value reward inversely proportional to end effector's s distance to object, and a lift bonus when starts to successfully lift up the object (i.e. when object is above the table). In PixMC-Sparse agent does not receive any distance-to-object-aeee reward, i.c. no reward until starts to lift up the object.
A.4 Wall-Clock Time and GPU Memory Usage Comparisons
We compare the training wall-clock time for MIMEx (ours) against that of 3 exploration baselines: random action noise (noise), intrinsic curiosity module (icm), and random network distillation (rnd). We include 3 variations of MIMEx for comparison: (1) standard model size and masking once (smaller, nm/ in Figure D. (2) standard model size and masking times (smaller, nm5 in Figure 4 (3) larger model size and masking once (smaller, nml in Figure 4h. We find that MIMEx not only achieves superior sample efficiency and performance, but also adds little wall-clock time overhead compared to the baselines.
Method Hours 10M steps
ours(1)
ours(2)
ours(3)
rnd
icm
noise 5.9
6.018
6.990
6.822
6.724
6.390
Similarly, we compare the average GPU memory usage for MIMEx against the baselines, using the same 3 variations of MIMEx for comparison as above. We find that while MIMEx does consume more memory resources during training, the increased memory usage less than 36% even when comparing the most compute-intensive MIMEx version (ours(2)) against the least compute- intensiv baseline (noise).
Method
ours( /)
ours(2)
ours(3)
rnd
icm
noise
GPU Mem Usage M
18034
24218
18402
17930
17954
17900
Fach number reported above was averaged from experiment runs on the KukaReach task in PixMC-Sparse; each run trained on single NVIDIA A100 GPU.
13

---

A.5 Additional Results on More Diverse Tasks and Environments
To check if MIMEX generalizes other domains, with ICM and RND on the ALE [6] PRIVATE EYE and VENTURE environments. Below, we report all results with 95% confidence intervals, over 5 random seeds each method (ICM, RND, MIMEx).
additional experiments compare MIMEx
per stop
10M
25M
50M
100M
200M
RND
-164.8+ 279.0
-188.4 281.9
-160 + 413.9
4105.8 + 112.0
6036.4 #7010.1
ICM
103.7
2334.8
+2095.9
7864.6 +4605.5
11963.8 746.
23377.6 9720.1
MIMEX
-529 + 416.4
104.6 658.8
2344.8 +2776.8
7859.6 + 6866.5
7717.2 6239.4
Table 4: Mean episodic return (R) at 10, 25, 50, 100, and 200 million steps of training on the Private
Eye task.
R per step
10M
25M
50M
100M
RND
227
1080
144
980 = 73
1660 + 159
ICM
380 243
1240
202
1400
310
1560 + 133
MIMEx
420
1080 114
1040 100
1660 48
Table 5: Mean episodic return (R) at 10, 25, 50, and 100 million steps of training the Venture task.
Our results for ICM and RND are similar to those reported in prior works. For MIMEX, we only tuned the exploration beta parameter; without extensive hyperparameter tuning, MIMEx still performs comparably with ICM/RND.
A.6 Additional Results on Varying Mask Distribution
We perform additional ablation studies better understand the importance of mask distribution in MIMEX formulation, especially on PixMC-Sparre tasks where existing methods (ICM, RND) fail. We interpolate between the MIMEx masking distribution and masking distribution closest to ICM/RND on the KukaReach task.
The following exploration sequence length and mask distributions used: (1) 70% uniformly random mask with sequence length (I5 umiform_77) (2) 70% fixed mask, sequence length (15 fixed 70). (3) 50% fixed mask, sequence length 5 (15 fixed 50), (4) 50% fixed mask, sequence length 2 (12 fixed 50). For "fixed mask", we only always mask out the last X positions of the input sequence based the mask ratio, effectively making the reconstruction problem future-state prediction problem. All results are reported with 95% confidence intervals and over 7 random seeds.
KukaReach
100
Success
20
0.00
0.25
1.25
1.50
Steps (M)
1e7
15 uniform_
15. fixed 70
15 fixed 50
12 fixed_50
Figure 7: Interpolating between the MIMEX masking distribution and ICM/RND masking distribution. Performance drops when going from uniformly random mask fixed mask when keeping exploration sequence length the same.
14


Reviews:
[{"review_id": "7tpP1CFjoy", "reviewer": "Reviewer_kmMp", "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. ", "questions": "1. What motivates using masked prediction loss as the intrinsic reward? \n2. 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?", "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.\n2. Stronger bias from other sources, like the expert demonstrations, is needed to improve the exploration ability of MIMEx further. \n\n\n------------\nAfter rebuttal\n---------------\nI would like to thank the authors for addressing part of my concerns. I agree to increase my score slightly. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This work proposed a general framework for deriving intrinsic rewards, which can be applied to various hard-exploration tasks. \n2. 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.\n3. 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. \n", "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. \n2. Line 81 cited a retracted paper. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback. We addressed individual comments and questions below, in particular your main concern on the validity of our method due to the citation of a retracted paper. We clarified that the key principle our method depends on is independent of the error in the retracted paper. Please let us know if there are any remaining issues that prevent our paper from getting a better score.\n\n**”The smoothness of writing can be improved.”**\n\nIn section 3.1, we describe details of MIMEx’s formulation and architecture in the form of a sequence autoencoder; we see that the section title could be potentially confusing, and will update it in our revised manuscript to “Sequence-Level Masked Autoencoders” to make for a better summary of the section and a smoother transition from the former paragraph. We will also update the writing to strengthen the motivation of using masked prediction loss as the intrinsic reward (as we will elaborate under “Questions”) and improve the overall smoothness in our next revision.\nPlease let us know if you find any other specific places where writing could be improved, and we will look into improving them accordingly.\n\n**”Line 81 cited a retracted paper.”**\n\nAs we noted on line 81 (with a footnote that points to Appendix A.1) in our submission, while [40] was later retracted due to an error, we are aware of the error and our claims remain valid despite the error in cited work. In [40], the authors mistakenly viewed BERT [9] as a Markov Random Field (MRF). While the original goal of [40] was to derive a procedure for sampling from masked language models (MLMs) by viewing them as MRFs, the work also inspired the use of MLM prediction error as a way for scoring sequences when decoding from a language model. The way we propose to use MLMs is similar to the latter, i.e. as a proxy metric for scoring sequence-level predictions of trajectory information. In other words, we do not formally treat the MLM in MIMEx as an MRF, and we do not attempt to obtain conditional distributions from which one generates samples. We also note that a correct energy-based view was later proposed in [13], which does not change the argument that we put forth either. While sampling from an energy-based model is expensive, we only seek to obtain a useful stochastic estimate of the energy function for the purpose of scoring. We will clarify the relationship between our method and [40] further in our next revision of the paper.\n\n**Questions**\n\n**1.** As we discussed in Section 3, the use of masked prediction loss is primarily motivated by its flexibility and effectiveness (as demonstrated by empirical results). Masked prediction can be applied to input sequences with arbitrary length and at arbitrary mask distribution; such a framework naturally lends itself to greater control over the difficulty of the underlying conditional prediction problem. By setting up conditional prediction problems on trajectories, we can obtain intrinsic rewards that consider transition dynamics across longer time horizons and extract richer exploration signals. We can also easily tune the difficulty of the prediction problem, by varying both the input length and the amount of conditioning context given a fixed input length. (Meanwhile, existing approaches framed as conditional prediction often consider one-step future prediction problems, which can saturate early as a useful exploratory signal; longer time-horizon prediction problems capture more complex behavior, but they can suffer from high variance.)\n\nAdditionally, we are motivated by the generality and scalability of the masked autoencoding objective. Masked autoencoding relies on less domain knowledge compared to methods like contrastive learning, and has proven success across many different input modalities.  We can also leverage standard architectures such as those used in masked language modeling and masked image modeling, for which the scalability and stability have been tested.\n\n\n**2.** Masked language modeling is connected to pseudo-likelihood through the estimation of conditional distributions. Pseudo-likelihood is a way of approximating likelihood by modeling the conditional distributions of variables given all other variables. Masked language modeling, such as the masked autoencoding objective in models like BERT, can be seen as a form of stochastic maximum pseudo-likelihood estimation, as we illustrated in section 2.3.\n\nIn masked language modeling, a model is trained to predict masked or missing tokens in a sequence. The model approximates the underlying joint distribution among variables by modeling the conditional distributions of the masked tokens given the rest of the sequence. By optimizing the model to predict the masked tokens, it effectively estimates the likelihood of the observed sequence. Approaches that estimate novelty can therefore be viewed as modeling different conditional prediction problems or masked prediction problems with different mask distributions; Table 1 provides several such examples. By applying different mask distributions, which determine the pattern of masking in the input data, various aspects of novelty can be captured. For example, masking a subset of state variables can measure novelty in states, while masking current or next-step features can estimate novelty in state transitions. Each approach represents a different way of modeling the conditional distributions and approximating the likelihood of the masked components.\n\n**Limitations**\n\nThank you for summarizing the limitations that we mentioned in Section 7 of the manuscript. We have been working on improving our exploration framework such that it can be generalizable to scenarios where exploration is not required and draw on stronger bias from sources like expert demonstrations, and will report the method and results in our follow-up works."}, {"review_id": "fiqidtX7ff", "reviewer": "Reviewer_mFvm", "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.", "questions": "1. Should the equality in the display mode equation in line 84 be an approximation since $|X|$ is not necessarily infinite?\n2. 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].\n3. Why choose the three chosen exploration baselines, random action noise, intrinsic curiosity module, and random network distillation?\n4. 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?\n5. 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.\n6. 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.\n7. Do you have ideas on extensions of Online masked prediction to offline RL and any advantages or disadvantages using MIMEx with this paradigm?\n\n[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.  \n[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.", "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. \n\nThe 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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**Originality**  \nThe 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.\n\n**Quality**  \nThe 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.\n\n**Clarity**  \nThe 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.\n\n**Significance**  \nMIMEx 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.  \n\n---\nOverall, 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.\n\nThank you for providing code with your submission for both sets of benchmarks!\n\n**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.\n\n- 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. \n- 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.\n- 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.\n\n- 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.\n\n- 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.\n\n[1] Xiao, T., Radosavovic, I., Darrell, T., & Malik, J. (2022). Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173.  \n[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.  \n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback. We addressed your main concern on the lack of diverse domains in our empirical study through experiments on two additional discrete-action environments. Our approach performs competitively against baselines, demonstrating its generalizability beyond continuous control tasks. Regarding your concern on the novelty, we would like to highlight that our main contribution is a novel framework that draws connections between masked autoencoding, pseudo-likelihood, and intrinsic rewards, as well as unifies many existing intrinsic bonus approaches. While each component of MIMEx has been explored in the past, the combination of them is what enables superior results. To the best of our knowledge, we are the first to explore using a random masked prediction objective for intrinsic bonus in the context of RL exploration.\n\n**evaluate on more diverse types of sparse rewards e.g. discrete instead of continuous**\n\nWe agree with this point, and would like to clarify that our experimental results have already covered (1) tasks with a mix of discrete and continuous rewards and (2) tasks with only discrete rewards. Examples of the former: Pick, quadruped_run; examples of the latter: Reach, cartpole_swingup_sparse. We will add a table summarizing the detailed reward terms and indicating each reward term’s type (discrete/continuous) in our revised manuscript.\n\n**more diverse tasks and environments**\n\nThank you for the suggestion. To check if MIMEx generalizes to other domains, we ran additional experiments to compare MIMEx with ICM and RND on the ALE *PRIVATE EYE* and *VENTURE* environments. Below, we report all results with 95% confidence intervals, over 5 random seeds on each method (ICM, RND, MIMEx).\n\n[*PRIVATE EYE*]\n\n| mean episodic return\\env steps | 10M | 25M | 50M | 100M |200M|\n|----------------------------|-----|-------|-------|---|---|\n| RND                        | -164.8± 279.0  | -188.4 ± 281.9   |  -160 ± 413.9   | 4105.8  ± 112.0  | 6036.4  ±7010.1  |\n| ICM                        | -31.4 ±103.7 | 2334.8 ±2095.9 | 7864.6 ±4605.5 | 11963.8 ± 746.1 | 23377.6  ± 9720.1  |\n| MIMEx                      | -529  ± 416.4 | 104.6  ± 658.8  | 2344.8  ±2776.8  | 7859.6  ± 6866.5 | 7717.2  ± 6239.4 |\n\n[*VENTURE*]\n\n| mean episodic return\\env steps | 10M  | 25M | 50M | 100M  |\n|----------------------------|------|------|------|---|\n| RND                        | 580 ± 227  | 1080 ± 144 | 980 ± 73 | 1660 ± 159 |\n| ICM                        | 380 ± 243 | 1240 ± 202 | 1400 ± 310 | 1560 ± 133  |\n| MIMEx                      | 460  ± 320  | 1080 ± 114 | 1040 ± 100  | 1660  ± 48 |\n\nOur results for ICM and RND are similar to those reported in prior works. For MIMEx, we only tuned the exploration beta parameter; without extensive hyperparameter tuning, MIMEx still performs comparably with ICM/RND. We hope these results improve confidence on the generalizability of MIMEx to other domains.\nWe will include the updated results and additional code in our revised work.\n\n**MIMEx as a \"lightweight\" module**\n\nWe used the term “lightweight” to imply programmatic simplicity rather than low memory consumption or wall-clock time. We will remove this sentence in the next revision to avoid confusion.\n\n**later exploration methods not used as baselines**\n\nThank you for the references; we find [2, 3] relevant and will include them in our revised paper. Still, we find our method and baselines more general in formulation, thus more amenable to fair comparison. The strength of MIMEx is that it is a simple module that generalizes a number of prior approaches that are still commonly used in practice, while being more flexible than these approaches. We therefore compare MIMEx against ICM and RND, two existing works that are general in formulation and strong in performance, to obtain performance difference comes from differences in algorithmic formulation rather than implementation details. An interesting direction of future work would be to use other ideas in intrinsic motivation together with our method, such as [2,3]. \n\n**better justify the baselines**\n\nWe will update the writing to elaborate on why we chose the baselines (as discussed above).\n\n**potential applications beyond the current scope**\n\nWe are particularly excited about applying MIMEx to real-world sensorimotor learning; many challenging robotic tasks are difficult to be specified with a dense reward function (e.g. screwing a water bottle cap where un-grasping and re-grasping are needed). We hope MIMEx’s promising results in simulation could transfer to real-world RL. We will add these discussions in the next revision.\n\n**Questions**\n\n1. $|X|$ is actually always finite here since we define $X$ as categorical variables. In this case, the masking distribution is a discrete uniform distribution and the equality holds.\n\n2. We use transformer blocks as in [4], with the same order and properties.\n\n3. We chose these baselines because of their effectiveness and generality across a wide range of domains and tasks.\n\n4. We include this sample curriculum to illustrate how PixMC-Sparse could be easily modified to avoid saturation as exploration algorithms get more sophisticated. We provide one example of such evaluation in Figure 2 (on *KukaPick*).\n \n5. [5] derives intrinsic rewards through predicting temporally extended general value functions, though being substantially different in terms of the environment it evaluates on (MiniGrid) and the focus (partial observability). We will include this reference in our next revision.\n\n6. Thank you for the suggestion. We will do so in our revised version.\n\n7. In the context of offline RL, online exploration is not possible since offline RL only uses previously collected data without additional online data collection. It is therefore unclear how intrinsic bonus can be applied in that setting.\n\n**Limitations**\n\nWe included additional results (see above) and runtime/memory analysis (see global PDF) to address these limitations."}, {"review_id": "DOJahJ1ph0", "reviewer": "Reviewer_if59", "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.", "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?\n\nIn 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.", "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.\n\nAnother 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.\n", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n\nI 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.\n\n---- After Author Rebuttals ----\n\nAfter 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback.\n\n**Strengthening the masked autoencoding argument**\n\nThank you for the suggestion. We agree that our method only generalizes the conceptual formulation of RND or ICM as summarized in Table 1, and does not capture differences in representation learning. We will clarify this in the revised paper. To further highlight the importance of having a flexible mask distribution, we ran additional experiments where we emulate RND/ICM through specific masks.\n\nSpecifically, we interpolated between the MIMEx masking distribution and a masking distribution closest to ICM/RND. Experiments with the following exploration sequence length and mask distributions are added:\n- l5_uniform_70 - seq length 5, uniformly random mask, 70% [MIMEx]\n- l5_fixed_70 - seq length 5, fixed mask, 70%\n- l5_fixed_50 - seq length 5, fixed mask, 50%\n- l2_fixed_50 - seq length 2, fixed mask, 50%\n(For “fixed mask”, we only always mask out the last X positions of the input sequence based on the mask ratio, effectively making the reconstruction problem a future-state prediction problem.)\n\n\nIn the global PDF, we present results on the *KukaReach* environment, each curve with 3 seeds. We observe that our method performs better than the masking strategy corresponding to RND/ICM, indicating that the flexible mask distribution enabled by our method is beneficial:\n1. Performance drops when going from uniformly random mask to fixed mask when keeping exploration sequence length the same.\n2. Performance drops further when reducing the length of exploration sequence from 5 to 2. This is consistent with our ablation results presented in Fig 4.\n3. The MIMEx variant with mask distribution closest to existing one-step exploration approaches, “l2_fixed_50”, performs the worst out of all the settings evaluated, which confirms our hypothesis that the additional flexibility in tuning masking distribution provided by MIMEx framework can positively contribute to its performance.\n\nWe believe that adding representation learning techniques such as inverse models from ICM is possible in our framework and is an interesting direction for future work. We will include the additional results and update the writing in our revised paper.\n\n**Reference to BYOL-Explore**\n\nBYOL-Explore indeed incorporates sequence-level information in its prediction error, though being very different from our work in both how the information is processed (by using a recurrent neural network) and used (to build the agent’s internal representation used for future prediction). We will include BYOL-Explore in the related work section of our revised paper and clarify these differences.\n\n**Questions**\n\nRegarding how pixel observations are encoded into latents, we follow the implementation of each respective baseline (MVP [41] for PixMC-Sparse and DrQv2 [42] for DMC). Specifically, for PixMC-Sparse, the pixels are encoded into latents via a frozen pretrained ViT model into latents; for DMC, the pixels are encoded into latents via a convolution network trained online with random shift data augmentation. In the former case, the representation does not collapse since a frozen encoder model is used; in the latter case, RL is used to train the latent representations and the random shift augmentation is used to prevent collapse. Empirically, each representation learning method has been state-of-the-art in its corresponding environment.\n\nRegarding including action in the prediction objective, we chose to mask only the latent observations because doing so would make the model more general and easier to tune (as the data modality being masked is kept consistent). We also reasoned that since we mask a history of observations, it is possible that adding action information results in more redundancy as some action information can be implicitly inferred. However, we think adding action is an interesting idea and will look into this design choice in our follow-up work.\n\n**Limitations**\n\n*partial observability*\n\nThank you for your suggestion. We agree that partial observability could be an important challenge to RL exploration methods, and that providing more information to the agent will help alleviate the challenge. Transformer-based models have been used to handle longer context length than recurrent models and to address partial observability (e.g. Chen’21, Reed’22, Brohan’22); one strength of MIMEx is therefore that it could be flexibly extended to include more data modalities or handle longer context length. Recent works have reported positive results where masked autoencoding on more general trajectory data is used for representation learning (e.g. Radosavovic’23); this is a promising sign that MIMEx could benefit from richer data too and we are excited to pursue this direction in our follow-up work.\n\nBrohan, Anthony, et al. \"Rt-1: Robotics transformer for real-world control at scale.\" arXiv preprint arXiv:2212.06817 (2022).\n\nChen, Lili, et al. \"Decision transformer: Reinforcement learning via sequence modeling.\" Advances in neural information processing systems 34 (2021): 15084-15097.\n\nRadosavovic, Ilija, et al. \"Robot Learning with Sensorimotor Pre-training.\" arXiv preprint arXiv:2306.10007 (2023).\n\nReed, Scott, et al. \"A generalist agent.\" arXiv preprint arXiv:2205.06175 (2022).\n\n*stochasticity*\n\nIndeed, for now we are tackling this problem by averaging multiple masks, and for PixMC-Sparse implementation we to some extent delegate this representation learning problem to the frozen pretrained ViT encoder. We do not focus on representation learning in the scope of this work, but will update our writing in Section 7 to include a discussion of this important problem in our revised paper.\n"}, {"review_id": "bVF0I99hHq", "reviewer": "Reviewer_SVND", "summary": "The paper introduces a novel method for Exploration in RL, called MIMEx (Masked Input Modeling for Exploration).\n\nPrevious 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.\n\nThe 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.\n\nThe paper then demonstrates the effectiveness of the approach with evaluations, baseline comparisons and ablation studies on a “PixMC-Sparse” benchmark suite.", "questions": "A few points:\n* **[a]** Can the authors address the point raised in “Weaknesses”, related to computational cost?\n* **[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?\n* **[c]** Minor point: For all plot figures, font sizes are too small.", "limitations": "I believe the authors appropriately addressed the limitations of their work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n\nThe 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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the very positive feedback. We appreciate your evaluation of our work and address the questions below.\n\n**”What is the overall runtime / resource utilization of MIMEx compared to other baselines?”** (also **Questions [a]**)\n\nThank you for the suggestion. We have included tables for both the wall-clock time and GPU memory utilization of MIMEx compared to other baselines in the global PDF, and will add this to our revised paper. While MIMEx could be more computationally demanding, the overhead in both wall-clock time and GPU memory usage is comparable to those of the baselines, so we also agree that its benefit in sample efficiency could outweigh the computational challenge.\n\n\n**Questions**\n\n**[b]** This is an interesting observation that we noticed too. We hypothesize that there is always a “sweet spot” that achieves the optimal balance of this tradeoff, rather than the longer sequence length the better: while longer exploration sequences may encourage more complex exploratory behaviors and tackle harder exploration problems, using longer exploration sequences in MIMEx may also increase the variance of prediction task and worsen performance. Our results in Figure 4, where MIMEx with exploration sequence length 6 underperforms MIMEx with shorter exploration sequence length, empirically confirm this hypothesis.\n\n\n**[c]** Thank you for the suggestion. We will increase the font sizes of all figures in our revised manuscript.\n"}]

Author rebuttal:
We thank all the reviewers for their valuable feedback. Below, we respond to each reviewer individually.

In the PDF attached below, we include additional results on wall-clock time/GPU memory usage (reviewer mFvm, reviewer SVND) and mask distribution ablation studies (reviewer if59).


Meta-review: While there is a bit of variance in the reviews, everyone recommended the paper for acceptance and the authors did a good job in rebuttal pushing the ratings up. I also agree with this assessment. There are clear room for improvement as mentioned by reviewers but the paper's contribution passes the bar to get accepted.
Decision: Accept (poster)