CLEAR: An Information-Theoretic  Framework for Distraction-Free Representation Learning in Visual Offline RL

Authors: Tri Wahyu Guntara, Daiki E. Matsunaga, HyeongJoo Hwang, Kee-Eung Kim
Primary area: reinforcement learning
Keywords: Visual Offline Reinforcement Learning, Information-Theoretic Representation Learning

Abstract:
Visual offline RL aims to learn an optimal policy for visual domains, solely from the pre-collected dataset comprised of actions taken on visual observations. Prior works on visual RL typically learn a dynamics model by extracting a latent state representation. However, the learned representation would contain factors irrelevant to control when there are distractions in the visual observations. These nuisance factors introduced by the distraction further exacerbates the difficulties of learning a good policy in the offline RL setting. In this work, we formalize the visual offline RL setting as a Partially Observable Markov Decision Process with exogenous variables (ExoPOMDP) and identify  these problems with previous approaches under an information-theoretic lens. To overcome these challenges, we propose CLEAR (**C**ontrollable **L**atent State **E**xtr**A**cto**R**) for visual offline RL, which learns the dynamics model of a succinct agent-centric state representation that is consistent with the underlying ExoPOMDP. We empirically demonstrate that CLEAR is able to outperform baselines on the DeepMind Control Suite with various types of distractions and perform consistently well across these distractions. We further provide qualitative analysis on the results showing that our approach successfully disentangles the distraction factors from the agent-centric state representation.

Decision: reject
Reviewer scores: 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1729759553656

thread_type:
official_review

<thread_object>
creator: reviewer_udww

<content>
summary:
This paper proposes CLEAR, an agent-centric state-representation method for vision-based offline reinforcement learning. The authors present an information theoretical approach to achieve this goal, disentangling exogenous variables while enforcing the Markovian properties. For computation, the author proposed a lower bound of compounded mutual information in Eq. (3) to  Eq. (4) and designed a corresponding encoder-decoder architecture that can be trained with image representation techniques. The authors demonstrated that learned state representations by the proposed CLEAR method show good performance when there is uncontrollable visual distraction in offline RL settings.

strengths:
1. The paper is well-written. It was straightforward to understand the core idea; despite dealing with complex concepts. Figures and equations are appropriately placed. 
2. The proposed method is sound and properly designed to solve the offline RL problems. 
3. In the visual domain, circumstances that ExoPOMDP describes are very common. This work can directly contribute to solving RL in real-world problems. 
4. The authors have clearly put significant effort into reducing computation costs and simplifying architecture design. The work focus on enhancing efficiency and practicality of information theoretical approaches. I think most of the results could be implemented and reproduced even though source code is not provided.

weaknesses:
1. Although the authors' claim primarily in the visual offline RL, the conducted experiments are strictly limited to MuJoCo variants. The authors are strongly encouraged to strengthen the experiment section by including (pure) vision-based RL problems with distractions.
2. Extensibility. I think the adaptation ability of CLEAR representation for different visual distractions from the training environments is only partially validated by multiple videos. The authors are encouraged to analyze general extensibility of CLEAR in these aspects for real-world scenarios.
3. It would be better if intuitive motivation or a toy experiment empirically verifies the ExoPOMDP setting.

questions:
1. In the graphical model in Fig 2, what happens when e_t is affected by s_{t-1}? Does this modification change Eqs. (2) and (3)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730407823139

thread_type:
official_review

<thread_object>
creator: reviewer_ucli

<content>
summary:
The paper addresses the problem of visual offline reinforcement learning in settings with distractions in visual observations. The authors formalize this setting as an ExoPOMDP and propose CLEAR, a method that learns the dynamics model of a agent-centric state representation consistent with the ExoPOMDP framework. Specifically, they take an information-theoretic approach, introducing mutual information-based loss functions to disentangle distraction factors from the agent-centric representation. Their method is evaluated on three tasks in the DeepMind Control Suite with diverse types of distractions.

strengths:
1/ The high-level motivation of the method and the information-theoretic foundation look solid. Specifically, the regularization term that encourages s_hat to encode an agent-centric state controllable by actions is particularly interesting, and its effectiveness is well supported by the experiments.

2/ They evaluate the methods on various types of distractions, with 2x2 grid distractions being especially notable. This setup serves as a great testbed for assessing whether a method can distinguish controllable parts from non-controllable ones. The experiments show that their method consistently outperforms other methods in this setting, indicating that the baselines lack the ability to discern controllability, while their method succeeds.

3/ Their experiments include a comprehensive set of prior work as baselines, covering not only offline RL but also model-based RL. This strengthens the paper by demonstrating that their method outperforms approaches across various frameworks.

weaknesses:
1/ Three tasks from a single domain seem too few, particularly since some baselines are comparable to the proposed method in some tasks (e.g. InfoGating on Hopper Hop). Expanding the evaluation to include more tasks (e.g., 5-6 tasks) would significantly strengthen the paper. A larger set of tasks would provide a more robust assessment of the method's performance and demonstrate its versatility across different scenarios.

2/ Although a large set of prior works is presented as baselines, some important and closely related studies are missing. Notable examples include [1] and [2]. [1] also employs information theory to address controllability. A detailed comparison between [1] and CLEAR is necessary, both methodologically and experimentally. [2] leverages causality to learn disentangled representations across four different categories, focusing on controllability and reward relevance.

