Paper: MIND: Masked and Inverse Dynamics Modeling for Data-Efficient Deep Reinforcement Learning (ICLR-2024)

Reviewer comment (Writing):
**Clarity**. Section 3 does not explain why a Transformer module is introduced to help modeling.

Author response:
To address the high correlation among consecutive frames in RL, we introduced a transformer. It helps capture spatio-temporal information effectively while considering high correlation.

For inverse dynamics modeling, actions are predicted using a classifier. This classifier utilizes the current state representations from the online network and the next state representations from the IT network. We have revised the manuscript to clearly mentioned in the manuscript, enhancing clarity on how actions are predicted.

Regrading the online network, it simultaneously performs state representation regression for the reconstruction task and action prediction for inverse dynamics modeling. The state representations output from the online and IT networks are used to predict actions through a classifier, while the online and MT networks focus on minimizing state representation errors. We have separate neural networks for each task, ensuring clear task-specific processing.

Reviewer comment (Writing):
**Clarity**. Section 3 does not explain how the online network performs state representation regression and action prediction simultaneously. Are there separate heads for each auxiliary task?

Author response:
To address the high correlation among consecutive frames in RL, we introduced a transformer. It helps capture spatio-temporal information effectively while considering high correlation.

For inverse dynamics modeling, actions are predicted using a classifier. This classifier utilizes the current state representations from the online network and the next state representations from the IT network. We have revised the manuscript to clearly mentioned in the manuscript, enhancing clarity on how actions are predicted.

Regrading the online network, it simultaneously performs state representation regression for the reconstruction task and action prediction for inverse dynamics modeling. The state representations output from the online and IT networks are used to predict actions through a classifier, while the online and MT networks focus on minimizing state representation errors. We have separate neural networks for each task, ensuring clear task-specific processing.

Reviewer comment (Experiments):
**Baselines**. Clearly there are missing baselines like RAD [1], DrQv2 [2] according to this study [3] on self-supervised learning with RL, and possibly [4] which uses Transformer for next frame prediction, and vanilla MAE [5]. With the current baselines I can not evaluate the significance of proposed method, considering the novelty of the proposed method is limited.

Author response:
In response to the reviewer's suggestion, we have included comparison in our revised Table 2. We conducted experiments across six DMControl environments with a limit of 100k interactions. Our method demonstrated notable performance, winning in three out of six environments and outperforming both the mean and median across all environments. Furthermore, we experimented on five tasks of medium-difficulty DMControl with 1M interactions, choosing comparison methods in line with the model-free approach. The comparison methods reported their average scores for six different seeds, while our results are based on two seeds. The results show improved data efficiency of our method at medium-difficulty compared to the baselines. We will conducted additional experiments to ensure a fair comparison and report the results in the final version. Detailed results are described in Appendix C.2.

Reviewer comment (Experiments):
The author is recommended to provide a comparison at 1M/2M of harder tasks instead of the six tasks presented in the paper.

