Towards Control-Centric Representations in Reinforcement Learning from Images

Thanks for your comments, we would like to clarify some points and explain your questions.

1. Explanation of frame stack.

We use the commonly used frame stack method, which involves stacking $(o_t, o_{t+1}, o_{t+2})$. This helps capture information about object motion, velocity, and direction, favoring the tasks with an underlying POMDP.

2. Explanation of the relative position embedding.

We did utilize relative position embeddings in our model, drawing inspiration from the prevailing techniques in the domain of transformer architectures. By doing so, we adopt the relative positional embeddings to encode relative temporal positional information into both state and action tokens. We will consider offering a more detailed description in subsequent revisions.

3. Explanation of the Block-wise masking.

The block masking technique we employed is inspired by the method described in [1], which utilizes both spatial and temporal masking blocks. We divide the sequence of image frames into several cubic segments. This segmentation considers both spatial and temporal dimensions, ensuring that each cube encompasses a portion of an image as well as its temporal extension over consecutive frames. The cubes are selected by initially defining a 2-D block at a random time step in the sequence of frames. Once a 2-D block is defined, it is extended along the temporal dimension. For the actual masking process, we randomly select these cubes until we achieve a predetermined masking ratio. Recognizing the importance of making the paper self-contained, we will provide a more detailed description in the future version of our paper.

[1] Wei, C., Fan, H., Xie, S., Wu, C. Y., Yuille, A., & Feichtenhofer, C. (2022). Masked feature prediction for self-supervised visual pre-training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 14668-14678).

4. The key difference between MLR and the proposed method ReBis.

We have provided a detailed analysis in Appendix C, comparing the ReBis with existing methods, particularly MLR. Firstly, it is essential to clarify that while ReBis shares some similarities with MLR, such as the adoption of a mask strategy and a transformer-based encoder, these elements are not unique to MLR and are widely used in current models. The key to differentiate a model is how it defines its objectives and training process when applying these frameworks. In contrast to MLR, ReBis is fundamentally based on the principles of bisimulation. This influences our method to tailor its objectives and training process specifically to enhance bisimulation. As a result, ReBis achieves an effective balance between efficiency and capability, particularly in capturing control-centric (reward-associated) state dynamics. This is a significant departure from MLR, which tends to overlook task-related information.

From the model architecture perspective, MLR employs a more complex network structure, including additional two projection heads, and one prediction head. While these components result in higher complexity and computational demands and are integral to MLR, ReBis does not require such complex structures to effectively perform its tasks.

Therefore, both in terms of design target and model architecture, ReBis significantly diverges from MLR, showcasing its unique approach and contribution to the field.

5. Why not other methods (e.g., Dreamer)?

We did consider and experiment with structures like Dreamer to enhance bisimulation performance. However, we encountered several issues. Firstly, methods in the Dreamer series are model-based, which inherently differs from the approach we aimed for in our research. Secondly, Dreamer involves pixel-level reconstruction. Our goal was to minimize irrelevant information as much as possible. Therefore, the reconstruction nature of Dreamer essentially conflicts with the objectives of bisimulation, which seeks to abstract important information away from such low-level details.

Given these considerations, it is currently not feasible to apply bisimulation to methods like Dreamer. On the other hand, MLR, with its focus on latent space reconstruction, aligns more closely with the principles of bisimulation. Furthermore, it is indeed possible that there are other approaches that could potentially yield equal or better results than ReBis. Our focus has been on establishing a paradigm in which certain modules are interchangeable. We plan to explore the potential of integrating ReBis with other methods in future work.

6. Ablation studies.

Thank you for your comment regarding the need for ablation studies to validate the effectiveness of the algorithm design, particularly the use of masked inputs and momentum encoder. Masking inputs is a fundamental step in our methodology, as it is integral to the subsequent reconstruction process, which plays a crucial role in our model's learning mechanism. We would like to emphasize that to explore its effectiveness, we have conducted thorough ablation experiments focusing on different masking ratios. The detailed results and analyses of these ablation studies are presented in Appendix D.3 of our paper.

The momentum update technique has become a common practice due to its effectiveness in enhancing the stability of the learning process. Our inspiration to use the momentum encoder for representation learning comes from its successful adaptation in references [2, 3, 4]. In response to your suggestion, we have conducted an ablation experiment focusing on the momentum coefficient $m$. The results of this experiment are as follows:

