Learning Versatile Skills with Curriculum Masking

Authors: yao  tang
zhihui  xie
zichuan  lin
deheng  ye
shuai  li
block  size=b
mask  ratio=r
size=b n block
ratio=r n mask
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:23:21

Abstract: Masked prediction has emerged as a promising pretraining paradigm in offline reinforcement learning (RL) due to its versatile masking schemes, enabling flexible inference across various downstream tasks with a unified model. Despite the versatility of masked prediction, it remains unclear how to balance the learning of skills at different levels of complexity. To address this, we propose CurrMask, a curriculum masking pretraining paradigm for sequential decision making. Motivated by how humans learn by organizing knowledge in a curriculum, CurrMask adjusts its masking scheme during pretraining for learning versatile skills.  Through extensive experiments, we show that CurrMask exhibits superior zero-shot performance on skill prompting tasks, goal-conditioned planning tasks, and competitive finetuning performance on offline RL tasks. Additionally, our analysis of training dynamics reveals that CurrMask gradually acquires skills of varying complexity by dynamically adjusting its masking scheme.

Main paper: 
 introduction: 
 Humans distinguish themselves from machines by their capacity to adapt and generalize. One crucial factor behind this discrepancy is the drive to acquire reusable knowledge (e.g., concepts and behaviors) even in the absence of explicit reward (White, 1959) . This has motivated research in unsupervised reinforcement learning (RL) (Laskin et al., 2021; Chebotar et al., 2021) , in which the agent is required to learn from reward-free offline data (Carroll et al., 2022; Schwarzer et al., 2021) or online interaction (Liu and Abbeel, 2021; Yarats et al., 2021) for pretraining.To build generic decision-making agents, great efforts have been made recently to apply selfsupervised learning objectives for unsupervised offline pretraining (Schwarzer et al., 2021; Sun et al., 2023) . Among these studies, one popular approach is masked prediction, a simple but versatile selfsupervision framework that has proven its effectiveness in domains like language (Devlin et al., 2019) and vision (He et al., 2022) . By masking a portion of the input trajectory and predicting it conditioned on the remaining unmasked tokens, the model can not only capture rich representations but also learn transferable behaviors. For example, given a masked trajectory (s 1 , [MASK], s 2 , a 2 , [MASK], a 3 ), a model learned by masked prediction is forced to reason about both dynamics (i.e., masked state s 3 ) and behaviors (i.e., masked action a 1 ).Given the effectiveness of masked prediction, an important question arises: how can we design and arrange the masking schemes for decision-making data to maximize its benefits? To explore this question, our research stems from the finding that models trained with token-wise random masking, a widely adopted masking strategy in natural language modeling, fall short in modeling long-term dependencies (see Figure 4 ). While randomly masked words describe semantically dense information, state-action sequences in decision-making data contain heavy information redundancy (e.g., consecutive states are usually similar). This causes the model to predict masked tokens solely based on their neighboring unmasked tokens. Based on the framework of masked prediction modeling in (a), the block-wise masking scheme in (b), we design a masking pool M in (c) consisting of masking schemes at different levels of complexity, which CurrMask evaluates the learning progress of the model θ and samples masking schemes from during pretraining in (d) .unique pattern of interleaved modality, unlike single-modal word sequences. These discrepancies in information density and sequence pattern require a reevaluation of masking scheme design for RL pretraining.To this end, we propose CurrMask to address the above considerations. CurrMask leverages blockwise masking and a curriculum-based approach arranging the order of different masking schemes to boost pretraining. An illustration of our approach is shown in Figure 1 . First, we posit that masking for decision-making data needs to be designed in blocks rather than tokens. A block of consecutive state and action tokens forms a semantic entity of skill (Ajay et al., 2021; Pertsch et al., 2021) . For example, applying block-wise masking of size 3 results in (s 1 , [MASK], [MASK] , [MASK] , s 3 , a 3 ). By using a block-wise masking scheme, the model is compelled to prioritize global dependencies over basic local correlations when making masked predictions. Moreover, the combination of multiple block sizes and mask ratios incentivizes the model to acquire adaptable skills that can be utilized for diverse downstream tasks.With different block-wise masking schemes representing skills at varying levels of complexity, we further propose a curriculum-based approach to arrange the order of learning them for better capturing semantically meaningful relations. Our main intuition is that since humans obtain knowledge by organizing it into different subjects and learning the entire curriculum in an easy-to-hard order, the ability of long-horizon reasoning can be developed by first learning how to act locally. This motivates us to measure the learning progress and schedule proper masking schemes at regular intervals to boost pretraining. Ideally, the curriculum can boost the pretrained model performance on downstream tasks. However, obtaining downstream information is usually intractable during pretraining. Instead, we consider the pretraining task, i.e., the reconstruction loss, as a proxy reference for learning progress. Utilizing this proxy for skill learning, the goal of selecting masking schemes can be interpreted as maximizing the decrease in reconstruction loss (Graves et al., 2017) . Considering the exploitation of the reconstruction loss decrease and the uncertainties in pretraining that require exploration of different masking schemes, the masking scheme selection task can be formulated as a multi-armed bandit problem (Lattimore and Szepesvári, 2020) , where each arm represents a masking scheme.We aim to enable the model to acquire versatile skills through CurrMask, meaning that the pretrained model becomes adaptable to different potential downstream tasks. We conduct a series of empirical studies on various MuJoCo-based control tasks, including locomotion and robotic arm manipulation.Our results demonstrate that CurrMask enables the learning of a versatile model that achieves superior performance in zero-shot skill prompting, zero-shot goal-conditioned planning as well as competitive fine-tuning performance on offline RL. Further analysis reveals that CurrMask effectively mitigates the issue of local correlations and is better at capturing long-term dependencies. These findings shed light on the design and arrangement of masking schemes that can effectively balance the acquisition of reusable skills at varying levels of temporal granularity and complexity. 
