Learning Versatile Skills with Curriculum Masking
Authors: Yao Tang
Zhihui Xie
Zichuan Lin
Deheng Ye
Shuai Li
Venue: NeurIPS 2024 poster (2024)
Keywords: reinforcement learning
unsupervised pretraining
masked prediction
curriculum learning
Primary area: reinforcement_learning
TL;DR: We propose a curriculum-based masked prediction approach for unsupervised RL pretraining, which acquires skills at different complexity and achieves superior performance in various downstream tasks.

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.

Learning Versatile Skills with Curriculum Masking
Yao Tang"! Shanghai Jiao Tong University
Zhihui Xie
Zichuan Lin?
Deheng Ye?
Shuai Li#1
Tencent Al Lab
Abstract
Masked prediction has emerged as a promising pretraining paradigm in offline reinforcement learning (RL) due to its versatile masking schemes, enabling flexi- ble inference across various downstream tasks with 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. Moti- vated 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 perfor- mance 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. Code is available at here
Introduction
Humans distinguish themselves from machines by their capacity to adapt and generalize. One crucial factor behind this discrepancy the drive 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 [0..] Chebotar al. 202T), which
agent is required to learn from reward-free free offline data (Carroll et al. 2022 Schwarzer et al. 2021
or online interaction (Liu and Abbeel 2021 Yarats et al. 1021) for pretraining.
To build generic decision decision-making making agents, great efforts have been made recently to apply self- supervised learning objectives for unsupervised offline pretraining (Schwarzer al. [02 Sun et al.) 2023). Among these studies, one popular approach is masked prediction, a simple but versatile sell- supervision framework that has proven its effectiveness in domains like language (Devlin 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 masked trajectory (#1, [MASK], 82, 02, [MASK], aa), model learned by masked prediction forced reason about both dynamics (i.e., masked state 83) and behaviors (i.e., masked action a1).
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. Besides, state-action sequences naturally come with
'Equal contribution.
Correspondence Shuai Li(shaaili8 @st.isisusuuuuu...........
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

mask
3: state
Daation
Masking Pool M
Output:
$
Bg
$
a,
a₂
My;
Random Masking
Black Sizemby Mask Ratiour,
Bidirectional Transformer 0
Masked Input:
5
a
a₂
3
Mx²
5
a2
a
Block-wise Masking
Block Sizemby, Mask Rationry
(a) Masked Sequence Modeling
(b) Masking Schemes
(c) The Masking Pool
Evaluate
e
sample from J
Evaluate
sample from M
Training Process
$
S
S
S
(d) CurrMask (ours)
Figure I: Illustration of CurrMask. Based on the framework of masked prediction modeling in (a), the block-wise masking scheme (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 0 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, propose CurrMask address the above considerations. CurrMask leverages block- wise masking and a curriculum-base 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- 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 al. 2021 Pertsch et al. 2021). For example, applying -iie masking of size 3 results in (31, [MASK]. [MSK]] [MASK] $3, az). ). By using a block-wise masking scheme, the model 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, further propose a curriculum-bass 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 al., (017). 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 masking scheme.
We aim to enable the model 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 zero-shot skill prompting. zero-shot goal-conditioned planning as well as competitive fine-tuning performance on offline RI. Further analysis reveals that CurrMask effectively mitigates the issue of local correlations and is better at capturing long-term dependencies. These findings shed light the design and arrangement of masking schemes that can effectively balance the acquisition of reusable skills at varying levels of temporal granularity and complexity.
2 Related Work
Masked Prediction as a Self-Supervision Task. Masked prediction requires the model to predict 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 al. 2022; Bao et al. 2022; Xie al., 2022), and decision making (Liu 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. (023) or by learning temporal
information (Liu et al., 20r2oCarrrol 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, 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-guit guided masks (Li et al.]
[02T]] [akogeorggouu et al., [2022 Li et al.
2022) and adversarial masks (Shi al. 2022 Tomar et al. 2023) to force 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 al. 2023) or a combination of random masking with other representation learning objectives (Sun et al. (023). In this work, we aim to illustrate the connection between masking schemes and skill learning, in search for proper automatic learning curricula for masked prediction for decision 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 reward- free RI,, in which the agent can interact with the environment in the absence of rewards (Eysenbach) et al. 2019 Yarats et al.) 202T). Another setting to utilize unlabeled offline data for representation learning [Schwarzer et al. 2021 Stooke et al., 2021) or skill learning (Ajay et al.) [222]] Jiang et al.)
2022, Xie et al., (023). Masked prediction presents a promising framework to enjoy the best both world. Our work focuses on leveraging masked prediction for unsupervised RL pretraining.
Curriculum Learning. Inspired by how humans learn faster when knowledge ordered by easiness, curriculum learning (Elman 1993; [engio et al.) 2009) has been formulated for machine learning algorithms improve training efficiency. While curriculum learning has been actively explored in the context of online RL (Jabri et al. [019 Fang et al.) (202T), in this work we show that offline RL pretraining also benefits from a proper learning curriculum.
3 Preliminaries
To better contextualize our method, we provide an overview of essential background knowledge on masked prediction and curriculum learning in this section.
3.1 Masked Prediction
Let (st, ( (s1, d1, 82, a2,.... ,8,,07 ar) denote trajectory consisting of state-action sequences and DD denote the training dataset. The self-supervised task of masked prediction is to re- construct from a masked view masked(7)) (7). where (asked(-) represents 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(7) (s1, [MASK], .2, a2, ST, [MASK]). Here, [MASK] represents a special learnable token. The learning objective is then given by:
max ENND B
Σ log Po (st, ag masked
where Po 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 concrete task the model trained on. Therefore, 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 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 al. 2022). For the latter, since undesirable to specify the tasks of interest when pretraining. the random masking scheme is widely
3

---

used (i.e., uniformly sampling subset of tokens to mask). These principles form the basis of our proposed masking approach, which is elaborated in Section
3.2 Automated Curriculum Learning
Automated curriculum learning considers how arrange the order of tasks during training by adapting the selection of learning scenarios to match the learner's 's abilities. Consider a series of tasks represented by loss functions Li.... Lk. The objective is find a time-varying sequence of tasks to accelerate training. To this end, a proper automatic curriculum needs to specify two factors: I) 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 summarizes the overall pipeline. At the core of CurrMask is masked prediction as versatile self-supervised learning objective and automatic learning curriculum over masking schemes to enable fast skill discovery. Once pretrained on offine 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.
4.1 Block-wise Masking Enhances Long-term Reasoning
Our research 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 al. 2020 Bao et al. 2022) that masks the trajectory in blocks instead of individual tokens. By doing so, CurrMask pushes the model 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 al. 2022), which has been shown 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 notion of skills or primitives. Prior works in offline skill discovery (Ajay et al.) 2021 Diang et al.] 2022) typically use variational inference partition trajectories into skills. 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 explore automatic curricula that can aid in learning skills.
4.2 Learning over a Mixture of Masking Schemes
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 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 & as:
Cu(T; Σ log Pe (St, az masked (7)),
where masked, M denotes specific masking scheme. CurrMask aims to minimize the multi-task learning objective Liarget (T: Σx= (T(
4.3 Automated Curriculum Learning Boosts Training Efficiency
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 short horizon. This motivates to consider curriculum learning to facilitate masked prediction.

---

Algorithm 1: Curriculum Masking
Input: : masking pool M with cardinality K, uniformly initialized probability distribution
Tw. training steps T, evaluation interval 7, offline training dataset DD validation dataset yatt bidirectional transformer Pe
Initialize exponential weights wi & 1 for k == , K;
Compute initial target loss Lrarget (0) on validation dataset Doal with Equation 4.2 Sample initial masking scheme k ~ Twi
for < 1..
/* Training Sample masks with k compute training loss C3 on D; Update 0 by gradient descent;
i. mod - then
/* Evaluation & update of masking scheme
Compute target loss Larrget (0) on Deat's
Calculate reward using scaled target loss decrease in Equation !
Update weights w with Equation 4
Update probability distribution Tw with Equation
Update masking scheme ~ nw:
end
end
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, 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 al. 2017) proxy signal for learning progress:
= fscale (target (0) - Ltarger(f'))
(I)
where @ and 6' denote the model parameters before and after training on 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, using the 20-th percentile r/o and 80-th percentile r of unscaled history values R; = {f,li 0, 21, ..., t). where represents the evaluation interval and t some interval timestep (i.e., is divisible by T):
2 2 - (1))
(fx) max min(1,
1)).
r r
Task Selection. To schedule tasks in meaningful order, we aim to minimize Ltarget achieved after training on them sequentially, which is equivalent to maximizing the reward defined in Equation I 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 masking scheme and the goal is to maximize the total reward earned over time. Since the reward distribution induced by Equation T shifts as the network learns, we use the EXP3 algorithm (Auer al. 2002), non-stochastic multi-armed bandit algorithm that mixes the probability distribution computed using exponential weights w with the uniform distribution constituting an e fraction of the total probability distribution. The sampling probability distribution for each masking scheme is then given by:
wi
nw = (1
i i= K,
Σラ
k
W
K
where wi represents the weight of the i-th arm (or masking scheme), w/ represents the updated w; in Equation 4 K denotes the number of arms available in the masking scheme pool and represents the
5

---

sampling probability distribution for each masking scheme. This equation shows how to sample each arm (i.e., each masking scheme) using weights. The weights w 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 by the policy ww. observes reward r, and uses the importance- [(r to update its weights according to the following formula:
weighted estimator á, Tw (3)
w = wi exp (y0./K)
= = K.
This 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 Cfor discussions on how Curr addresses non-stationarity.
5 Experiments
In this section we conduct an empirical study to answer the following questions: (Q1) Can CurrMask learn 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?
5.1 Environment Setup
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 total of tasks that are associated with different environments (walker, quadruped and jaco). At evaluation, on offline datasets adapts to different downstream tasks. For more details and experimental results, please refer to Appendix [B.2 and Appendix D respectively.
test how well the model pretrained for each environment
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 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 1eft, bottom right, top 1eft, and top right). These tasks are sparse-reward tasks given that nonzero rewards are provided only when the current position is within certain distance threshold of the target position.
Dataset Collection. For each environment, 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. (022).
construct a multi-task dataset by collecting trajectories
Implementation Details. We consider multiple mask ratios R- {15%, 35%, 55%, 75%, 95%} and = multiple block sizes 8 = {1, 2,... 20} to construct the masking pool M ((r b) r E R, be B} ([M) 100) for CurrMask. For all the evaluated masked prediction methods, we use the same bidirectional encoder-decode following prior works (He al. [022]] Liu et (022). 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.
transformer architecture with a 3-layer encoder and 2-layer decoder,
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 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

---

Table 1: Skill prompting results. We report the zero-shot performance of models pretrained with different masking schemes. Results are averaged over 10 random seeds. The best and second results bold and underlined, respectively.
Reward 1
walker
waiker
walker
quad
quad.
Jaco
jaco br
II
Jaco
Average
MaskDP
1032 26
58.4 #:23
29.3 = 1.4
36.0 #22
45. 244
38.1
444
38.4
#:30
56.9 = 39
04.0 3.
36.7
MIM
107.1
58.8
27.3
37.1 #:3
42.8
72.0
71.8 4:39
72.5 5.1
77.6
Mixed inv
103.3 AM
59,5 1:30
23.8 1:2
45.1 1:30
43.2 2:
51.8 53
53.0
1:28
56.8
s 16
59.7 1:48
55.1
proed-prog
103.5
47.
25.8
40.5 18
45.6
17
85.3 2.
85.4
1:77
84.2
4.44
88.5
1:77
Mixed
110.8 22
52.2
30.5 122
43.3 +27
$1.3
=2
66.0
=6.0044
61.6 +37
62.3
36
66.5
4.
60.7
GPT
101.8 19
34.6 133
21.6
1.0
41.9 #29
48.8 4:32
86.1 55
83.1
#:7
83.9 =
85.7
3.00
65.3
CurMask
111.2 24
79.9 1:2
38,9
38.0
$1.0 34
88.4
88.5 4:36
86.0
92,9
75.0
stage, block sizes are sampled from {1, ........... 5]. {1, ....., 103, 21, 153. and (1, 20}
respectively, while the mask ratio is sampled from R; Mixed-inv adopts an inverted approach to the mask curriculum of bixee-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 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 goal state as well as past states and actions as input and predict the next state or action to reach the goal.
5.2 Downstream Tasks
To demonstrate the versatility of CurrMask, 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 diverse multi-task data is expected 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 (so; ao, $1, 01 82, 02 [MASK], [MASK], ..... [MASK]) ) and during evaluation, we set the environment state to s2, 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.
coal-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 tar- get 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 (Sstart, MASK, ..., MASK, $goal1, MASK ..., MASK, 8 gool22 MASK, we set the environment state to Sstart 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 ctt al. 2018) to perform offline RL training, following prior work (Liu et al. 2022). Although TD3 originally designed an off- ooolpc RL algorithm. Yarats et al. (2022) show that it achieves very competitive performance on offline datasets of diverse behaviors. The offline dataset collected from the entire replay buffer of a ProtoRL agent (Yarats et al.) (022)) trained for 2M environment steps. Notably, the datasets consist of highly exploratory data, which emphasizes the importance of having good representations.
5.3 Main Results
We test the versatility of CurrMask over variety of downstream tasks, in answer to Q1 and Q2
Table 2: Goal-conditioned planning results. We report the zero-shot performance of models pretrained with different masking schemes. Results are averaged over 20 random seeds. The hest and second results are bold and underlined, respectively. The lower the better.
Distance
walker_s
walker_w
walkers
quad__w
quad_
Ja
Ja
Jaca_11
Jaco
Average
Ma
4,8 2.48
10.10
6:20
15.52
1.3
20.7 0.5
21.62 0:07
1.2 2.0
1.42 2005
1.39
0.06
006
1 n.06
8.71
MIM
6.05 as8
12.20
641
17.92
0:5
23.93
25.09
2.38
2.12=0.10
2.35=
2.30 0.10
10.59
Mixed-inv
5.32 40.53
11.25 0:31
16.51 2:3
22.63 03
23.31 2m
1.55 2006
1,53 2005
1.57= 000
1.53 2000
9.47
Mixed-preg
4.96 0:48
10.18 0:2
15.77
+ 044
23.49 032
24.28 026
1.46 0.44
1.44 2.04
1.44=00
1.44 2:000
9.38
Mixed
483 1.47
10.15
15.47
20.67
21.66
1.47
1.47
1.43_
1.44 1.1
8.73
Goal-GPT
7.47 4074
15.15 4:
21.04
27.36
0m7
28.76
0.0
3.34 20000
3.58 2011
3126=01
3.50
2:111
12.61
CurMask
4.85 1.47
9.90 _037
15.31 1.0
20.47 1.3
21.39 1.67
1.39 1.16
1.38 Lasi
1.33 use
1.34 1
8.60
7

---

300
curnelative reward
250
MIN
800
700
Martine ine
200
Mard pro
600
1.0
Cuntrace (uan)
400
1.00
200
100
0
30K
15x
25k
30k
30K
5k
10k
15k
20k
25k
30k
30k
5k
10k
20k
25k
35k
ste ps (run)
Steps (stand)
steps (walk)
Figure Offline RL results. We report the finetuning performance of models pretrained with different masking schemes. Results are averaged over 10 random seeds.
Skill Prompting Results. Table summarizes the zero-shot performance for skill prompting. We observe that CurrMask achieves the best performance in 7 out of 9 tasks and also attains the highest average performance, outperforming MaskDP by margin of 32%. suggesting that CurrMask is proficient in mirroring skill prompts accomplish specific tasks. Besides, other baseline methods that incorporate block-wise masking (i.e., Mixed, Mixed-prog, and Mixed-inv) generally outperform random masking. This matches our expectation that blocks form more semantically meaningful entities than individual tokens and can be utilized by masked prediction facilitate skill learning. The only exception is Mixed-inv. The poor performance of Mixed-inv sends strong signal that proper curriculum important for masked prediction training.
Goal-conditioned Planning Results. Next, we evaluate how capable CurrMask for long-horizon planning. As shown in Table 2 CurrMask
roll out better goal-reaching trajectories than the baselines out of 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 pood-poog-aoo and Mixed-ins have worse performance than MaskDP. This indicates that the superior performance of CurrMask is not only due to block-wise masking but rather consequence of dynamically balancing what mask during training.
Offline RL Results. Finally, we present offline RL results 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 autoregressive architecture based baselines, which suggests that CurMassk 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.
5.4 Analysis
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 Curraask, conduct an ablation study the choice of block sizes. Figure BalureB 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
40
bloc
Cumalative reward
80
15
11:33
0.22
probability
0.26
6.21
60
0.20
25
0.19
CanMaal
30-
40-
0.24
0.18
madk
35%
15%
1226
15%
15
15
block size (run)
block size (reach top left)
steps
steps
(a) Impact of wlick-wise masking
(b) Impact of masking curricula (walker run)
Figure 3: Both block-wise masking and curriculum masking contribute to CurrMask's perfor- mance. Left: the performance of zero-shot skill prompting as a function fixed block size. Right: the probabilities of choosing different block sizes and mask ratios during pretraining with CurMask.

---

100
40
cumulative reward
35
Mond prog
-.8
30
APT
0.2
0.2
60
Random
CurrMask (ours)
length (run)
length (reach top left)
(a) Attention maps (reach top left)
(b) Skill prompting performance vs. rollout length
We visualize the attention map (L2-normalized over different heads) of the first decoder layer, when the model conducting skill prompting. The horizontal axis represents the keys, and the vertical axis represents the queries. Right: the performance of zero-shot skill prompting as function of rollout length (tasks; walker run and jaco reach top left).
Figure 4: Analysis of long-term prediction capability. Left:
sizes to perform zero-shot skill prompting. Besides, mixing different block sizes uniformly for training, referred to Mixed, leads 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 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 -axed-prog does not consistently lead to performance improvements compared to Random. Specifically, in most coal-conditioned planning tasks and in offline RL (shown in Figure 2 and skill prompting tasks of the walker domain, Mixed-prue -rog performs worse than Random. In contrast, for CumMask, 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 prog and CurrMask. This comparison emphasizes that a proper masking scheme and aurriculum-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 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 baxe- area highlights the model's '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 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 B
the
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 visualizes the time- varyiig 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 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 makinn 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. for further reference.
6 Conclusion
In this work, we propose CurrMask, curriculum masking approach for unsupervised RI. 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 the computational overhead, as CurrMask relies on an extra bandit model to schedule masking schemes for training"] Furthermore, the advantages offered by CurrMask could be affected by the underlying structure the environment. This encourages us to extend our method to more challenging settings like image-based RL in future research.
Acknowledgements
The corresponding author Shuai Li is sponsored by CCF-Tencent Open Research Fund.
References
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofrr Nachum. 2021. [OPAI.]: Offine Primitive Discovery for Accelerating Offline Reinforccment Learning In International Con- ference on Learning Representations.
Peter Auer, Nicolo Resa-Bianchi, Yoav Freund, and Robert E Schapire. 2002. The nonstochastic multiarmed bandit problem. SIAM journal on computing 32, (2002), 48-77.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. 2022. BEÎT: BERT Pre-Training of Image Transformers. In International Conference on Learning Representations. httops///pperrevier net/forumiidæp- -ih--id----d-----ddd BhZSz5904
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. 2009. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning. 41-48.
Raphael Boige, Yannis Flet-Berliac, Arthur Flajolet, Guillaume Richard, and Thomas PIERROT. 2023. PASTA: Pretrained Action-Statt Transformer Agents. In NeurlPS 2023 Foundation Models for Decision Making Workshop. https://openreviee wtwp:/rwwwwwwwwweeeeeeeeee
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin. Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (Eds.), Vol. 33. Curran Associates, Inc., 1877-1901. https: https://proceedinggg nwwr cc/ 2wwpee444-Papee.pppppp
paper paper_filaee fwwes/ilepppe /tt.////pppp//////////////////6.
Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anca Dragan, al. 2022. Uni [mask]: Unified inference in sequential decision problems. Advances in neural information processing systems 35 (2022), 35365-35378.
We observe training time (wall clock time) overhead of 4.7% for 100k gradient steps.
10

---

Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jacob Varley, Alex Irpan, Benjamin Eysenbach. Ryan Julian, Chelsea Finn, et al. 2021. Actionable Models: Unsupervised Offline Reinforcement Learning of Robotic Skills. In International Conference on Machine Learning. PMLR, 1518-1528.
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. 2019, BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume / (Long and Short Pupers). Association for Computational Lin- guistics, Minneapolis, Minnesota, 11166 hhttps: hopps://do2t http///o..orgg/0//// wwg/10.g///ooo http:////.o.orrr...../////////// 11t.00/111///11111111
Jeffrey L Elman. 1993. Learning and development in neural networks: The importance of starting small. Cognition 48, 1 (1993), 71-99.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. 2019.
Diversity is All
You Need: Learning Skills without a Reward Function. In International Conference on Learning Representations. https::/openzrevior:
Kuan Fang. Yuke Zhu, Silvio Savarese, and L. Fei-Fei. 2021. Adaptive Procedural Task Generation for Hard-Exploration Problems, In International Conference on Learning Representations. https
Scott Fujimoto, Herke Hoof, and David Meger. 2018. Addressing function approximation error in actor-critic methods. In International conference on machine learning. PMLR, 1587-1596.
Alex Graves, Mare G Bellemare, Jacob Menick, Remi Munos, and Koray Kavukcuoglu. 2017. Automated curriculum learning for neural networks. In international conference machine learning. PMLR, 1311-13200
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. 2022. Masked autoencoders scalable vision learners. In Proceedings of the IEEE/CVF Conference Com- puler Vision and Patlern Recognition. 16000-16009.
Allan Jabri, Kyle Hsu, Abhishek Gupta, Ben Eysenbach, Sergey Levine, and Chelsea Finn. 2019. Unsupervised curricula for visual meta-reinforreement learning. Advances in Neural Information Processing Systems 32 (2019).
Yiding Jiang, Evan Liu, Benjamin Eysenbach, Zico Kolter, and Chelsea Finn. 2022. Learning Options Compression. Advances in Neural Information Processing Systems 35 (2022), 21184- 21199.
Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S Weld, Luke Zettlemoyer, and Omer Levy. 2020. Spanbert: Improving pre-training by representing and predicting spans. Transactions of the Association for Computational Linguistics 8 (2020), 64-77.
Ioannis Kakogeorgiou, Spyros Gidaris, Bill Psomas, Yannis Avrithis, Andrei Bursuc, Konstantinos Karantzalos, and Nikos Komodakis. 2022. What to hide from your students: Attention-guicl guided masked image modeling. In Computer Eccon-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27 2022, Proceedings. Part XXX. Springer. 300-318.
Alex Lamb, Riashat Islam. Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy. and John Langford. 2022. Guaranteed discovery of controllable latent states with multi-step inverse models. arXiv preprint arXiv:2207.0822 (2022).
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. 2021. URLB: Unsupervised Reinforcement Learning Benchmark. In Deep RL Workshop NewrIPS 2021.
Tor Lattimore and Csaba Szepesvári. 2020. Bandit algorithms. Cambridge University Press.
Gang Li, Heliang Zheng, Daqing Liu, Chaoyue Wang, Bing Su, and Changwen Zheng. 2022. SemMAE: Semantic-Guided Masking for Learning Masked Autoencoders. In Advances in Neural Information Processing Systems, Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun
Cho (Eds.). https://openreee

---

Zhaowen Li, Zhiyang Chen, Fan Yang, Wei Li, Yousong Zhu, Chaoyang Zhao, Rui Deng, Liwei Wu, Rui Zhao, Ming Tang, et al. 2021. Mst: Masked self supervise transformer for visual representation. Advances in Neural Information Processing Systems 34 (2021), 13165-13176.
Fangchen Liu, Hao Liu, Aditya Grover, and Pieter Abbeel. 2022. Masked autoencoding for scalable and generalizable decision making. Advances in Neural Information Processing Systems 35 (2022), 22608-12618
Hao Liu and Pieter Abbeel. 2021. Behavior from the void: Unsupervised active pre-training, Advances in Neural Information Processing Systems (2021), 18459-18473.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. 2021. Accelerating reinforcement learning with learned skill priors. In Conference on robot learning. PMIR, 188-204.
Ilija Radosavovic, Tete Xiao, Stephen James, Pieter Abbeel, Jitendra Malik, and Trevor Darrell. 2023, Real-world robot learning with masked visual pre-training. Conference on Robot Learning. PMLR, 416-426.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R De- von Hjelm, Philip Bachman, and Aaron C Courville. 2021. Pretraining representations for data- efficient reinforcement learning. Advances in Neural Information Processing Systems 34 (2021), 12686-12699.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel, 2023. Masked world models for visual control. In Conference on Robot Learning. PMLR, 1332-1344.
Yuge Shi, N Siddharth, Philip Torr, and Adam R Kosiorek, 2022. Adversarial masking for self- supervised learning. In International Conference on Machine Learning. PMLR, 20026-20040.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. 2021. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning. PMLR, 9870-9879.
Yanchao Sun, Shuang Ma, Ratnesh Madaan, Rogerio Bonatti, Furong Huang, and Ashish Kapoor. 2023. SMART: Self-supervised Multi-task pretr/ining with contRol Transformers. In The Eleventh International Conference on Learning Representations. https:///peerrever. htwp:/errrrieee.neeffrururr Hd=9piH3Hg80Ef
Manan Tomar, Riashat Islam, Sergey Levine, and Philipp Bachman, 2023. Ignorance is Bliss: Robust Control via Information Gating, arXiv preprint arXiv:2303.0612) (2023).
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. 2020. dm imn_cotrol:: Software and tasks for continuous control. Software Impacts 6 (2020), 100022.
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. 2008. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning. 1096-1103.
Robert W White. 1959. Motivation reconsidered: the concept of competence. Psychological review 66, 5 (1959), 297.
Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. 2023. Masked Trajectory Models for Prediction, Representation, and Control. In International Conference on Machine Learning.
Zhihui Xie, Zichuan Lin, Deheng Ye, Qiang Fu, Yang Wei, and Shuai Li. 2023. Unsupervised Pretraining for Decision Transformer. Proceedings of the 40th International Conference on Machine Learning (Proceedings of Machine Learning Research, Vol. 202), Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (Eds.). PMLR, 38187-38203. httpss//proceedingg: mlr wwt.ppressvv02/xiieeesss html
12

---

Zhenda Xie, Zheng Zhang, Yue Cao. Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. 2022. Simmim: simple framework for masked image modeling. In Proceedings of the IEFF/CVE Conference Computer Vision and Pattern Recognition. 9653-9663.
Denis Yarats, David Brandfonbrener,, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. 2022. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprint arXiv:2201. (2022).
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lemrel Pinto. 2021. Reinforcement learning with prototypical representations. In International Conference on Machine Learning. PMLR, 11920-11931.
Xiang Zhou, Yi Xiong, Ningyuan Chen, and Xuefeng GAO. 2021. Regime Switching Ban- dits. In Advances in Neural Information Processing Systems, M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (Eds.), Vol. 34. Curran Associates. Inc., 4542-4554 https::/prooceddnngs http://rrocodddngg.ourooop..pppp cc/ cc/paper_fileeteeeeeeeeeeeeeeeee files pppprr.flless/pppp/////////// /aper/2021/f1lol/ 290043377eee2001110000766622222e pdf
13

---

A Pseudocode of Block-wise Masking
Algorithm demonsttates the block-wise masking mechanism, which is employed as an intermediate step for masking in CurrMask and other baseline models.
Algorithm 2: Block-wise Masking
Input
sequence length L, mask ratio p. block size b binary mask matrix m E {0, 1 } (0 for masked, for unmasked),
Output:
b then
return random raadom_maskinae masking(I.p):
end
1 p- et 1
/+ length of masked tokens */ /* number of blocks in sequence */
Initialize mask m + 1:
Randomly choose start index random(0, .b..).)..) 1): Shuffle block indices bs t- shuffle(0, ...,. C 1);
Expand block indices token indices t- (1 b + + for in bs for j in (1,... -)):
Mask tokens <- set set_zero(m, [s[-/ :]);
Len((s) then Mask remaining tokens set set_zero(m, ((O, ... L 1) (s)(1en((s) ::)):
end return m?;
B Experimental Details
B.1 Data Collection
Pretraining Datasets For all the environments, we create a multi-task dataset by gathering tra- jectories from the replay buffer of TD3 agents. We collect total of 12M steps from the replay buffer for each environment. Each task in the walker environment is trained for 4M environment steps, each task in the jaco environment is trained for 3M steps and each task in the quadruped environment is trained for 6M steps. By following this procedure, we ensure that the pretraining datasets encompasses experiences of varying quality.
For zero-shot evaluation of both skill prompting and goal-conditioned planning, construct separate validation for each environment using the same collection protocol for pretraining datasets but with different random seeds.
Validation Datasets
Training Datasets for Offline RL Each offline dataset is obtained from the complete replay buffer of ProtoRL agent, which was trained for 2M environment steps. For each task, the collected dataset is relabeled with task- rewards during offline RL. It is worth mentioning that these datasets contain highly exploratory data, emphasizing the significance of having effective representations. Table summarizes the statistics.
Table 3: Episodic return statistics of training datasets used for offline RL.
task
min
max
mean
stand
27.07
408.59
198.85
walk
4.81
199.95
72.95
run
4.55
79.00
38.61
B.2 Implementation Details
Hyperparameters Our CurrMask implementation is based the MaskDP codebase Table summarizes the hyperparameters used by CurrMask for training and evaluation.
14

---

Table 4: Hyperparameters used for model training and evaluation.
model
value
# encoder layers
3
# decoder layers
2
# autoregressive transformer layers
5
attention heads
context length
64
hidden dimension
256
mask ratio
35%, 55%, 75%, 95%]
block size
1,2,...,
training
optimizer
Adam
batch size
384
learning rate
le-4
gradient steps
300k
EXP3
0.2
EXP3 y
0.1
evaluation interval T # evaluation samples
100
10
skill prompting
# seeds
10
# trajectories sampled per seed
100
prompt length
8
rollout length
120
goal-conditioned planning
# seeds
20
# trajectories sampled per seed
100
offline RL
# seeds
10
# training steps
35k
Baselines and Implementation Details Our method CurrMask focuses on masked prediction. Therefore, we aim to compare CurrMask with the most relevant skill learning methods that are also based on masked prediction. For all the baselines, we use the same model architecture and common hyperparameters as CurrMask. The implementation of Mixed-pra prog involves evenly dividing the pretraining process into four stages based on the ratio of training steps to total training steps. We se- quentially apply four distinct mask schemes, denoted as ((r,b) r E R, bE (1.2....5)) ( E R, E (1,2,... 15)), and ((r, E R. b
r E R. b (1, 2................. 1O}}, ((r,b) (1,2..... 20}} ). This deliberate control enables the progressive increase in block size of the mask, posing greater challenges to the training procedure as it unfolds. The implementation of Mixed-inv shares significant similarities with pixod-prog. Both methods adopt four-stage approach to partition the training process. The key distinction lies the sampling of sub-sequence lengths training progresses. In the case of Mixed-ina inv. these lengths follow descending pattern: ((r,b) E R, bE {1,2..... 20) ((r,b) r R, b E (1,2. 15)), {(r.b) R, be 2..
R, e
(1,2.... ..,,, 10}), and ((r.b) For all the evaluated bidirectional masked prediction methods, we use the same encoder-decoder transformer architecture with 3-layer encoder and 2-layer decoder and the same learning objective following prior works (Liu et al. (022)). The learning objective is a design choice to compute loss on the entire input (Vincent et al.) 2008) only the masked tokens (Devlin et al., (019). In this work, we apply the former as it has been shown to work better with sequential decision-making data (Liu al. 2022).
Compute Resources CurMaask is intended to approachable the RL research community. The entire workflow can be executed on single GPU. Utilizing single RTX 3090 graphics card, the pretraining on the assembled datasets takes approximately 7-8 hours for 300k gradient steps.
Evaluation of Skill Prompting facilitate skill prompting, the agent provided with short state-action segment randomly extracted from a trajectory in the validation dataset. The agent is
15

---

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 prompt length of timesteps and the initial position of each prompt randomly sampled seajjjettrryeenngge 0.85 saaeettooy.eentthh Therefore, the prompt
within the range of [0.1 may be located at the beginning of trajectory or skewed towards the later stages, resulting in the agent's s state being in low-speed starting phase or 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 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 L.2 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 input, while the critic takes a sequence of state-action pairs as input. Both components operate without any masking. Then perform RL training using the modified architecture.
C Discussions on Non-stationarity
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. 202I). 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, 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.
alleviate this issue by rescaling rewards using historical
D Additional Experimental Results
D.1 Impact of Masking Curricula
In Figure 5 we plot the cumulative reward of each 30 steps of the generated trajectory. The patterns of Mixed -pro-rooggg 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. Curraask however exhibits faster skill acquisition and adapts its masking scheme dynamically during training.
0:30
cumulative reward
30-50
60-90
90-120
20.
20
20
10
10-
0-
50k
100k
200k
200k
100k
200k
25ck
Look
15.c
250k
300k
steps (Random)
steps (Mixed-prog)
staps (CurrMask (ours))
Figure 5: Skill prompting performance Walker run in the pretraining phase.
Figure 6 illustrrate the mean block size and mean mask ratio used by CurrMask during pretraining, as function of training steps. We can see that during the pretraining process, CurrMask continually adjusts the block size and mask ratio according to the training progress. Overall, there is upward
16

---

10.05
10.00
0.5500
beg
0.5475
0.5425
0.5375
0
50c
10ck
15ck
2002
2502
300k
50k
100k
150k
200k
250k
300k
Steos
Steas
(a) Mean block size
(b) Mcan mask ratio
Figure 6: Visualizations of mean block size and mean mask ratio during curriculum pretraining on walker domain.
Table 5: Reward results. We report the performance across Token-wise masking and Block-wise masking with different block sizes. The best and second-best results are bold and underlined, respectively.
Reward
walker,s
walker_w
walker_r
quad_w
quad_r
jaco_bl
jaco_br
jaco_1l
jaco_tr
Average
Token
1032
23.4
29.3 14
36.6
45.1 2.
58.1 4.4
58.4 30
56.9 1s
64.0
56.7
Block-5
22
314414
14.7 =5
=
33.8 1.9
24.8 14
28.2 12
35.0
Block-10
107.2
35.6111
18.3
48.3
44.
56.6
56.7
61.0
53.3
Block-15
1112 11
83.9 =37
32.2 =2
48.4=10
38.4
71.2 5.6
70.4 +
70.0 +51
76.2 1:
66.9
Block-20
112.0
19
35
49.2
41.6
84.7
85.6
84.5
73,4
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 mask schemes, CurrMask ensures steady improvement on downstream tasks.
D.2 Impact of Block size
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- prolptino are in Table D.2]
We can observe from the table that 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 skill-prompting also improves. These observations support our intuition that wise masking with relatively high block sizes helps the model better capture long-term dependencies.
E Attention Visualization
In this section, we provide 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 tokens, while goal masking masks all tokens except two randomly sampled state tokens.
Specifically, we evaluate the aforementioned masking methods on 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 truncated token sequence of length 32, resulting in a final attention map of size 32 x 32 for clearly demonstrating the differences.
Additional Results We provide additional visualization results in Figure Apart from the obser- vation that Curraask better captures long- termttermmm dependencies than Random, we find that increasing the block size for Block-wiss leads to greater capabilities long-term prediction, which supports our intuition regarding the benefits brought by block wise masking.
17

---

100
1.6
-24
22
12
Miues-ping
-22
1.0
128
16
16
22
Block-was
Block wise (30)
Stoceise (15)
Btock Mise (20)
(a) reach top left
-24
404
133
12
2.1
Random
Maeed
MixeSppng
CurMask tours)
26
25
366
0
3A
Block-iie
Block wise (30)
Bloceisse (15)
Block wise (20)
(b) walker run
26
106
45
2.4
-2.
0.2
-.1
Random
Mixed
MixeSprg
CurMMsk tours)
100
56
25
26
34
22
22
-02
23
-41
(c) quadruped run
Figure 7: Attention visualization with prompt masking.
Impact Statement
In this paper, we present a curriculum-based masked prediction approach for unsupervised RI. pretraining. Although our method is not expected to pose direct social risks, the use of extensive pre- training 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.
18

---

150
125
122
1.0e
3.00
1.00
4.15
0.1
675
3.50
R30
0.25
825
Rondom
CurtMask(ours)
350
-1S
3.0
122
+1.0
140
675
15)
(a) reach top left
se
-12
Random
Mixed
dood prog
CarMask (ouns)
L.5
Los
1.00
4A
15
Brockwisse
(b) walker run
1.25
3.00
1.00
1.0
635
130
625
prog
(onMaask(ouss)
-125
125
-125
125
1.00
1:00
1000
675
-1.0
1358
450
125
1225
Socewise (5)
Block-wlse 11.00
Blocwwise (15)
Blockwise Q211
(c) quadruped run
Figure 8: Attention visualization with goal masking.
19

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide clear statement of the paper's contributions and scope in the abstract and introduction.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of the work in Section 6
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide the detailed descriptions on the assumptions of our CurrMask in Section
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide implementation details of our method in the Section 5 and Ap- pendix
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide detailed instructions for data collection and code implementation.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all the training and test details in the Appendix B
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report all the experimental results with appropriate errors or error bars in figures.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide details experiments compute resources in the Appendix B.2
20

---

9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httpsh/htitpp...apppp....pp..ppa
Question:
Answer: [Yes]
Justification: Our research conducted in the paper conforms in every respect with the NeurIPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss both potential positive societal impacts and negative societal impacts of CurrMask the Appendix
Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [Yes]
Justification: We describe how to avoid misuse of CurMaask in Appendix F
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We respect and cite the creators or original owners of assets used in the paper.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We communicate the details of CurnMask in the paper well as the Ap- pendix
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We do not include crowdsourcing experiments and research with human subjects in this paper.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: We do not include crowdsourcing experiments and research with human subjects in this paper.
21


Reviews:
[{"review_id": "JdMzN6WxcC", "reviewer": "Reviewer_f2Vb", "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.", "questions": "See weaknesses.", "limitations": "The limitations are discussed in the last section of this paper. There is no potential negative societal impact.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1.\tThe paper is well written and easy to follow.\n2.\tThis work uncovers that the optimal combination of block size and mask ratio requires adaptive selection.\n3.\tThe analytical experiments demonstrate the ability of this method to capture long-term dependencies.", "weaknesses": "1.\tThe 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.\n2.\tThere is no time complexity analysis of this method. If possible, the authors could present how much time each baseline and this method consume respectively.\n3.\tIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer f2Vb,\n\nThank you for your valuable feedback! We appreciate your positive remarks about the soundness, presentation, and contribution of our work. Below, we address the specific points you raised:\n> Comparing token-wise and block-wise masking\n\nWe 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. We will include this experiment in our revised appendix. The results of skill-prompting are in the following table. We **bold** the highest value and *italicized* the second highest value in each column.\n\n| Reward $\\uparrow$ | walker\\_s        | walker\\_w        | walker\\_r        | quad\\_w          | quad\\_r          | jaco\\_bl         | jaco\\_br         | jaco\\_tl         | jaco\\_tr         | Average |\n|-------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|---------|\n| Token             | 103.2 ± 2.6      | 58.4 ± 2.3       | 29.3 ± 1.4       | 36.6 ± 2.2       | **45.1 ± 2.4**     | 58.1 ± 4.4       | 58.4 ± 3.0       | 56.9 ± 3.9       | 64.0 ± 3.3       | 56.7    |\n| Block-5           | 95.8 ± 2.2       | 31.4 ± 1.4       | 14.7 ± 0.5       | 37.1 ± 1.8       | 33.8 ± 1.9       | 24.4 ± 2.0       | 24.8 ± 1.4       | 25.3 ± 1.2       | 28.2 ± 1.2       | 35.0    |\n| Block-10          | 107.2 ± 2.9      | 35.6 ± 1.1       | 18.3 ± 0.9       | 48.3 ± 2.8   | 40.6 ± 2.4       |  56.6 ± 5.5       | 56.7 ± 3.3       | 55.0 ± 3.4       | 61.0 ± 3.1   |  53.3   |\n| Block-15          | *111.2 ± 1.9*    | **83.9 ± 3.7**   | *32.2 ± 2.0*     | *48.4 ± 3.0*     | 38.4 ± 2.3       | *71.2 ± 5.6*     | *70.4 ± 3.5*     | *70.0 ± 5.1*     | *76.2 ± 3.5*     |  *66.9*   |\n| Block-20          | **112.0 ± 2.5**  | *78.4 ± 1.9*     | **33.4 ± 1.5**   | **49.2 ± 2.7**   | *41.6 ± 2.6*     | **84.7 ± 5.8**   | **85.6 ± 2.0**   | **84.5 ± 4.2**   | **91.7 ± 3.0**   | **73.4**|\n\nWe 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.\n\n>  Time complexity analysis\n\nFirstly, we would like to clarify that we briefly compare the training time overhead of CurrMask versus random masking (MaskDP) in the last line of page 9. Specifically, we observed a training time (wall clock time) overhead of 4.7% for 100k gradient steps.\nWe wholeheartedly agree with the reviewer's suggestion to present the exact time consumption for each baseline and our method. To address this, we have run all the baselines on a single NVIDIA 3090 GPU and recorded pertaining time for 300k steps, as shown in the table below.\n\n| Baseline     | MaskDP |   MTM   |   Mixed   | Mixed-prog | Mixed-inv | CurrMask | \n|------------|-----------|-----------|-----------|-----------|-----------|-----------|\n| Time (seconds) | 27,450  | 28,457| 27,733 | 27,864     | 27,931    | 28,998    | \n\nAll the baselines in the table share the same model architecture. From this table, we can see that the pretraining time of CurrMask is approximately 5% longer than the fastest baseline, MaskDP. This additional time is mainly due to the overhead of evaluating pretraining progress every 100 steps and adjusting the mask scheme sampling distribution accordingly.\n\n> Visualizations of mean mask ratio in the training process and impact of mask ratio\n\nFor the added figure, we submit a PDF file in the global rebuttal. We kindly ask you to refer to the PDF for detailed figures.\n\nWe have included both the mean mask ratio and mean block size of the whole training process in **Fig. 1**. We can see that during the pretraining process, CurrMask continually adjusts the block size and mask ratio according to the training progress. Overall, there is an upward trend in block size, suggesting that CurrMask progressively enhances masking difficulty to better capture the long-term dependency. For the mask ratio, CurrMask also continuously adjusts with the mean mask ratio slightly decreasing from 0.55 to 0.54.\n\nWe also include the visualization of the impact of mask ratio on mask prediction in **Fig. 2**. Similar to Figure 3(a) in the main paper, we compare masking with constant mask ratios 0.15,0.35,0.55,0.75,0.95  with masking with Mixed_ratios and CurrMask. Mixed_ratio uniformly samples a mask ratio from [0.15,0.35,0.55,0.75,0.95] in each step. We can observe that when sampling with a constant mask ratio, a larger mask ratio tends to benefit masked prediction more. However, when sampling with mixed mask ratios, the performance surpasses all baselines that sample with a constant mask ratio. This indicates that low and high mask ratios have irreplaceable benefits for masked prediction, making sampling from multiple mask ratios an appropriate choice. This observation is similar to the conclusion  mentioned in previous work[1] that the mixed mask ratio strategy outperforms fixed mask ratio baselines.\n\n[1]Liu F, Liu H, Grover A, et al. Masked autoencoding for scalable and generalizable decision making.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7xOrrBcLQl", "reviewer": "Reviewer_znVZ", "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.", "questions": "- The \"from scratch\" baseline is TD3?\n- 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?\n\n---\nEdit after rebuttal: Questions have been addressed, I will keep the overall positive score.", "limitations": "Apart from the mentioned overhead the proposed method is rather straightforward and simple, i.e., not many limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Well written paper that is easy to follow, with clean formalization and good, readable balance between text and material (plots, diagrams, tables, algorithms). \n- Generally decent evaluation, 9 environments / tasks, each evaluated with 20 runs. Fine-tuning potential included as well, although not quite fairly assessed (see weaknesses).\n- The choice of random-masking baseline variants does cover a good range of the ablation information of CurrMask, and the results analysis (Tab.1, Fig.3&4) provides insightful understanding of the experiments.", "weaknesses": "- 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.\n- 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.\n- 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. \n- Experiments could use some different domain for comparison (apart from the mujoco domains).\n- 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.\n\n\n---\nMinor issues:\n- 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.\n- l26 missing word \"conditioned on the remaining (?), ...\"\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer znVZ,\n\nThank you for your insightful review of our work! We appreciate your valuable suggestions and will respond to your questions in detail.\n\n>W1 Explanations on \"versatile\" or \"diverse\" skills\n\nIn our study, we use the terms \"versatile skills\" or \"diverse skills\" to refer to the model's ability to learn multiple skills. According to the dictionary definition, versatile means \"able to adapt or be adapted to many different functions or activities\". In our context, versatile skills indicate that through CurrMask, the pretrained model becomes adaptable to different potential downstream tasks. We will clarify this in our updated paper.\n\n> W2 Clarifications on the starting step and number of training steps in fine-tuning experiments (Fig.2)\n\n*The starting step in Fig.2*: In the fine-tuning experiments, we compare the pretrained baselines with the “from scratch” baseline to see whether the pretraining process is helpful to offline RL. Therefore, all the baselines should be compared using the same fine-tuning steps (instead of pretraining steps) as the x-axis, starting from 0. Such comparison is common in previous offline RL pretraining works[1][2].\n\n*Training steps*: Our model architecture, data collection, pretraining, and finetuning protocols follow the previously established protocols in previous work[1], where the number of finetuning steps was set to 25k and also experimented on similar mujoco datasets. \n\n>W3 Interpretation of Fig.4a\n\nThanks for pointing this out. We will revise the explanations of Fig. 4a in the main paper. Below are the improved explanations:\n\"Fig. 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. It can be observed that the attention of CurrMask in this region is more pronounced, indicating that CurrMask is better at recalling prompt context when predicting the future compared to random masking.\"\n\n>W4 Experiments on domains apart from mujoco\n\nThank you for your suggestion regarding experiments on domains apart from Mujoco. While our current experiments focus on the Mujoco domain, we will consider extending our evaluation to other domains in future work. \n\n>W5 Performance of random masking reduces the significance of CurrMask\n\nWe would like to highlight that the effectiveness of CurrMask is demonstrated by its consistent outperformance across various tasks, indicating its ability to better capture long-term dependencies and improve model generalization. This advantage is particularly evident in scenarios where capturing complex patterns, such as long-distance predictions, is crucial, proving CurrMask to be a highly valuable technique.\n\n>Minor issues\n\nWe appreciate your attention to detail and pointing out all the minor issues.\n\n*Repeated mentions of previous work*: We would like to clarify that our primary goal in mentioning the work is to highlight the established protocols and settings we followed to ensure a fair comparison. Thanks to your advice, we will revise the wording to avoid repeated mentions in the paper.\n\n*l26 Missing Word*: Thanks for pointing out the missing word. We will rewrite this sentence for clarity as follows: “By masking a portion of the input trajectory and predicting conditioned on the remaining unmasked tokens, …”.\n\n*Title Specificity*: We appreciate your suggestion regarding the title. To better reflect the scope of our work, we will consider revising the title to be more specific to our focus on RL pertaining.\n\n>Q1 The \"from scratch\" baseline is TD3?\n\nYes.\n\n>Q2 Factor in CurrMask learning overhead\n\nIt is indeed a very interesting aspect! We experiment on MaskDP and Mixed  see the impact of the overhead computation, as the training time for CurrMask is 5% higher than these baselines. Specifically, we assess how the performance of these baselines would change if their training steps were increased by 5% (from 300k pretraining steps to 315k). The results are summarized as follows: \n\n| Reward $\\uparrow$ | walker\\_s  | walker\\_w| walker\\_r  | quad\\_w  | quad\\_r   | jaco\\_bl  | jaco\\_br  | jaco\\_tl | jaco\\_tr  | Average |\n|-------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|---------|\n|MaskDP_300k   | 103.2 ± 2.6      | 58.4 ± 2.3       | 29.3 ± 1.4       | 36.6 ± 2.2       | 45.1 ± 2.4     | 58.1 ± 4.4       | 58.4 ± 3.0       | 56.9 ± 3.9       | 64.0 ± 3.3       | 56.7    |\n|MaskDP_315k  |  103.1 ± 2.7 |52.4 ± 2.3 | 30.6 ± 0.9 | 39.3 ± 2.7 | 47.0 ± 2.8   |  59.2 ± 4.6             | 59.6 ± 3.3              | 59.5 ± 5.2          | 68.3 ± 3.7  | 57.6    |\n|Mixed_300k| 110.8 ± 2.2 | 54.2 ± 2.0 | 30.5 ± 1.2 | 43.3 ± 2.7 | 51.3 ± 2.8 | 66.0 ± 6.4 | 61.6 ± 3.7 | 62.3 ± 3.6 | 66.5 ± 4.0 | 60.7 |\n|Mixed_315k| 109.7 ± 2.3 | 55.4 ± 2.1 | 32.1 ± 1.7 | 43.1 ± 2.4 | 53.2 ± 3.1 | 67.2 ± 5.7 | 63.0 ± 4.2 | 65.3 ± 3.3 | 67.7 ± 4.1 | 61.9 |\n|CurrMask_300k    | 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.0 |\n\nWe can see that from 300k to 315k pretraining steps, the average reward improvement of MaskDP and Mixed is very limited (within 2%). After accounting for the overhead computation factor, their performance still lags behind CurrMask, demonstrating that CurrMask's benefits are not solely due to additional computation time.\n\n[1]Liu F, Liu H, Grover A, et al. Masked autoencoding for scalable and generalizable decision making.\n\n[2]Wu P, Majumdar A, Stone K, et al. Masked Trajectory Models for Prediction, Representation, and Control", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yANakquX4V", "reviewer": "Reviewer_CzXA", "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.", "questions": "Figure 1 (a), the inputs are all masked?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The idea of designing different masking curriculum to learn different skills is generally interesting and makes sense.\n\n2. The experiments, although in limited domains do showcase that the method works well for the most part.\n\n3. Overall writing is great. Hyperparameters used in the experiments are provided in the appendix for reproducibility.", "weaknesses": "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.\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer CzXA,\n\nWe sincerely appreciate your time to review our paper and your valuable feedback. We will address your concerns in detail below.\n\n>Comparison to existing skill learning methods with offline data\n\nFirstly, we would like to clarify that our paper focuses on masked prediction. Therefore, we aim to compare our method with the most relevant skill learning methods that are also based on masked prediction. In our paper, MaskDP and MTM are the closest baselines that fit this criterion, providing a fair comparison.\n\nApart from MaskDP and MTM we have provided, regarding the two papers you mentioned, Ajay et al., 2021 and Jiang et al., 2022:\n\nAjay et al., 2021: We are unable to reproduce the paper, especially without publicly available code.\n\nJiang et al., 2022: The experimental environments used in this work are grid world and tasks in the visual domain, which are significantly different from our state-based mujoco domain. We follow the hyperparameters in the paper and adapt their method LOVE to our dataset and environment and train 35k steps for a fair comparison. The evaluated results are in the following table.\n\n| Reward $\\uparrow$ | walker\\_stand | walker\\_walk | walker\\_run |\n|-------------------|------------------|------------------|------------------|\n| MaskDP |  790.9  ±  141.8   |   439.8   ± 178.9    |   174.7   ±  38.0      |\n| MTM | 786.5  ±   140.3  |   459.8   ± 150.5    |    177.1      ±  39.3  |\n| LOVE | 201.4  ±  122.7   |  135.8   ±  61.7    |  74.6    ±  47.8      |\n| CurrMask |  833.9 ±  89.5   | 699.2   ±   79.4    |   274.4    ±   61.3    |\n\nFrom the table, we can observe that LOVE cannot perform reasonable results under our experimental setting and requires further tuning. We hope this clarifies our rationale and the efforts we made to provide a fair and relevant comparison in our paper. To enhance clarity, we will revise our paper to include more detailed considerations regarding the selection of baselines.\n\n>Explanations on notations and intuition of the \"task selection\" method\n\n*Notations in Section 4.3*:\n\n$\\omega_i$ and  $\\omega'_i$: $\\omega_i$ represents the weight of the $i$-th arm (or masking scheme) and $\\omega'_i$ represents the updated $\\omega_i$  in Equation 4.\n\n$K$: This denotes the number of arms (or masking schemes) available in the masking scheme pool $\\mathcal{M}$.\n\n$\\pi$: This represents the sampling probability distribution for each masking scheme.\n\n*Intuition of Equation 3 and 4*:\n\nEquation 3: This equation shows how to sample each arm (i.e., each masking scheme) using weights. The weights $\\omega$​ determine the probability distribution $\\pi$. We sample the masking schemes based on $\\pi$. This probabilistic approach allows us to balance exploration and exploitation by dynamically adjusting the focus on different masking schemes based on their performance.\n\nEquation 4: This equation describes how we update the arm weights based on the observed rewards. By updating the weights $\\omega$ 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.\n\nThanks again for providing the detailed suggestions, and we will rewrite this section in the updated version accordingly.\n\n> Visualization of masking curriculum\n\nWe would like to clarify that we have included the visualization curriculum masking and the analysis of the impact on performance in Appendix D. We kindly ask the reviewer to refer to the section for more details. We also include a more detailed visualization of the masking curriculum in Figure 1 of the global rebuttal pdf. \n\nFrom Appendix D and Fig. 1 in the global rebuttal pdf, we can see that during the pretraining process, CurrMask continually adjusts the block size and mask ratio according to the training progress. Overall, there is an upward 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 will revise the paper to highlight this aspect thanks to this advice.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QbySdZcUZ8", "reviewer": "Reviewer_Wq3n", "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.", "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\n- 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? \n- 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": "The authors adequately discussed the limitations of the proposed method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed framework creatively introduces curriculum skill learning to masked prediction for RL pretraining, which facilitates long-term reasoning and training efficiency.\n- The proposed method generally brings performance improvements across different downstream tasks.\n- Comprehensive analysis is performed to understand the effectiveness of each component in the proposed method.\n- The paper is well-written and well-structured.", "weaknesses": "- 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\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer Wq3n,\n\nWe sincerely appreciate your time to review our paper and your valuable feedback. We will address your concerns in detail below.\n\n>W1 Non-stationary reward distribution may increase instability of the training process\n\nRegarding the concern about non-stationary reward distribution potentially increasing the instability of the training process, we would like to kindly direct you to Appendix C of our paper, where we have a detailed discussion. Despite the non-stationary nature of the reward distribution, our method leverages the inherent adaptability of EXP3 and rescales rewards using historical percentiles to ensure that the learning process progresses gradually without sudden regime shifts.\n\nAdditionally, in Appendix D, we have a detailed discussion on the impact of the masking curriculum during the pretraining process. We encourage the reviewer to refer to Figure 5, which demonstrates that CurrMask enables more stable and consistent improvement in skill prompting performance during the pretraining phase compared to other methods.\n\n>W2 Masking blocks let the model spend capacity on memorizing arbitrary state-action segments\n\nWe understand your concern that the model may spend capacity on memorizing arbitrary state-action segments. However, during pretraining, it is inherently challenging to determine which information will be effective for downstream tasks, regardless of how we adjust the pretraining paradigm. Given this limitation, our strategy is to explicitly provide a more versatile pool of potential reusable behaviors during pretraining. By doing so, we aim to enhance the model's learning by constructing a masking pool consisting of masking schemes of different mask ratios and block sizes and let the model adaptively adjust the learning progress with curriculum masking. \n\n\n>Q1 Necessity of block-wise masking\n\nTo achieve the goal of mixing masking schemes, CurrMask constructs the masking pool with masking schemes that are a combination of high/low ratios and large/small block sizes. Simply mixing different mask ratios intuitively lacks the scenario of combining low mask ratios with large block sizes. To empirically demonstrate the necessity of block-wise masking, we can compare the performance of MaskDP and Mixed on skill-prompting. Both MaskDP and Mixed samples multiple mask ratios. MaskDP leverages token-wise masking and Mixed leverages block-wise masking. Specifically, both MaskDP and Mixed randomly sample a mask ratio from [0.15, 0.35, 0.55, 0.75, 0.95] and Mixed samples a block size from [1,2,...,20] at each pertaining step. The results of skill-prompting are as follows. We **bold** the highest reward in each column.\n\n| Reward $\\uparrow$ | walker\\_s  | walker\\_w | walker\\_r  | quad\\_w   | quad\\_r  | jaco\\_bl  | jaco\\_br  | jaco\\_tl | jaco\\_tr  | Average |\n|-------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|---------|\n| MaskDP| 103.2 ± 2.6 | **58.4 ± 2.3**  | 29.3 ± 1.4  | 36.6 ± 2.2 | 45.1 ± 2.4  | 58.1 ± 4.4 | 58.4 ± 3.0  | 56.9 ± 3.9       | 64.0 ± 3.3    | 56.7    |\n| Mixed | **110.8 ± 2.2** | 54.2 ± 2.0  | **30.5 ± 1.2**   | **43.3 ± 2.7**       | **51.3 ± 2.8**   | **66.0 ± 6.4** | **61.6 ± 3.7**       | **62.3 ± 3.6**       | **66.5 ± 4.0**    | **60.7**    |\n\nComparing MaskDP and Mixed, we can see that even with mixed mask ratios, applying block-wise masking still benefits the downstream task performance. This observation suggests that merely combining mask ratios cannot replace the effect of leveraging block-wise masking. \n\n>Q2 Explanations of Figure 4(a) and discussion of longer prediction horizon\n\n*Explanations of Figure 4(a)*: We will revise the explanations of Fig. 4a in the main paper. Below are the improved explanations. Fig. 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. It can be observed that the CurrMask's attention in this region is more pronounced, indicating that CurrMask is better at recalling prompt context when predicting the future compared to random masking. \n\n*Longer prediction horizon*: We visualize the attention map where CurrMask predicts the attention map for all 64 tokens based on the first 8 unmasked tokens in Figure 3 of the global rebuttal PDF. The trend that CurrMask's attention in the black-boxed region is more pronounced still holds true.\n\n>Q3 Stability of reward calculation and overall training dynamics with different evaluation interval $I$s\n\nFirstly, in Section 4.3 (L75-179), we introduced our approach of rescaling rewards into the [-1, 1] range using historical percentiles. This ensures that for different $I$s, the reward distribution remains stable within this range.\n\nTo assess the impact of different $I$s on overall training dynamics, we choose $I$ from \n[20,50,100(CurrMask),200,500] for visualization. In the global rebuttal PDF, Figure 4 illustrates the pretraining dynamics under these different  $I$s. From the figure, we can observe that the training dynamics are relatively normal for different $I$ values. Comparing the impact of different $I$s on pretraining dynamics, it is evident that larger $I$ values ($I = 200,500$) result in greater fluctuations in loss compared to smaller $I$ values ($I = 20, 50, 100$). This is intuitively understandable, as a larger $I$ means the model pretrains for more steps under the same mask scheme, leading to a greater reduction in loss for the current mask scheme. When the model resamples a new mask scheme after training for $I$ steps, the loss exhibits more noticeable upward fluctuations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We give our sincere thanks to all the reviewers for their insightful suggestions and positive feedback of our work. For all the figures that we add for further discussion, we submit a PDF file in this global rebuttal.

Meta-review: 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.
Decision: Accept (poster)