[2] He, K., Fan, H., Wu, Y., Xie, S., & Girshick, R. (2020). Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 9729-9738).

[3] Laskin, M., Srinivas, A., & Abbeel, P. (2020, November). Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning (pp. 5639-5650). PMLR.

[4] Zang, H., Li, X., & Wang, M. (2022, June). Simsr: Simple distance-based state representations for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 36, No. 8, pp. 8997-9005).

7. Comparison with model-based RL method.

Indeed, ReBis is a model-free method, as we does not construct the reward function and the explicit reconstruction in observation space. In our framework, we only reconstruct the state embeddings in latent space, which means categorizing our method into model-based RL seems not to be appropriate. Nevertheless, we still constructed an experiment on comparing our method and a model-based representation method TIA [5] (one method that incorporates Dreamer structure and also aims to learn task-specific information), to show the effectiveness of our model:

Please note that due to the tight rebuttal timeframe, these experimental results are preliminary. We plan to include more model-based methods into comparison in subsequent versions, to further demonstrate the advantages of our model.

[5] Fu, X., Yang, G., Agrawal, P., & Jaakkola, T. (2021, July). Learning task informed abstractions. In International Conference on Machine Learning (pp. 3480-3491). PMLR.

Thanks for your comments, we would like to clarify some points and explain your questions.

1. Do we need to update the encoder every RL update? Can we update it less frequently than the policy?

Our current method does indeed require updating the encoder at every RL update. This frequent updating is integral to ensuring that the encoder accurately captures the evolving state representations, which is crucial for the effectiveness of our RL algorithm. However, inspired by your suggestion, we explored the possibility of updating the representation encoder less frequently.

To test this, we conducted experiments where the encoder was updated every 10, 100, or 1000 environmental interactions. The results of these experiments are as follows:

These findings suggest that while our current approach of updating the encoder at every RL update is effective, there may be potential in exploring alternative update frequencies. The need for frequent updates stems from the dynamic nature of reinforcement learning environments. Updating the encoder with every RL update ensures that the state representations remain current and relevant, thereby facilitating more effective learning and decision-making. Less frequent updates may lead to a decreased ability of the encoder to keep pace with the rapidly changing environment and agent dynamics. These findings open avenues for future research to optimize the balance between update frequency and computational efficiency, potentially through adaptive updating mechanisms or more efficient methods.

2. Is the method complementary to other representation learning approaches?

Our method can be effectively combined with data augmentation techniques, such as random crop. Prior research [1, 2] has uncovered that applying proper data augmentation can significantly improve the sample efficiency of reinforcement learning models.

The integration of video prediction techniques with our method may offer additional contextual information, aiding in a better understanding of temporal dynamics. Similarly, time contrastive learning, which focuses on understanding the temporal relationships between different observations, could also complement our approach. By leveraging the temporal contrasts, our method might further enhance its ability to discern meaningful patterns in temporal sequences.

In our future work, we believe that incorporating these approaches could potentially lead to significant improvements in the model's performance, particularly in scenarios with complex temporal dependencies.

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

[2] Kostrikov, I., Yarats, D., & Fergus, R. (2020). Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649.

3. Can we use it for pre-training a representation without task specification?

Our model is primarily designed to extract task-specific features, which is central to its effectiveness in sparse reward environments. This focus ensures that the model could capture relevant information that directly contributes to the task.

Despite this task-specific orientation, we conducted exploratory experiments to assess the model's adaptability to different environments. In one such experiment, we used an encoder trained in the 'Finger Spin' environment to evaluate its performance across 'Finger Turn Easy’ environment. The results of this experiment are as follows:

These findings suggest that the learned representation will generalize to unseen reward functions, it also exhibits a degree of flexibility and can adapt to new environments to some extent. Based on these results, it appears that our model can be used for pre-training without explicit task specifications. However, the effectiveness of this pre-training may vary depending on the similarity between the training and target environments.

We appreciate the reviewer’s feedback to our replies and the concerns about the generalizability and applicability of our method.