related work: 
 Masked Prediction as a Self-Supervision Task. Masked prediction requires the model to predict a missing portion of the input that has been held out. Pretraining via masked prediction has been explored in natural language processing (Devlin et al., 2019; Joshi et al., 2020) , computer vision (He et al., 2022; Bao et al., 2022; Xie et al., 2022) , and decision making (Liu et al., 2022; Carroll et al., 2022; Sun et al., 2023; Wu et al., 2023; Boige et al., 2023) . One of the most notable applications is masked language modeling (Devlin et al., 2019) for learning transferable text representations.Recently, it has been shown that masked prediction can also facilitate decision making, either by training visual backbones (Radosavovic et al., 2023; Seo et al., 2023) or by learning temporal information (Liu et al., 2022; Carroll et al., 2022; Sun et al., 2023; Wu et al., 2023; Boige et al., 2023) .In line with this research direction, we investigate the impact of masking schemes on pretraining.Masking Schemes. For masked prediction, a central question is what is masked. One common scheme is to randomly mask some of the tokens from the input. Apart from random masking, recent studies have proposed attention-guided masks (Li et al., 2021; Kakogeorgiou et al., 2022; Li et al., 2022) and adversarial masks (Shi et al., 2022; Tomar et al., 2023) to force the model to focus specific parts of the input for better performance. In the domain of decision making, previous work usually considers the simplest random masking scheme (Liu et al., 2022) , manually designed task-specific masks (Carroll et al., 2022) , random autoregressive masking (Wu et al., 2023) or a combination of random masking with other representation learning objectives (Sun et al., 2023) . In this work, we aim to illustrate the connection between masking schemes and skill learning, in search for a proper automatic learning curricula for masked prediction for decision-making data.Unsupervised RL Pretraining. Our work also falls into the category of extracting prior knowledge without extrinsic human supervision for sample-efficient RL. Previous work has vastly studied rewardfree RL, in which the agent can interact with the environment in the absence of rewards (Eysenbach et al., 2019; Yarats et al., 2021) . Another setting is to utilize unlabeled offline data for representation learning (Schwarzer et al., 2021; Stooke et al., 2021) or skill learning (Ajay et al., 2021; Jiang et al., 2022; Xie et al., 2023) . Masked prediction presents a promising framework to enjoy the best of both world. Our work focuses on leveraging masked prediction for unsupervised RL pretraining.Curriculum Learning. Inspired by how humans learn faster when knowledge is ordered by easiness, curriculum learning (Elman, 1993; Bengio et al., 2009) has been formulated for machine learning algorithms to improve training efficiency. While curriculum learning has been actively explored in the context of online RL (Jabri et al., 2019; Fang et al., 2021) , in this work we show that offline RL pretraining also benefits from a proper learning curriculum. 
preliminaries: 
 To better contextualize our method, we provide an overview of essential background knowledge on masked prediction and curriculum learning in this section.Let τ = (s t , a t )T t=1 = (s 1 , a 1 , s 2 , a 2 , • • • , s T , a T) denote a trajectory consisting of state-action sequences and D denote the training dataset. The self-supervised task of masked prediction is to reconstruct τ from a masked view masked(τ ), where masked(•) represents a specific masking function. For example, if masked(•) represents a deterministic scheme that masks the initial and final actions of the input, the resulting masked trajectory is masked(τ ) = (s 1 , [MASK], s 2 , a 2 , • • • , s T , [MASK]).Here, [MASK] represents a special learnable token. The learning objective is then given by:max θ E τ ∼D T t=1 log P θ (s t , a t | masked(τ )) ,where P θ is parameterized by a bidirectional transformer (Devlin et al., 2019) . By reconstructing state-action sequences, the model learns to reason over temporal dependencies.Importantly, the choice of masked(•) specifies a concrete task the model is trained on. Therefore, it is crucial to design an appropriate masking scheme that enables learning of general relationships in state-action sequences. This goal boils down to two aspects: 1) how much is masked, and 2) what is masked. For the former, it has been shown that a high mask ratio (e.g., 95%) is meaningful for decision-making data due to its low information density (Liu et al., 2022) . For the latter, since it is undesirable to specify the tasks of interest when pretraining, the random masking scheme is widely used (i.e., uniformly sampling a subset of tokens to mask). These principles form the basis of our proposed masking approach, which is elaborated in Section 4.Automated curriculum learning considers how to arrange the order of tasks during training by adapting the selection of learning scenarios to match the learner's abilities. Consider a series of tasks represented by loss functions L 1 , . . . , L K . The objective is to find a time-varying sequence of tasks to accelerate training. To this end, a proper automatic curriculum needs to specify two factors: 1) how to measure learning progress, in order to adjust its task schedule dynamically, and 2) how to perform task selection based on progress signals. We describe our design in Section 4. 
