SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We have conducted extensive ablation studies to 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-Sparse, an easier KukaReach task and a harder KukaPick task. We summarize our findings below.
Trajectory-level exploration To our knowledge, MIMEx is the first framework that successfully incorporates sequence-level intrinsic reward to solve hard exploration tasks. Intuitively, doing so
Figure 3: Comparison with baselines. We compare our method (ours) against 3 exploration baselines: random action noise (noise), intrinsic curiosity module (icm), and random network distillation (rnd). The top figure shows results on 8 tasks from the PixMC-Sparse suite, and the bottom figure shows results on 6 tasks from the DeepMind Control Suite.
encourages agent to search for new trajectories rather than one-step 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 4, 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 a classical curse to RL algorithms [36]. In deriving intrinsic reward using MIMEx, a source of high variance is 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-based 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 it enables extremely flexible control over the difficulty of conditional prediction 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 a wide range of settings. We are therefore particularly interested in studying how varying mask distribution affects the performance 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.
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 using mask ratios  $20\%$  and  $50\%$  in Figure 5. We hypothesize that, because the difficulty of exploration varies from task to task, there exists a 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 a fixed mask ratio of  $70\%$  for our default implementation since we find it achieve a good trade-off between hyperparameter tuning and robustness of performance.
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. nm1 is masking once, and nm5 is masking 5 times followed by taking 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.
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 a mask ratio of  $70\%$  and (2) masking on 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, so opted to use the more computationally efficient time-only masking in our default implementation.
50%
Figure 5: (left) Varying mask ratio. Results from masking out  $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 a mask ratio of  $70\%$  ( $70\%$ ) and (2) masking on the feature dimension with mask ratios of  $80\%$ ,  $90\%$ , and  $95\%$  ( $80\%(a)$ ,  $90\%(a)$ ,  $95\%(a)$ ).
70% 80%(a)
90%(a) 95%(a)
Undirected exploration [41] approaches add noise to individual actions (e.g. random action noise) [11], action policy (e.g.  $\epsilon$ -greedy), or even network weights [13]. 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 count-based exploration to high-dimensional state spaces, [4] and [24] propose pseudo-count as a 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.
[8] 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 a BYOL-like [15] prediction objective. Relatedly, [28] derives intrinsic rewards through predicting temporally extended general value functions, though focusing on partial observable tabular environment. [18, 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 it is usually hard to get a good estimate of the reduction in uncertainty with all the stochasticity inherent in modern neural networks and optimization tools.
Inspired by a 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 on variable-length input sequences, MIMEx enables flexible control over the difficulty of its underlying pseudo-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 [38]. 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 performance 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.
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.
A.1 Clarification on reference to [43]
As we make a reference to [43] – 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 error in cited work.
In [43], the authors mistakenly viewed BERT[9] as a Markov Random Field (MRF). While the original goal of [43] 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 [14], 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.
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 a threshold value; action penalty.
KukaReach: distance to goal (from humanoid hand); goal bonus when distance to goal is smaller than a threshold value; action penalty.
FrankaCabinet: distance to handle (of cabinet); handle bonus when mesh of parallel gripper intersects mesh of handle; open bonus when cabinet is open; distance to goal (from parallel gripper); open pose bonus when parallel gripper is within a 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 a 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 to goal (from parallel gripper); goal bonus when distance to goal is smaller than a threshold value; action penalty.
KukaPick: 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.
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 a threshold value; action penalty.
KukaMove: 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.
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 a variable-valued reward at every time step; the reward value is continuous and inversely
Figure 6: Visualization of tasks. (Top-left) KukaReach. (Top-right) FrankaCabinet. (Bottom-left) KukaPick. (Bottom-right) FrankaMove.
proportional to end effector's distance to goal location. In PixMC-Sparse, agent only receives a fixed-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 a variable-valued reward inversely proportional to end effector's distance to cabinet handle. In PixMC-Sparse, agent only receives a 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 a cube-shaped object using its end effector and lift up the object until it reaches a goal height. In PixMC, agent constantly receives a variable-valued reward inversely proportional to end effector's distance to object, and a lift bonus when it 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-based reward, i.e. no reward until it starts to lift up the object.
For Move tasks, agent needs to grasp a cube-shaped object using its end effector and lift up the object until it reaches a goal location. In PixMC, agent constantly receives a variable-valued reward inversely proportional to end effector's distance to object, and a lift bonus when it 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-based reward, i.e. no reward until it 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, nm1 in Figure 4), (2) standard model size and masking 5 times (smaller, nm5 in Figure 4), (3) larger model size and masking once (smaller, nm1 in Figure 4). 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	ours(1)	ours(2)	ours(3)	rnd	icm	noise
Hours / 10M steps	6.018	6.990	6.822	6.724	6.390	5.9
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 is less than  $36\%$  even when comparing the most compute-intensive MIMEx version (ours(2)) against the least compute-intensive baseline (noise).
Method	ours(1)	ours(2)	ours(3)	rnd	icm	noise
GPU Mem Usage / M	18034	24218	18402	17930	17954	17900
Each number reported above was averaged from 10 experiment runs on the KukaReach task in PixMC-Sparse; each run was trained on a single NVIDIA A100 GPU.
A.5 Additional Results on More Diverse Tasks and Environments
To check if MIMEx generalizes to other domains, we ran additional experiments to compare MIMEx 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 on each method (ICM, RND, MIMEx).
R per step	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	-31.4±103.7	2334.8±2095.9	7864.6±4605.5	11963.8±746.1	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  $(\bar{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	580 ± 227	1080 ± 144	980 ± 73	1660 ± 159
ICM	380 ± 243	1240 ± 202	1400 ± 310	1560 ± 133
MIMEx	460 ± 320	1080 ± 114	1040 ± 100	1660 ± 48
Table 5: Mean episodic return  $(\bar{R})$  at 10, 25, 50, and 100 million steps of training on 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 to better understand the importance of mask distribution in MIMEx formulation, especially on PixMC-Sparse tasks where existing methods (ICM, RND) fail. We interpolate between the MIMEx masking distribution and a masking distribution closest to ICM/RND on the KukaReach task.
The following exploration sequence length and mask distributions are used: (1)  $70\%$  uniformly random mask with sequence length 5 (l5.uniform_70);, (2)  $70\%$  fixed mask, sequence length 5 (l5_fixed_70), (3)  $50\%$  fixed mask, sequence length 5 (l5_fixed_50), (4)  $50\%$  fixed mask, sequence length 2 (l2_fixed_50). 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. All results are reported with  $95\%$  confidence intervals and over 7 random seeds.
Figure 7: Interpolating between the MIMEx masking distribution and ICM/RND masking distribution. Performance drops when going from uniformly random mask to fixed mask when keeping exploration sequence length the same.