In terms of generalizability for in-domain tasks, our approach holds distinct advantages. These arise from our method's integration of both task-specific and spatiotemporal dependencies, a claim substantiated by our latest experiments mentioned in the Answer to Q3. These findings further demonstrate that the ReBis encoder maintains superior performance over current baselines, even after adopting the encoder from another in-domain task. Furthermore, we recognize that the reviewer's concern may arise from bisimulation theory's inherent focus on integrating behavioral similarity into representation learning, which naturally leads to task or domain specificity. Adhering to the 'no free lunch theorem', we understand that a model tailored for a specific task may not universally adapt to all environments, which is why our evaluations are concentrated solely on in-domain tasks. And models with higher generalizability with cross-domain tasks often necessitates more in-domain-specific tuning. Hence, the trade-off between generalizability and task-specific (or domain-specific) optimality is a direct consequence of the model's design principles.

Additionally, it is important to note that all baseline approaches used for comparison in our study update their encoder with the same frequency as ours. This frequency setting is also consistent with the theoretical requirement of a high update frequency for the encoder, given the dependence of the bisimulation objective on the policy, and it will not hinder the effectiveness of our model in complex environments. Besides, this does not affect inference complexity as we only use CNN encoder during inference time. Nevertheless, the time complexity deserves attention here, as the incorporation of a transformer architecture in our training process results in considerable time requirements for each gradient update. Therefore, future work could interestingly focus on accelerating model forward speed and achieving model distillation without compromising accuracy.

Thanks for your comments, we would like to clarify some points and explain your questions.

1. Can this loss be used as a pretraining loss for learning an encoder?

Thank you for your question. We have experimented with this possibility. We used an encoder pre-trained in one environment (Finger Spin) to evaluate its performance in different environment (Finger Turn Easy). The result of the experiment is as follows:

This experiment demonstrates that our method is versatile enough to be used for pre-training representations without task-specific constraints. The trained encoder shows promising adaptability and generalization across various environments, indicating its potential for broader applications in reinforcement learning.

2. The method of representation learning seems way more complex than the RL algorithm itself. If this is an auxiliary loss, the transformer model capturing the dynamics is never used by the RL algorithm which seems a waste of resources.

While the transformer model primarily focuses on capturing the dynamics for representation learning, its contribution extends beyond being just an auxiliary loss. The learned representations, albeit indirectly, play a vital role in the overall RL process. They provide the RL algorithm with a more nuanced understanding of the environment. Regarding the concern about resource utilization, Rebis aims to balance the computational overhead with the benefits brought by the transformer model. We understand the importance of efficient resource usage and continuously strive to optimize our model to make the best use of available resources. The current design reflects a trade-off where the added complexity is counterbalanced by the substantial improvements in learning quality.

We acknowledge the need to explore ways to directly leverage the dynamics captured by the transformer model within the RL algorithm. Future iterations of our research will focus on more tightly integrating these components, ensuring that the dynamics captured by the transformer are more directly utilized in decision-making processes.

3. What if the distracting images in distracting dm control are added during training as well? Can the model get distracted then?

It is a consideration that aligning the noise levels in training data with those in testing data, could potentially lead to better adaptation and performance in noisy environments. This approach might allow the model to become more accustomed to distractions.

In our current research, we deliberately chose a more complex setup where the model is trained without distractions and tested in the presence of distracting images. This setup was intended to assess the model's ability to generalize and perform effectively in environments different from those encountered during training.

In response to your query, we conducted additional experiments where distracting elements were included during both training and testing. We found that the results obtained from both experimental setups were similar. This outcome underscores the effectiveness of our method, indicating its robustness and generalizability across different levels of environmental noise.

6. Explanation of Block-wise Masking.

We would like to clarify that the use of block-wise masking is not to simulate the sampling of stochastic dynamics as the reviewer assumed. As we have discussed in Section 4, we consider the environment nearly deterministic, we employee the transformer as a deterministic dynamics model. Besides, block-wise masking in our approach does not serve as the sampling of stochastic dynamics, it in fact is used to get rid of spatio-temporal redundance of the observations, coupled with bisimulation-based objective. Intuitively, we don't think random cropping, drop-out and injection of noise can serve the similar purpose, due to their inadequacy of capturing the observation-level spatio-temporal information.

