Learning Versatile Skills with Curriculum Masking

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.

Strengths

- The proposed framework creatively introduces curriculum skill learning to masked prediction for RL pretraining, which facilitates long-term reasoning and training efficiency. - The proposed method generally brings performance improvements across different downstream tasks. - Comprehensive analysis is performed to understand the effectiveness of each component in the proposed method. - 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 - A skill usually refers to a meaningful abstraction beyond simply consecutive states and actions (e.g. a state-action sequence completing a subtask/subgoal). Therefore, not all state-action sequences are useful in downstream tasks. Masking blocks randomly may let the model spend lots of capacity on memorizing arbitrary state-action segments, instead of capturing reusable behaviors that can be efficiently transferred to downstream applications.

Questions

- It puzzles me whether block-wise masking is needed necessarily. If mask ratios are randomly chosen from a range, high mask ratios would naturally generate masked blocks of varying sizes, and low mask ratios would also likely ask the model to perform token-level reconstruction, equivalently implementing the mixture of masking schemes, which seems to make the block-wise masking redundant - It remains unclear to me how Figure 4(a) serves as evidence of long-term prediction capability. And how would the attention map change if the prediction horizon gets longer? - Given that the reward is computed every $I$ training step(s), how stable the reward calculation and the overall training dynamics is with different $I$s?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors adequately discussed the limitations of the proposed method.

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.

Strengths

1. The idea of designing different masking curriculum to learn different skills is generally interesting and makes sense. 2. The experiments, although in limited domains do showcase that the method works well for the most part. 3. 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. 2. There is not enough explanation for the proposed "task selection" method, which I believe is the central part of the proposed approach. Specifically, in section 4.3, what do \omega_i, \omega'_i, K denote? Does \pi represent the policy? What is the intuition behind the two equations in the task selection subsection? Without explaining these, it is hard for me to understand how the proposed approach learn the masking curriculum. 3. There is no visualization of the proposed masking curriculums. As this is the central contribution of the paper, it would be very interesting to see what the actual masking curriculum is for those continuous control tasks and how is affect the numerical results.

Questions

Figure 1 (a), the inputs are all masked?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes

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.

Strengths

- Well written paper that is easy to follow, with clean formalization and good, readable balance between text and material (plots, diagrams, tables, algorithms). - Generally decent evaluation, 9 environments / tasks, each evaluated with 20 runs. Fine-tuning potential included as well, although not quite fairly assessed (see weaknesses). - 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. - The fine-tuning experiments (Fig.2) seem to be plotted unfairly against the "from scatch" baseline, since pretrainined models do not "start" at step 0 (but at -#pretraining-steps) on the x-axis. For the training itself, 25k steps for mujoco domains is rather brief and all training curves seem to be stopped mid-training. As such the actual suitability for fine-tuning is rather questionable, or at least not shown with significance. - The attention-map result of Fig.4 a) could be better interpreted in the main paper, there is a visual difference but its not quite clear to me how these maps correlate to the claimed useful long-term dependencies skills. The appendix provides some more insight, but all relevant information and explanations of main paper plots should be in the main paper as well. - Experiments could use some different domain for comparison (apart from the mujoco domains). - While the improvement of CurrMask compared to the baselines is shown, the significance of the results is lowered by the fact on how good the random masking techniques still perform. --- Minor issues: - Repeated mentions to "following previous/prior work Liu et al.2022" which give the impression of self-reference. If this is not a self-citation please clarify the wording for the double blind reviewing standard. - l26 missing word "conditioned on the remaining (?), ..." - The title could be more specific, claiming generally versatile skills for (basically only) the mujoco domain feels a bit far-fetched. The scope of this work is not broad enough to warrant such a sweeping claim.

Questions

- The "from scratch" baseline is TD3? - Although CurrMask mostly outperforms the random baselines, did you perhaps try factor in the overhead of the CurrMask learning into the evaluation? I.e., if the random variants would increase their training time by the observed 4.7% wall-clock time overhead, how would your estimation be on the performance comparison given in Table 1? --- Edit after rebuttal: Questions have been addressed, I will keep the overall positive score.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Apart from the mentioned overhead the proposed method is rather straightforward and simple, i.e., not many limitations.

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.

Strengths

1. The paper is well written and easy to follow. 2. This work uncovers that the optimal combination of block size and mask ratio requires adaptive selection. 3. The analytical experiments demonstrate the ability of this method to capture long-term dependencies.

Weaknesses

1. The authors have not conducted experiments to compare token-wise and block-wise masking before they choose the block-wise masking. Despite the theoretical analysis in the Introduction section, it would be better to prove this choice through experimentation. 2. There is no time complexity analysis of this method. If possible, the authors could present how much time each baseline and this method consume respectively. 3. In figure 6, the authors report the mean block size of the whole training process, so I am curious about the mean mask ratio of the whole training process. Furthermore, is it possible to visualize the impact of mask ratio just like Figure 3(a), so that we can know whether masked prediction benefits from larger mask ratio or smaller ratio.

Questions

See weaknesses.

Rating

5

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

The limitations are discussed in the last section of this paper. There is no potential negative societal impact.

Rebuttal Acknowledged

I'd like to thank the authors for their detailed responses to my concerns and questions. I think a more efficient paradigm for skills pre-training should be explored but I do agree with the author's rationale on this problem. Most of my concerns have been addressed, and I will keep my positive rating.

Rebuttal Acknowledgement

I'd like to thank the authors for taking the time and discussing the weaknesses in detail. I think the overall positive trend leaning accept is adequate and I will keep my score.