Learning Versatile Skills with Curriculum Masking

Dear Reviewer Wq3n,

We sincerely appreciate your time to review our paper and your valuable feedback. We will address your concerns in detail below.

W1 Non-stationary reward distribution may increase instability of the training process

Regarding 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.

Additionally, 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.

W2 Masking blocks let the model spend capacity on memorizing arbitrary state-action segments

We 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.

Q1 Necessity of block-wise masking

To 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.

Comparing 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.

Q2 Explanations of Figure 4(a) and discussion of longer prediction horizon

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.

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.

Q3 Stability of reward calculation and overall training dynamics with different evaluation interval $I$s

Firstly, 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.

To assess the impact of different $I$s on overall training dynamics, we choose $I$ from [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.

Dear Reviewer CzXA,

We sincerely appreciate your time to review our paper and your valuable feedback. We will address your concerns in detail below.

Comparison to existing skill learning methods with offline data

Firstly, 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.

Apart from MaskDP and MTM we have provided, regarding the two papers you mentioned, Ajay et al., 2021 and Jiang et al., 2022:

Ajay et al., 2021: We are unable to reproduce the paper, especially without publicly available code.

Jiang 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.

From 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.

Explanations on notations and intuition of the "task selection" method

Notations in Section 4.3:

$\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.

$K$: This denotes the number of arms (or masking schemes) available in the masking scheme pool $\mathcal{M}$.

$\pi$: This represents the sampling probability distribution for each masking scheme.

Intuition of Equation 3 and 4:

Equation 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.

Equation 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.

Thanks again for providing the detailed suggestions, and we will rewrite this section in the updated version accordingly.

Visualization of masking curriculum

We 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.

From 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.

Dear Reviewer znVZ,

Thank you for your insightful review of our work! We appreciate your valuable suggestions and will respond to your questions in detail.

W1 Explanations on "versatile" or "diverse" skills

In 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.

W2 Clarifications on the starting step and number of training steps in fine-tuning experiments (Fig.2)

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].

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.

W3 Interpretation of Fig.4a

Thanks for pointing this out. 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 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."

W4 Experiments on domains apart from mujoco

Thank 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.

W5 Performance of random masking reduces the significance of CurrMask

We 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.

Minor issues

We appreciate your attention to detail and pointing out all the minor issues.

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.

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, …”.

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.

Q1 The "from scratch" baseline is TD3?

Yes.

Q2 Factor in CurrMask learning overhead

It 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:

We 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.

[1]Liu F, Liu H, Grover A, et al. Masked autoencoding for scalable and generalizable decision making.

[2]Wu P, Majumdar A, Stone K, et al. Masked Trajectory Models for Prediction, Representation, and Control

Dear Reviewer f2Vb,

Thank 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:

Comparing token-wise and block-wise masking

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. 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.

We can observe from the table that on most tasks, Block-15 and Block-20 achieve higher rewards compared to token-wise masking. Additionally, as the block size increases, the performance of block-wise masking on skill-prompting also improves. These observations support our intuition that block-wise masking with relatively high block sizes helps the model better capture long-term dependencies.

Time complexity analysis

Firstly, 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. We 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.

All 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.

Visualizations of mean mask ratio in the training process and impact of mask ratio

For the added figure, we submit a PDF file in the global rebuttal. We kindly ask you to refer to the PDF for detailed figures.

We 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.

We 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.

[1]Liu F, Liu H, Grover A, et al. Masked autoencoding for scalable and generalizable decision making.