Integrating random cropping with our method may enhance the model's ability to generalize and adapt to different visual inputs. However, employing random cropping as a standalone technique is not as effective in our specific application. While random cropping often results in large, contiguous sections of the image being visible, block-wise masking creates a more scattered pattern of visible and masked regions. This scattered pattern is more effective in our context for capturing control-centric information, as it challenges the model to reconstruct and understand fragmented visual inputs.

For dropout and injecting noise, although they can serve as sampling from stochastic dynamics, they are not applicable currently in our framework as we mainly focus on deterministic settings. We would like to investigate the potential of applying our framework in stochastic setting in our future work.

7. The explanation of how ReBis overcome the challenge of sparse rewards.

ReBis overcomes the challenge of sparse rewards by employing an asymmetrical architecture in its dynamics model. As mentioned in our paper, particularly in Theorem 4, we acknowledge that bisimulation tends to collapse in sparse reward environments. However, it's important to note that agents can still interact with these environments, which means the agent can still learn control relevant informations from the interactions. This understanding forms the primary motivation behind our work with ReBis.

In the framework of ReBis, an asymmetric component within the Siamese architecture is specifically tailored to prevent feature collapse in environments where rewards provide limited information. The effectiveness of this approach is underpinned by Theorem 5, which demonstrates how an asymmetrical architecture can alleviate feature collapse. It does so by increasing the effective feature dimensionality throughout the training process. Specific details regarding Theorem 5 are presented in the appendix.

8. Concern about the novelty of our paper.

While our model utilizes structural components that are commonly recognized in the field, the true claim of our research lies in our unique application of these elements. We've conducted an in-depth examination of bisimulation within the context of reinforcement learning, with a specific focus on its limitation in complex environments with sparse rewards.

Our main contribution is the development of a straightforward yet highly effective approach, designed specifically to address the limitations we identified in traditional bisimulation methods. We have meticulously tailored the objectives and training processes of our model to cater to the unique challenges of reinforcement learning, with a particular focus on learning robust, control-centric representations.

Complementing our theoretical advancements, our model has been empirically validated to demonstrate its effectiveness. Our research reveals that such theoretically-grounded solutions can effectively mitigate challenges that more complex traditional methods struggle to overcome. In essence, our model is an effective framework that can match or outperform existing algorithms, which may provide new insights in the field.

Thank you for your valuable comments. We regret that the initial manuscript did not adequately convey our assumptions, partly due to page constraints. We deeply appreciate your feedback, which has prompted us to revisit these aspects and provide more clarity. In this paper, we assume that the dynamics in real-world environments tend to be nearly deterministic, and therefore we focus primarily on deterministic settings. We have thoroughly revised the related theories and descriptions, marking these updates in blue for your convenience.

1. Explanation of Definition 1.

It is important to note that the term 'bisimulation-based' does not necessarily imply the inclusion of the Wasserstein distance. The bisimulation-based methods, such as MICo or SimSR, also do not incorporate the Wasserstein distance. Instead, they employ independent coupling as an alternative to mitigate the computational cost associated with the Wasserstein distance. Besides, back to the early developments in bisimulation, as illustrated in [1, 2], they originally presented both Wasserstein-based and Total-Variation-based objectives. Therefore, ‘bisimulation-based' approaches are not inherently required to include the Wasserstein distance.

Additionally, we would also like to kindly remind the reviewer that the defined measurement does not strictly qualify as a 'metric', as it does not fulfill the requirements of the triangle inequality. Regarding the absence of the expectation operator in our definition, this omission stems from the assumption of deterministic transition dynamics, as outlined in Sections 3 and 4. This assumption influenced our choice of a transformer as the dynamics model, ensuring deterministic reconstruction and prediction within our framework. We appreciate your feedback on the need for clearer, self-contained explanations and have accordingly revised the relevant sections in our manuscript.

[1] Ferns, N. F. (2003). Metrics for markov decision processes.

[2] Ferns, N., Panangaden, P., & Precup, D. (2012). Metrics for finite Markov decision processes. arXiv preprint arXiv:1207.4114.

2. Explanation of Theorem 2 and 3.