[1] Homanga Bharadhwaj, et al. "Information prioritization through empowerment in visual model-based RL." ICLR 2022

[2] Yuren Liu, et al. "Learning world models with identifiable factorization." Neurips 2023

questions:
1/ Which representation is used for offline RL (i.e., policy learning)? I assume s_hat is used and not e_hat, but I would like to clarify this.

2/ The method requires extensive hyperparameter search over the variance and balancing factors of the KL divergence for each task, domain, and distraction type (Table 9). It would be much more convenient to eliminate this necessity. Do you have any ideas on how to achieve this?

3/ Since the method aims to improve representation learning, it could potentially be applied within an MBRL framework. How do you think it would perform when integrated with an MBRL framework?

4/ How can this method handle other forms of nuisance information, such as color perturbations of task-relevant parts, where the assumption in L278 does not hold?

5/ Including more detailed descriptions of each column in the caption of Table 3 would make it much easier to interpret the results.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730674609460

thread_type:
official_review

<thread_object>
creator: reviewer_7k7j

<content>
summary:
The paper aims at disentangling distractions in the visual observations from agent-related information for visual offline RL tasks. To this end, the paper formalizes the visual offline RL setting as an ExoPOMDP and derive the corresponding objective with the tools of information theory. Specifically, the disentanglement is achieved by a VAE-like model, where the decoder network has a compositional architecture that blends the reconstructed foreground and background visual parts with masks.

strengths:
- The idea is simple and effective.
- The experiments seem thorough on the DeepMind Control Suite.
- The paper is clearly written and well-organized.

weaknesses:
- I am admittedly not an expert of visual offline RL, and would like to hear my colleague reviewers' opinions on this.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1732282703102

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
# General Response

We would like to thank all reviewers for taking their time to provide constructive feedback on our paper. We are very encouraged that the reviewers have identified the following strengths in our work:

* Simple and effective idea with good motivation (Reviewer $\text{\textcolor{red}{7K7j}}$, $\text{\textcolor{blue}{UCLi}}$, $\text{\textcolor{green}{UDWw}}$)
* Extensive experiments on the DeepMind Control Suite (Reviewer $\text{\textcolor{red}{7K7j}}$, $\text{\textcolor{blue}{UCLi}}$)
* Comprehensive set of baselines and types of distractions (Reviewer $\text{\textcolor{blue}{UCLi}}$)
* Clearly written (Reviewer $\text{\textcolor{red}{7K7j}}$, $\text{\textcolor{green}{UDWw}}$)


We outline below the changes made to the manuscript:

1. Added InfoPower [1] as an additional baseline in the Related Work in Section 4 as well as the Experiments Section (Table 1 and 2 in Section 5) [$\text{\textcolor{blue}{UCLi}}$]

2. Added more detailed descriptions in the caption of Table 3 in Section 5.3 [$\text{\textcolor{blue}{UCLi}}$]

3. Added two additional environments: Finger-Spin and Cartpole-Swingup (Table 7 in Appendix F) [$\text{\textcolor{blue}{UCLi}}$, $\text{\textcolor{green}{UDWw}}$]

4. Added qualitative results for Finger-Spin (Figure 7 in Appendix F) [$\text{\textcolor{blue}{UCLi}}$, $\text{\textcolor{green}{UDWw}}$]

Finally, we have addressed each question as a separate response for each reviewer. Please let us know if there are any further questions or necessary clarifications and we would be happy to address them during the discussion period.

**References**

[1] "Information prioritization through empowerment in visual model-based RL." Bharadhwaj et al. (ICLR 2022)
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper looks at visual offline reinforcement learning under distractions in the visual observations. An ExoPOMDP (POMDP with exogenous variables) formulation is developed, leading to an information theoretic approach that disentangles distracting factors (that are not relevant for control) from agent-centric representations (that are). This leads to a dual-encoder architecture with a compositional decoder, and results are shown (initially) across three environments of DeepMind Control Suite, and subsequently add additional ones post-review. 

The reviewers appreciated the problem setting, information-theoretic formulation, and subsequent approach that was developed. However, there are several concerns (some common) including limited environments (UCLi, UDWw), generality of the method and understanding exactly what kinds of visual perturbations can be handled (UCLi, UDWw), missing related work and benchmarks (UCLi), and hyper-parameters and their sensitivity/tuning. The authors provided a rebuttal, including some additional environments, qualitative results, and additional baselines. Reviewr UCLi mentioned concerns about these results, and during discussion UDWw also continued to have reservations. 

  Overall, this paper is borderline and required considering all of the materials including the paper itself, reviews, rebuttals, author comments, and discussions. After considering these, I lean towards not accepting this paper at this time. While the paper's information-theoretic formulation and approach is interesting, the biggest overall limitation is the experimental execution. I agree with reviewers that the first three environments were indeed not sufficient to elucidate in a rigorous way when and how the method works compared to state-of-art baselines. The addition of new environments and baseline is appreciated, but there are significant concerns on the failure of InfoPower (original environments) and InfoGating (new environments). Unfortunately the rebuttal period doesn't lend itself to thorough examination of the correctness and cause of these results. Further, questions related to exactly what types of distractions the method supports should be better made explicit and analyzed throughout the paper. As a result, while we appreciate the author's significant work and rebuttal, the paper requires a more comprehensive empirical methodology and analysis to scientifically understand where and when this method works compared to state-of-art.