curriculum masked prediction: 
 In this section, we describe the proposed approach, CurrMask, for unsupervised RL pretraining. Algorithm 1 summarizes the overall pipeline. At the core of CurrMask is masked prediction as a versatile self-supervised learning objective and an automatic learning curriculum over masking schemes to enable fast skill discovery. Once pretrained on offline data, CurrMask can perform various downstream tasks in a zero-shot manner, or be finetuned for policy learning. In the following, we elaborate the design of CurrMask and provide sufficient explanation.Our research is based on the discovery that models trained using random masking, a commonly used strategy in natural language modeling, fall short in capturing long-term dependencies (see Figure 4 ). This is undesirable for decision-making agents that maximize long-term reward. To overcome this issue, CurrMask applies the block-wise masking scheme (Joshi et al., 2020; Bao et al., 2022) that masks the trajectory in blocks instead of individual tokens. By doing so, CurrMask pushes the model to focus on semantically meaningful abstractions rather than simple local correlations. Predicting missing blocks of state-action sequences also resembles multi-step inverse dynamics models (Lamb et al., 2022) , which has been shown to learn robust representations for decision making. We present pseudocode of our block-wise masking implementation in Appendix A.Blocks consisting of consecutive states and actions also form a notion of skills or primitives. Prior works in offline skill discovery (Ajay et al., 2021; Jiang et al., 2022) typically use variational inference to partition trajectories into skills. In this work, we argue that masked prediction with block-wise masking represents an alternative approach for offline skill discovery. The link between masked prediction and skill discovery inspires us to explore automatic curricula that can aid in learning skills.The block size explicitly determines the level of temporal granularity for masked prediction. To capture both short-term and long-term temporal dependencies, CurrMask employs a combination of masking schemes with varying block sizes and mask ratios during pretraining.Given a set of masking schemes M where |M| = K, we define the loss function for masked prediction task k as:L k (τ ; θ) = T t=1 log P θ (s t , a t | masked k (τ )) ,where masked k ∈ M denotes a specific masking scheme. CurrMask aims to minimize the multi-task learning objectiveL target (τ ; θ) = 1 K K k=1 L k (τ ; θ).A key feature of mixed masking schemes is their inherent variability in complexity. Intuitively, the ability of reasoning over global dependencies can be developed by first learning how to plan within a short horizon. This motivates us to consider curriculum learning to facilitate masked prediction. By scheduling masking schemes in a meaningful order, we expect that the model will learn more efficiently and quickly during training.Evaluation of Learning Progress. The first factor to determine is the measure of learning progress. Ideally, we would like the curriculum to maximize the rate at which the model learns to solve downstream tasks. However, it is usually intractable to measure without downstream task information. Hence, we consider target loss decrease (Graves et al., 2017) as a proxy signal for learning progress:EQUATIONwhere θ and θ ′ denote the model parameters before and after training on a masking scheme for an interval I, respectively. The underlying rationale is that the most efficacious masking scheme is evidenced by the greatest target loss decrease observed before and after training on it. To alleviate the issue of time-varying magnitudes, we follow Graves et al. (2017) to rescale values into [-1, 1] using the 20-th percentile r lo t and 80-th percentile r hi t of unscaled history values R t = {r i |i = 0, I, 2I, ..., t}, where I represents the evaluation interval and t is some interval timestep (i.e., t is divisible by I):EQUATIONTask Selection. To schedule tasks in a meaningful order, we aim to minimize L target achieved after training on them sequentially, which is equivalent to maximizing the reward defined in Equation 1.In the selection of masking schemes, we need to balance the maximization of the reward with the uncertainties inherent in pretraining dynamics, requiring deliberate exploration among diverse masking schemes. This can be formulated a multi-armed bandit problem (Lattimore and Szepesvári, 2020) , where each arm represents a masking scheme and the goal is to maximize the total reward earned over time. Since the reward distribution induced by Equation 1 shifts as the network learns, we use the EXP3 algorithm (Auer et al., 2002) , a non-stochastic multi-armed bandit algorithm that mixes the probability distribution computed using exponential weights w with the uniform distribution constituting an ϵ fraction of the total probability distribution. The sampling probability distribution π w for each masking scheme is then given by:EQUATIONwhere ω i represents the weight of the i-th arm (or masking scheme), ω ′ i represents the updated ω i in Equation 4, K denotes the number of arms available in the masking scheme pool and π represents the sampling probability distribution for each masking scheme. This equation shows how to sample each arm (i.e., each masking scheme) using weights. The weights ω determine the probability distribution π from which we sample the masking schemes. This probabilistic approach allows us to balance exploration and exploitation by dynamically adjusting the focus on different masking schemes based on their performance.Each time EXP3 samples an arm k by the policy π w , observes reward r, and uses the importanceweighted estimator xi = I{i=k}r πw(i) to update its weights according to the following formula:EQUATIONThis equation describes how we update the arm weights based on the observed rewards, which in this context are the computed pretraining progress. By updating the weights ω according to the rewards, we ensure that the more effective masking schemes (those that lead to better pretraining progress) are sampled more frequently in subsequent iterations. This adaptive mechanism helps the model to learn an effective masking curriculum over time. As such, exponential growth significantly increases the probability of choosing good arms (i.e., masking schemes). Please see Appendix C for discussions on how CurrMask addresses non-stationarity. 