We're grateful for the chance to elucidate our theoretical framework and how it aligns with our definitions. In our paper, we operate under the assumption that the dynamics in real-world environments are nearly deterministic. Therefore, our focus is primarily on deterministic settings. Specifically, we assume that for any latent state $s\in\mathcal{S}$, $a \in \mathcal{A}$, there exists a unique $\kappa(s,a)\in\mathcal{S}$ such that $P_{s}^a(\kappa(s,a))=1$. With this assumption in place, we show that the equivalence between the bisimulation measurement $d(\kappa(s_i,a),\kappa(s_j,a))$ and the W1 distance $\mathcal{W_1}(d)(P{s_i}^a,P_{s_j}^a)$ in deterministic settings.

As such, we can effectively apply Theorems 2 and 3 to the measurement defined in our paper. The deterministic nature of the transitions in our model aligns well with the theoretical underpinnings of these theorems, ensuring their relevance and applicability to our work. We hope this explanation clarifies how our assumptions and theoretical developments are congruent with the application of Theorems 2 and 3 to our measurement.

3. Explanation of Theorem 4.

We apologize for the typo error that we made. The right-hand-side of Equation 16 in our previous manuscript should be $\max_{x,y}|R_x^\pi-R_y^\pi|$ instead of $|r_{s_i}^{\pi} - r_{s_j}^{\pi}|$. We have fixed this error in new manuscript. We also provide a more detailed derivation in the appendix now. Regarding Theorem 4, our intention was to address the challenges posed by sparse reward environments when using bisimulation. MICo distance is a special case, inherently possesses a non-zero distance. However, as we described in the beginning of Section 3, we "primarily focus on a cosine distance-based bisimulation measurement", where the issue that we referred indeed exists. Besides, although the characteristic of MICo having non-zero self-distance makes it unsuitable for Theorem 4, the author of SimSR[3] pointed out that this attribute may encounter the failure mode of representation collapse, even without considering the sparse reward settings.

[3] Zang, H., Li, X., & Wang, M. (2022). Simsr: Simple distance-based state representations for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 36, No. 8, pp. 8997-9005).

4. Explanation of "multi-modals".

"Multi-modals" here refers to the presence of multiple modes in the behavior or decision-making patterns of an agent.

5. Description of the blue color in Tables 1 and 2.

The blue color in these tables is specifically used to highlight tasks that operate under sparse reward conditions. These tasks include Ball in Cup, Catch, Cartpole Swingup Sparse, Finger Turn Easy/Hard, and Pendulum Swingup. In the newly uploaded manuscript, we added the corresponding explanation.

Dear reviewer,

Thanks again for the time spent on our manuscript and your valuable feedback. We have revised the paper to address your questions and comments. We believe we have responded to the questions and concerns in full, but if something is missing please let us know and we would be happy to add it. If we have addressed your concerns we would appreciate it if the reviewer could re-evaluate our work in light of these clarifications.

3. Concern on Theorem 5

We apologize for the overclaim in our last reply of the sentence "ReBis overcomes the challenge of sparse rewards by ...". We meant to say ReBis can alleviate the representation collapse issues in sparse reward environments. In our context, the term 'asymmetric' refers not only to the network architecture, such as in siamese networks, but also encompasses the asymmetry in inputs. This is implemented by applying block-wise masking to one branch with the online encoder, while the other branch with the momentum encoder remains unmasked. We underscore that block-wise masking, as a means to reduce redundancy, is a crucial element of this asymmetric setup. Without block-wise masking, siamese architectures struggle to extract useful information in environments with sparse rewards. This is because they lack the inherent capacity to disentangle environmental information, leading to a collapse into identical representations. Conversely, with asymmetric inputs, the effective feature dimensionality improves over time during training. This allows for some information capture, even in settings where environmental rewards are extremely sparse.

4. Other issues

(1) MICo or SimSR don't call themselves "bisimulation"

We appreciate the reviewer's perspective on the term 'bisimulation' and understand that it may be perceived as having strict limitations, potentially leading to disagreements. Our intention is to highlight that the motivation of our work is from bisimulation-like objectives (or bisimulation-like measurements), rather than the specific stricted "bisimulation metric". We will modify these terms in the camera-ready version to ensure clarity and precision.

(2) $\kappa$ is mentioned in Equation 3 but defined after Equation 5

We first introduced $\kappa$ on the upper side of the second page, which might not be very noticeable. We have now reiterated in Equation 3.

(3) Others

Thanks! We have now provided the corresponding revisions in our updated manuscript.