Author response:
A1: Thank you for the valuable suggestions. We agree with the reviewer that further experiments with continuous control benchmarks is necessary. TACO believes that there is no obligation to compare because it is considered a recent work according to the policy of the conference (https://iclr.cc/Conferences/2024/ReviewerGuide). Nevertheless, to address the reviewer's concerns, we tested our method on five medium-difficulty DMControl tasks using 1M interactions, as reported by TACO. We selected these tasks to align with our model-free approach and make a comparable analysis with comparison methods. The comparison methods reported their average scores for six different seeds, while our results are based on two seeds. As shown in the table below, the results show improved data efficiency of our method at medium-difficulty compared to the benchmarks. In particular, in the environments Walker, run; Hopper, hop and Reacher, hard, our method outperformed TACO, which was the best performing method among the compared algorithms. We will conduct additional experiments to ensure a fair comparison and report the results in the final version. We have included the detailed DMControl-1M results excluding TACO in Appendix C.2 for further reference.

| Environment | CURL | DrQ | DrQ-v2 | SPR | ATC | A-LIX | TACO | MIND |
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
| Quadruped, run | 181 $\pm$ 14 | 179 $\pm$ 18 | 407 $\pm$ 21 | 448 $\pm$ 79 | 432 $\pm$ 54 | 454 $\pm$ 42 | **541 $\pm$ 38** | 520 $\pm$ 30 |
| Walker, run | 387 $\pm$ 24 | 451 $\pm$ 73 | 517 $\pm$ 43 | 560 $\pm$ 71 | 502 $\pm$ 171 | 617 $\pm$ 12 | 637 $\pm$ 21 | **659 $\pm$ 18** |
| Hopper, hop | 152 $\pm$ 34 | 192 $\pm$ 41 | 192 $\pm$ 41 | 154 $\pm$ 10 | 112 $\pm$ 98 | 225 $\pm$ 13 | 261 $\pm$ 52 | **293 $\pm$ 3** |
| Reacher, hard | 400 $\pm$ 29 | 471 $\pm$ 45 | 572 $\pm$ 51 | 711 $\pm$ 92 | 863 $\pm$ 12 | 510 $\pm$ 16 | 883 $\pm$ 63 | **954 $\pm$ 28** |
| Acrobot, swingup | 5 $\pm$ 1 | 24 $\pm$ 8 | 210 $\pm$ 12 | 198 $\pm$ 21 | 206 $\pm$ 61 | 112 $\pm$ 23 | **241 $\pm$ 21** | 217 $\pm$ 13 |
| Mean | 225.0 | 263.4 | 379.6 | 414.2 | 423.0 | 383.6 | 512.6 | **528.8** |
| Median | 181.0 | 192.0 | 407.0 | 448.0 | 432.0 | 454.0 | **541.0** | 520.0 |

Reviewer comment (Experiments):
Why not also use masking of the inputs to the IT network?

Author response:
Our method uses masking augmentation exclusively in the online network for the reconstruction task. This is to focus on missing information without the need for separate data augmentation for learning visual representations. On the other hand, the IT network is tailored for inverse dynamics modeling. It extracts representations for a series of next states, capturing essential information about rapidly changing states. While our method conducts both reconstruction and inverse dynamics modeling simultaneously, the objectives differ. Therefore, we input the original or unmasked data into the IT network to align with its specific purpose.

Reviewer comment (Writing):
In equation one, p, c, q, is not defined

Author response:
A1: Thank you for your careful reading. In equation one, the CNN, transformer, projector, and predictor are denoted by *c*, *t*, *p*, and *q*, respectively. We have provided description of these terms in Section 3.

Reviewer comment (Evaluation):
It would be good to report confidence intervals in results, for instance in Table 1

Author response:
A2: Thank you for your suggestion. We have included the standard deviations of the human-normalized scores for each game in Table 1. In addition, Figures 2 and 3, as well as Table 2, provide information about the confidence interval or standard deviation.

Reviewer comment (Novelty):
The paper doesn't seem very novel -- it uses momentum contrast masked loss + a standard inverse dynamics representation learning objective

Author response:
A3: Mask-based reconstruction and dynamics modeling are established techniques in representation learning. However, their application in RL, especially in rapidly changing environments, is still largely unexplored. Our work is novel in that we are among the first to combine dynamics modeling with masked modeling in a self-supervised multi-task learning framework. This approach specifically addresses the limitations noted in previous methods (referenced in [1]). Our focus extends beyond just learning visual representations. We also aim to capture the controllable elements in environments with rapid state changes. To handle the high correlation between consecutive frames in RL, we introduced a transformer. This aids in predicting a series of actions by analyzing correlations and spatio-temporal information of consecutive states. Furthermore, we conducted systematic experiments to establish best practices in this area. Our findings show that integrating inverse dynamics modeling in latent space, like MLR [1], enhances efficiency in rapidly evolving environments. This advancement in methodology represents a significant contribution to the field.

[1] Yu, T., Zhang, Z., Lan, C., Lu, Y., & Chen, Z. (2022). Mask-based latent reconstruction for reinforcement learning. Advances in Neural Information Processing Systems, 35, 25117-25131.

Reviewer comment (Evaluation):
Can the authors provide some intuition why masking + momentum contrast outperforms other baseline methods?

Author response:
A4: Data augmentation techniques work very sensitively in RL. In previous studies [1][2], they empirically experimented with various data augmentations and reported which augmentation can deliver the correct reward signal to the agent. On the other hand, masked modeling, which only considers masking augmentation, is a simple yet effective way to reconstruct missing data. Typical reconstruction tasks make predictions at the pixel-level. However, in RL that leverages latent representations to learn the agent's policy, it is important to reconstruct missing state information in the latent space. In terms of learning representations for states, momentum contrast (i.e., self-distillation approach [5]) has been shown to excellent at preventing representation collapse, benefiting from its asymmetric architecture and exponential moving average updating [3][4][5]. SPR [3], demonstrated that without a target network, performance declines due to representation collapse. MLR [6] combines the strengths of both masking and momentum contrast to reconstruct latent representations, proving to be more effective than previous studies.

[1] Yarats, D., Kostrikov, I., & Fergus, R. (2020, October). Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International conference on learning representations.

[2] Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., & Srinivas, A. (2020). Reinforcement learning with augmented data. Advances in neural information processing systems, 33, 19884-19895.

[3] Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., & Bachman, P. (2020, October). Data-Efficient Reinforcement Learning with Self-Predictive Representations. In International Conference on Learning Representations.

[4] Grill, J. B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., ... & Valko, M. (2020). Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33, 21271-21284.

[5] Balestriero, R., Ibrahim, M., Sobal, V., Morcos, A., Shekhar, S., Goldstein, T., ... & Goldblum, M. (2023). A cookbook of self-supervised learning. arXiv preprint arXiv:2304.12210.

[6] Yu, T., Zhang, Z., Lan, C., Lu, Y., & Chen, Z. (2022). Mask-based latent reconstruction for reinforcement learning. Advances in Neural Information Processing Systems, 35, 25117-25131.