experiments: 
 In this section we conduct an empirical study to answer the following questions: (Q1) Can CurrMask learn a versatile model that achieves good performance on a variety of downstream tasks, both in zero-shot and finetuning scenarios? (Q2) What role do block-wise masking and masking curricula play in CurrMask? (Q3) Does CurrMask better capture long-term temporal dependencies, and if so, what mechanism within CurrMask facilitate this capability?We evaluate our method on a set of environments from the DeepMind control suite (Tunyasuvunakool et al., 2020) . Each environment has several tasks specified by how the reward function is defined. Specifically, we consider a total of 9 tasks that are associated with 3 different environments (walker, quadruped and jaco). At evaluation, we test how well the model pretrained for each environment on offline datasets adapts to different downstream tasks. For more details and experimental results, please refer to Appendix B.2 and Appendix D, respectively.Environments. The walker environment consists of 3 locomotion tasks (run, stand, and walk) and the quadruped environment provides 2 locomotion tasks (run and walk). All the tasks provide a dense reward measure of task completion. For example, task run provides rewards encouraging forward velocity. We also conduct experiments on jaco, which is an environment for robot arm manipulation including 4 reaching tasks (bottom left, bottom right, top left, and top right). These tasks are sparse-reward tasks given that nonzero rewards are provided only when the current position is within a certain distance threshold of the target position. Dataset Collection. For each environment, we construct a multi-task dataset by collecting trajectories of 12M steps from the replay buffer of TD3 agents (Fujimoto et al., 2018) . This collection procedure ensures that the pretraining dataset contains experiences of varying quality. For zero-shot evaluation, we additionally construct a validation set for each environment using the same protocol but with different random seeds, following the setting in the prior work (Liu et al., 2022) . Implementation Details. We consider multiple mask ratios R = {15%, 35%, 55%, 75%, 95%} and multiple block sizes B = {1, 2, . . . , 20} to construct the masking poolM = {(r, b) | r ∈ R, b ∈ B} (|M| = 100)for CurrMask. For all the evaluated masked prediction methods, we use the same bidirectional encoder-decoder transformer architecture with a 3-layer encoder and a 2-layer decoder, following prior works (He et al., 2022; Liu et al., 2022) . The encoder input is unmasked states and actions and the decoder input is the whole trajectory including both masked and unmasked tokens.The evaluated autoregressive baselines consist of 5 layers for fair comparison.Baselines. We compare CurrMask with the following baselines: MaskDP (Liu et al., 2022) samples a mask ratio from R and randomly masks a portion of individual tokens in each training step; MTM (Wu et al., 2023) adopts random autoregressive masking which first randomly samples masked tokens and all future tokens of the last masked token are masked; Mixed randomly samples a masking scheme from M with a uniform distribution; Mixed-prog uses a manually designed mask curriculum that progressively increases the block size during pretraining, divided into four stages. In each 101.8 ± 2.9 34.6 ± 1.3 21.6 ± 1.0 41.9 ± 2.9 48.8 ± 3.2 86.1 ± 5.7 83.1 ± 2.7 83.9 ± 5.1 85.7 ± 3.0 65.3 CurrMask 111.2 ± 2.4 79.9 ± 1.2 38.9 ± 1.9 38.0 ± 2.2 51.0 ± 3.4 88.4 ± 5.1 88.5 ± 3.6 86.0 ± 4.3 92.9 ± 3.5 75.0stage, block sizes are sampled from {1, 2, . . . , 5}, {1, 2, . . . , 10}, {1, 2, . . . , 15}, and {1, 2, . . . , 20} respectively, while the mask ratio is sampled from R; Mixed-inv adopts an inverted approach to the mask curriculum of Mixed-prog, utilizing the block size pools in reverse order. Apart from masked prediction baselines, we also compare with GPT and Goal-GPT, which are auto-regressive models similar to GPT (Brown et al., 2020) . GPT takes the past states and actions as inputs and outputs the next state or action. To accommodate downstream tasks that require information about the goal state, Goal-GPT is modified to take a goal state as well as past states and actions as input and predict the next state or action to reach the goal.To demonstrate the versatility of CurrMask, we consider various downstream tasks that require different capabilities, including zero-shot inference by specifying certain masking schemes (i.e., skill prompting and goal-conditioned planning) and adaptation via finetuning (i.e., offline RL). Skill Prompting. A unified model trained on diverse multi-task data is expected to acquire various skills that can be invoked to perform certain tasks. Skill prompting tests this ability by requiring the model to generate consecutive behaviors given a short state-action sequence. An input example of which prompt context length is 3 looks like (s 0 , a 0 , s 1 , a 1 , s 2 , a 2 , [MASK], [MASK], ..., [MASK] ) and during evaluation, we set the environment state to s 2 , perform predicted actions and record rewards.For each task, we sample prompt contexts of length eight from the validation set, and evaluate the quality of the generated trajectory of length 120 by its task rewards.Goal-conditioned Planning. Another type of downstream task we consider is goal-conditioned planning. Starting from a given state, the model needs to roll out actions that can achieve target goals within a number of steps. We condition the pretrained model a start state and four goal states to evaluate the model's capability to generate long-term plans. With the input pattern of (s start , MASK, ..., MASK, s goal1 , MASK, ..., MASK, s goal2 , MASK, ...), we set the environment state to s start perform the predicted actions. The performance is assessed by the L2 distance between each goal and the state that is achieved by executing the predicted actions within the given time budget and is closest to the goal. We choose the goal states at distances of 20, 40, 60, and 80 future timesteps from the start states to evaluate the long-term planning capability of model. Offline RL. Finally, we study if the representations learned by CurrMask can accelerate offline RL.For each task, we add a critic head and actor head on top of the encoder and run TD3 (Fujimoto et al., 2018) to perform offline RL training, following prior work (Liu et al., 2022) . Although TD3 is originally designed as an off-policy RL algorithm, Yarats et al. (2022) show that it achieves very competitive performance on offline datasets of diverse behaviors. The offline dataset is collected from the entire replay buffer of a ProtoRL agent (Yarats et al., 2021) trained for 2M environment steps. Notably, the datasets consist of highly exploratory data, which emphasizes the importance of having good representations.We test the versatility of CurrMask over a variety of downstream tasks, in answer to Q1 and Q2. 2 , CurrMask can roll out better goal-reaching trajectories than the baselines in 8 out of 9 tasks. We can observe that Goal-GPT exhibits the worst performance in all the tasks, suggesting that the autoregressive model falls short in downstream tasks that necessitate the simultaneous use of bidirectional information, compared to directional models. Another notable observation is that, in contrast to skill prompting results, Mixed, Mixed-prog and Mixed-inv have worse performance than MaskDP. This indicates that the superior performance of CurrMask is not only due to block-wise masking but rather a consequence of dynamically balancing what to mask during training.Offline RL Results. Finally, we present offline RL results in Figure 2 . Compared with learning from scratch, learning with pretrained representations obtained by CurrMask results in training speedup and performance improvement. We observe that CurrMask generally outperforms other masked pretraining baselines as well as autoregressive architecture based baselines, which suggests that CurrMask not only learns diverse skills but also extracts transferable representations for policy learning. It should also be noted that in some cases (e.g., walk and run) pretraining with GPT leads to diminished performance for finetuning, whereas CurrMask is generally more stable.In this section, we investigate several aspects of CurrMask to further answer Q2 and Q3. Impact of Block-wise Masking. To better understand how block-wise masking contributes to CurrMask, we conduct an ablation study on the choice of block sizes. Figure 3a shows the influence of the block size, where masked prediction is combined with a fixed block size and mask ratios randomly sampled from R. With block-wise masking, masked prediction benefits from larger block sizes to perform zero-shot skill prompting. Besides, mixing different block sizes uniformly for training, referred to as Mixed, leads to average performance. By leveraging a proper combination of masking schemes with different block sizes, CurrMask boost the skill prompting performance matching or exceeding the best results among all the block-wise masking schemes shown in Figure 3a . This indicates that the block size does have a great impact over final performance, and a proper learning curriculum can be also crucial.Impact of Masking Curricula. Another important question is whether good masking curricula should be determined manually or found adaptively during training. We would like to emphasize that Mixed-prog does not consistently lead to performance improvements compared to Random. Specifically, in most goal-conditioned planning tasks and in offline RL (shown in Figure 2 ) and skill prompting tasks of the walker domain, Mixed-prog performs worse than Random. In contrast, for CurrMask, we consistently observe improvements over Random.We provide additional experimental results to better illustrate that CurrMask is not just rediscovering the programmatic curriculum. Figure 5 displays the skill prompting results versus training steps during pretraining, revealing noticeable differences in skill learning progress between Mixed-prog and CurrMask. This comparison emphasizes that a proper masking scheme and curriculum-based pretraining progress are unlikely to be predetermined and highlights the benefits of CurrMask for its adaptivitity.Evaluation of Long-term Prediction. One of the most important intuition behind CurrMask is that block-wise masking can enhance the model's capability to capture long-term dependencies. To verify this, we look into the attention maps during prediction with skill prompts even when they are far from current timesteps. Figure 4a shows how CurrMask predicts the attention map for all 32 tokens based on the first 8 unmasked tokens and the subsequent 24 masked tokens. The vertical axis represents the query, and the horizontal axis represents the key. The black-boxed area highlights the model's dependency on the keys of the first 8 unmasked tokens when predicting the attention map for the subsequent 24 masked tokens. We notice significant differences in how the approaches use the prompt. It can be observed that the attention of CurrMask in the bounded region is more pronounced, indicating that CurrMask is better at recalling prompt context when predicting the future compared to random masking. Besides, the predictions of CurrMask attend to prior actions more than they do to prior states. These findings support our intuition that CurrMask is more effective at extracting useful long-term dependencies. We offer a more comprehensive assessment in Appendix E. This discrepancy in attention patterns is further validated by the performance of long-horizon skill prompting. Figure 4b shows the skill prompting performance as a function of the rollout length. We observe that CurrMask outperforms the baselines significantly when the rollout length is extended. Notably, Random, Mixed and Mixed-inv have degenerated performance for long rollouts on the right task, supporting our hypothesis that CurrMask acquires non-trivial long-term prediction capacity.Visualization of Masking Curricula. Next, we investigate how automatic curricula steer masking schemes during training. Figure 3b visualizes the time-varying probabilities of choosing different block sizes and mask ratios during CurrMask pertaining. We can see that CurrMask gradually increases the probability of choosing large block sizes while also preferring a moderate mask ratio.The former observation reveals that CurrMask has a tendency to learn more complex skills, which aligns with our intuition. For the latter, we believe it reflects the degree of information redundancy in sequential decision-making data, also reported in previous work (Liu et al., 2022; He et al., 2022) .We also visualize the mean block size and mean mask ratio during pretraining in Appendix D.1 for further reference. 
conclusion: 
 In this work, we propose CurrMask, a curriculum masking approach for unsupervised RL pretraining. Motivated by the unique pattern of sequential decision-making data (i.e., low information density and interleaved modality), we propose to apply block-wise masking with mixed mask ratios and block sizes to capture temporal dependencies at both short-term and long-term levels of granularity. As different masking schemes naturally vary in prediction difficulty, we consider automated curriculum learning as the inner drive to facilitate training by scheduling these schemes in a meaningful order. We show through extensive experiments that CurrMask learns a versatile model that consistently outperforms the baselines in various downstream tasks. Our analysis of the impact of block-wise masking and curriculum learning emphasizes the adaptivity of CurrMask and its superior ability to extract global dependencies.Limitations. One limitation of CurrMask is the computational overhead, as CurrMask relies on an extra bandit model to schedule masking schemes for training1 . Furthermore, the advantages offered by CurrMask could be affected by the underlying structure of the environment. This encourages us to extend our method to more challenging settings like image-based RL in future research. For all the evaluated bidirectional masked prediction methods, we use the same encoder-decoder transformer architecture with a 3-layer encoder and a 2-layer decoder and the same learning objective following prior works (Liu et al., 2022) . The learning objective is a design choice to compute loss on the entire input (Vincent et al., 2008) or only on the masked tokens (Devlin et al., 2019) . In this work, we apply the former as it has been shown to work better with sequential decision-making data (Liu et al., 2022) .Compute Resources CurrMask is intended to be approachable to the RL research community. The entire workflow can be executed on a single GPU. Utilizing a single RTX 3090 graphics card, the pretraining on the assembled datasets takes approximately 7-8 hours for 300k gradient steps.Evaluation of Skill Prompting To facilitate skill prompting, the agent is provided with a short state-action segment randomly extracted from a trajectory in the validation dataset. The agent is then positioned at the final state of the segment and tasked with generating subsequent behaviors in an autoregressive manner. The quality of the generated sequence is evaluated by comparing its accumulated rewards with those obtained from the rollout of an expert with advanced skills. In detail, we employ a prompt length of 8 timesteps and the initial position of each prompt is randomly sampled within the range of [0.1 • trajectory length, 0.85 • trajectory length]. Therefore, the prompt may be located at the beginning of a trajectory or skewed towards the later stages, resulting in the agent's state being in a low-speed starting phase or a high-speed running phase in the cases of the walk/run task.Evaluation of Goal-conditioned Planning To implement goal-conditioned planning, we randomly sample a goal context of length 100 from the trajectories in the validation set. The position of the goal is set at specific locations [20, 40, 60, 80] . The agent is initially placed at the starting position of the goal context, and the rollout continues for the remaining tokens. We calculate the L2 distance between each goal state and its closest state token within the rollout length as a metric for evaluation.Evaluation of Offline RL In offline RL, the main objective is to train a model to maximize the return for a specific task, as defined by a reward function. This differs from our self-supervised pretraining objective, so additional finetuning is required. To align with the RL setting, we modify the bidirectional attention mask in the transformer to a causal attention mask. This change allows the model to attend only to previous states and actions during training, simulating the sequential nature of RL tasks. We also utilize a standard actor-critic framework similar to TD3 by incorporating a critic head and an actor head on top of the pretrained encoder. The actor takes a sequence of states as input, while the critic takes a sequence of state-action pairs as input. Both components operate without any masking. Then we perform RL training using the modified architecture. 
 
 Appendix: Non-stationarity is a major challenge for algorithm design in our context. We want to emphasize two important properties: 1) The reward distribution is non-stationary, and 2) Despite this, the learning process usually progresses gradually without sudden regime shifts (Zhou et al., 2021) .For the former, we want to emphasize that our method tackles the non-stationarity in two aspects. Firstly, EXP3, a special case of online mirror descent, is inherently adaptable to reward distributions that change over time. Secondly, we alleviate this issue by rescaling rewards using historical percentiles. For the latter, while abrupt distribution changes are not typically observed, we believe that our framework can easily accommodate other techniques like sliding windows and reward discounting to address significant non-stationarity.In Figure 5 , we plot the cumulative reward of each 30 steps of the generated trajectory. The patterns of Mixed-prog and CurrMask are substantially different. During the initial stage of pretraining, Mixed-prog (trained with small blocks only) struggles to learn skills at all levels of temporal granularity. CurrMask however exhibits faster skill acquisition and adapts its masking scheme dynamically during training. trend in block size, suggesting that CurrMask progressively enhances masking difficulty. For the mask ratio, CurrMask also continuously adjusts, CurrMask also continuously adjusts, slightly decreasing from 0.55 to 0.54. Through continual adjustments on mask schemes, CurrMask ensures steady improvement on downstream tasks.We have conducted experiments on token-wise masking and block-wise masking with constant block sizes of 5,10,15 and 20 to prove the significance of block-wise masking empirically. The results of skill-prompting are in Table D .2. We can observe from the table that on most tasks, Block-15 and Block-20 achieve higher rewards compared to token-wise masking. Additionally, as the block size increases, the performance of block-wise masking on skill-prompting also improves. These observations support our intuition that block-wise masking with relatively high block sizes helps the model better capture long-term dependencies.In this section, we provide more details about our attention map visualization and additional results.Setup To provide a clearer visualization of the differences between CurrMask and other baselines in zero-shot skill prompting and goal reaching, we visualize the attention maps of their first layer decoders.For comparison purposes, we employ two masking techniques: prompt masking and goal masking. Prompt masking masks all tokens except the first 8 tokens, while goal masking masks all tokens except two randomly sampled state tokens. Specifically, we evaluate the aforementioned masking methods on 10 trajectories randomly sampled from validation sets using the pretrained model. We then compute the average attention map for each technique, and finally apply L2 normalization to the attention maps of all four heads to obtain the first layer attention map. We focus on a truncated token sequence of length 32, resulting in a final attention map of size 32 × 32 for clearly demonstrating the differences.We provide additional visualization results in Figure 6 -7. Apart from the observation that CurrMask better captures long-term dependencies than Random, we find that increasing the block size for Block-wise leads to greater capabilities in long-term prediction, which supports our intuition regarding the benefits brought by block-wise masking. In this paper, we present a curriculum-based masked prediction approach for unsupervised RL pretraining. Although our method is not expected to pose direct social risks, the use of extensive pretraining datasets highlights the significance of avoiding harmful bias in training data. We emphasize the need for ethical responsibility to ensure that our method contributes positively to societal and technological progress. We observe a training time (wall clock time) overhead of 4.7% for 100k gradient steps.https://github.com/FangchenLiu/MaskDP_public 


Decision: The paper presents a curriculum learning method for RL, CurrMask, that learns skills from masked trajectories. The reviewers found the paper, among others,  well written, the idea appealing, and performance improvements on the downstream tasks. After the rebuttal, three out of four reviewers found most of their concerns addressed (the remaining reviewer did not respond), with all ratings keeping a positive trend leaning toward acceptance. Therefore, I recommend the paper to be accepted.

Reviews: review: paper_summary: paper_summary: This paper proposes a curriculum masking pretraining paradigm for RL training, which is based on the block-wise masking schemes and is able to decide the block size and mask ratio automatically. Specifically, the authors design a masking pool with different masking scheme of different block size and mask ratio. Given the target loss and the corresponding reward from the environment, this method formulates the masking selection as a multi-armed bandit problem and sample the masking scheme from the masking pool according to the updated policy. The experiments on control tasks demonstrate the effectiveness of this method.
strength_weakness: strength_weakness: 1.	The authors have not conducted experiments to compare token-wise and block-wise masking before they choose the block-wise masking. Despite the theoretical analysis in the Introduction section, it would be better to prove this choice through experimentation.
2.	There is no time complexity analysis of this method. If possible, the authors could present how much time each baseline and this method consume respectively.
3.	In figure 6, the authors report the mean block size of the whole training process, so I am curious about the mean mask ratio of the whole training process. Furthermore, is it possible to visualize the impact of mask ratio just like Figure 3(a), so that we can know whether masked prediction benefits from larger mask ratio or smaller ratio.
questions: questions: See weaknesses.
limitations: limitations: The limitations are discussed in the last section of this paper. There is no potential negative societal impact.
paper_summary: paper_summary: This work proposes a curriculum learning approach to reinforcement learning skills from masked trajectory sequences. Given a set of pre-collected environment samples (here from a TD3 agent in 9 different mujoco domains), the proposed MaskCurr curriculum treats the agent's progress (target-loss decrease) like a two-armed bandit problem, where one of the bandits is the amount of information that is covered up (mask-ratio) and the other bandit is the length of the covered gaps that the agent has to reconstruct (block-size). Evaluation is tested on two tasks, skill-prompting, which requires the trained agent to complete a starting sequence, and goal-planning, where masked trajectories with interspersed checkpoints (goals) have to be filled in such that the intermediate goals are reached. CurrMask is compared to a set of various masking techniques (i.e., variations of random masking) as well as GPT variants, and is shown to perform competitively in both of the above mentioned tasks.
strength_weakness: strength_weakness: - Not clear how "versatile" or "diverse" skills are classified here. I understand these properties w.r.t. to how different single tasks behaviors are emerging, rather than "learn multiple tasks", which I believe is meant here.
- The fine-tuning experiments (Fig.2) seem to be plotted unfairly against the "from scatch" baseline, since pretrainined models do not "start" at step 0 (but at -#pretraining-steps) on the x-axis. For the training itself, 25k steps for mujoco domains is rather brief and all training curves seem to be stopped mid-training. As such the actual suitability for fine-tuning is rather questionable, or at least not shown with significance.
- The attention-map result of Fig.4 a) could be better interpreted in the main paper, there is a visual difference but its not quite clear to me how these maps correlate to the claimed useful long-term dependencies skills. The appendix provides some more insight, but all relevant information and explanations of main paper plots should be in the main paper as well. 
- Experiments could use some different domain for comparison (apart from the mujoco domains).
- While the improvement of CurrMask compared to the baselines is shown, the significance of the results is lowered by the fact on how good the random masking techniques still perform.


---
Minor issues:
- Repeated mentions to "following previous/prior work Liu et al.2022" which give the impression of self-reference. If this is not a self-citation please clarify the wording for the double blind reviewing standard.
- l26 missing word "conditioned on the remaining (?), ..."
- The title could be more specific, claiming generally versatile skills for (basically only) the mujoco domain feels a bit far-fetched. The scope of this work is not broad enough to warrant such a sweeping claim.
questions: questions: - The "from scratch" baseline is TD3?
- Although CurrMask mostly outperforms the random baselines, did you perhaps try factor in the overhead of the CurrMask learning into the evaluation? I.e., if the random variants would increase their training time by the observed 4.7% wall-clock time overhead, how would your estimation be on the performance comparison given in Table 1?

---
Edit after rebuttal: Questions have been addressed, I will keep the overall positive score.
limitations: limitations: Apart from the mentioned overhead the proposed method is rather straightforward and simple, i.e., not many limitations.
paper_summary: paper_summary: The paper presents a method that learns skills through curriculum masking. Specifically, the approach CurrMask can automatically arranges the order of different masking schemes for training. The algorithm is tested on Deepmind Control Suite tasks, and show positive results in representation learning, zero-shot skill prompting, and zero-shot goal-conditioned planning.
strength_weakness: strength_weakness: 1. My main concern is that there is no comparison to other existing skill learning methods with offline data, e.g. the two papers (Ajay et al., 2021; Jiang et al., 2022) mentioned in the paper's related work section. As a skill learning/discovery paper, at least one of the existing approaches with the similar setting should be empirically compared with.

2. There is not enough explanation for the proposed "task selection" method, which I believe is the central part of the proposed approach. Specifically, in section 4.3, what do \omega_i, \omega'_i, K denote? Does \pi represent the policy? What is the intuition behind the two equations in the task selection subsection? Without explaining these, it is hard for me to understand how the proposed approach learn the masking curriculum.

3. There is no visualization of the proposed masking curriculums. As this is the central contribution of the paper, it would be very interesting to see what the actual masking curriculum is for those continuous control tasks and how is affect the numerical results.
questions: questions: Figure 1 (a), the inputs are all masked?
limitations: limitations: Yes
paper_summary: paper_summary: This paper presents CurrMask, a novel masked prediction approach for unsupervised RL pretraining, which learns skills of different complexity through block-wise masking and adaptively adjusts masking schemes in a curriculum for training efficiency. In contrast to previous methods that perform random masking at the token level, CurrMask applies a pool of masking schemes with different block sizes to capture temporal dependencies in various scales. In addition, CurrMask trains a bandit model to schedule the masking schemes to facilitate pre-training using the target loss decrease as the reward. The method is extensively evaluated in three downstream task settings across different environments from the DeepMind control suite. The experimental results demonstrate the strong empirical performance of CurrMask for both zero-shot inference and finetuning, which generally outperforms the compared baselines.
strength_weakness: strength_weakness: - The masked prediction pretraining is interleaved with a bandit training process with non-stationary reward distribution, which could increase the instability of the training process
- A skill usually refers to a meaningful abstraction beyond simply consecutive states and actions (e.g. a state-action sequence completing a subtask/subgoal). Therefore, not all state-action sequences are useful in downstream tasks. Masking blocks randomly may let the model spend lots of capacity on memorizing arbitrary state-action segments, instead of capturing reusable behaviors that can be efficiently transferred to downstream applications.
questions: questions: - It puzzles me whether block-wise masking is needed necessarily. If mask ratios are randomly chosen from a range, high mask ratios would naturally generate masked blocks of varying sizes, and low mask ratios would also likely ask the model to perform token-level reconstruction, equivalently implementing the mixture of masking schemes, which seems to make the block-wise masking redundant
- It remains unclear to me how Figure 4(a) serves as evidence of long-term prediction capability. And how would the attention map change if the prediction horizon gets longer? 
- Given that the reward is computed every $I$ training step(s), how stable the reward calculation and the overall training dynamics is with different $I$s?
limitations: limitations: The authors adequately discussed the limitations of the proposed method.

Comments: title: Rebuttal Acknowledgement
Rebuttal Acknowledged
comment: We would like to thank you for your positive response. Please let us know if we have addressed all your concerns or if you have any further questions, and we will respond promptly.
Thank you very much for your positive score. We are happy to discuss further if there are any additional questions.
I'd like to thank the authors for taking the time and discussing the weaknesses in detail. I think the overall positive trend leaning accept is adequate and I will keep my score.
Thank you for your detailed response. My concerns have been addressed and I will keep my overall positive score.
Thank you for taking the time to review our work and give insightful advice. We appreciate your suggestion regarding the exploration of more efficient paradigms for skills pre-training and agree that this is a promising direction for future research. We're glad that our responses were able to address most of your concerns, and we are grateful for your positive rating.
I'd like to thank the authors for their detailed responses to my concerns and questions. I think a more efficient paradigm for skills pre-training should be explored but I do agree with the author's rationale on this problem. Most of my concerns have been addressed, and I will keep my positive rating.