Paper: Information-Theoretic World Model learning for Denoised Predictions (ICLR-2024)

Reviewer comment (Evaluation):
The reconstructions aren't terribly compelling. While background info is removed, it's not obvious if all the information about the agent is correctly retained.

Author response:
Acknowledging the observations about the reconstructions, it is important to highlight from our results that the method retains the majority of the information about the agent. The process effectively darkens the background, predominantly isolating and preserving the agent's details. This indicates that while background information is removed, the critical aspects of the agent are largely maintained. Future iterations of our model will focus on refining this aspect to ensure even more precise retention of all critical information about the agent. For more analysis, we have also added the reconstruction of the Cartpole Swingup task in Figure 5 of the Supplementary Material.

Reviewer comment (Novelty):
Papers like ACRO (Islam 2022) and InfoGating (Tomar 2023, which I suppose is a bit recent) should be cited. InfoPower also feels relevant, and while it is cited, the difference doesn't seem to be sufficiently discussed.

Author response:
**[A2]** Thanks for bringing these papers to our attention. We have cited both the papers.

Reviewer comment (Experiments):
The integration with SAC feels overly specific.

Author response:
**[A3]** Thank you for pointing out the relevance of InfoPower in our discussion. While InfoPower focuses on maximizing empowerment, specifically $I(a_k;z_{k+1}|z_k)$, our objective diverges significantly from this. We deconstruct $I(z_{k+1},a_k;z_k)$ into two components: $I(z_k;z_{k+1})$ and $I(z_k,a_k|z_{k+1})$. Although the second term might appear similar to InfoPower's approach, it fundamentally differs in both concept and implementation. It's true that InfoPower and other methods we cited [1-3] employ an inverse dynamic model $q(a_k|z_k,z_{k+1})$, but our approach advocates a more generalized strategy. We propose that focusing on $I(z_{k+1},a_k;z_k)$ implicitly encompasses learning inverse dynamics, diverging from InfoPower's explicit modeling of this aspect. We have cited InfoPOWER in Section 4.3 now.

Reviewer comment (Evaluation):
The experimental results seem fairly weak. There is a clear improvement in cheetah run and a small improvement in walker walk, but other tasks are more modest.

Author response:
**[A4]** The detailed description of our SAC integration is crucial for both reproducibility and clear understanding, enabling other researchers to accurately replicate and build upon our work.

Reviewer comment (Experiments):
The results in the appendix on standard setup seem quite bad as the proposed method is badly losing to Dreamer.

Author response:
**[A5]** In our rebuttal, we present enhanced results in additional environments such as Walker Walk and Cartpole Swingup, achieved through weighing coefficient tuning conducted subsequent to the paper's initial submission. Our updated results now show a significant improvement over the baseline methods in the specified environments. We encourage reviewing the new graphs presented in Figure 4 for a detailed comparison of these enhancements. The details of these newly optimized weighing coefficient are available in Section 4 of the Supplementary Material. We have also added Cartpole swingup in the random settings (Figure 2 for results and Figure 5 for reconstruction in the Supplementary Material), which clearly aligns with our claims. Additionally, we deliberately highlight the limitations of all methods, including our own, in three challenging environments. This was done to underscore the current limitations of state-of-the-art methods in handling these complex environments, thereby emphasizing the need for further research in this area.

Reviewer comment (Evaluation):
The paper doesn't do much to analyze what is happening in the information bottleneck.

Author response:
**[A6]** DPI, although not outperforming Dreamer in static environments, is designed for effective generalization across natural settings. As illustrated in the graph, other methods in the same category as DPI also face challenges in these environments. Nevertheless, DPI maintains competitive performance compared to other methods (excluding Dreamer) in four out of the six tested environments.

Reviewer comment (Theory):
The paper refers to the setup as a lagrangian formulation, but for this to work, it seems insufficient to treat beta as a hyperparameter. So what exactly makes it a lagrangian setup, as opposed to just adding an extra penalty?

Author response:
Thanks for this question. We use the similar terminology utilises in Information Bottleneck Principle (IB) [4] to be uniform with the theory. In IB formulation, this parameter is responsible for the tradeoff between the complexity of representation and  amount of preserved relevant information. Similarly, in our model, these parameters are instrumental in balancing the relationship between past and future variables, with the present variable acting as the bottleneck.

Reviewer comment (Theory):
What is the reason to think that increasing mutual information between present and future states is compatible with removing distractors? For example, the background video will have temporal correlation, so we could potentially maximize that by keeping the TV in the state. See Efroni 2022 (PPE) as well as Lamb 2022 (AC-State) for some discussion of this issue.

Author response:
**[A2]** This is a great question and thanks for suggesting AC-State [5] especially. Our idea can also be viewed in this way: Forward models while adept at predicting future states, can falter in noisy environments as it can also model task-irrelevant details. Inverse dynamic models, on the other hand, are effective in determining actions based on state transitions but can suffer from a similar problem. Our approach, encompasses both the models, learning predictive models ($q(z_{k+1}|z_k,a_k,h_h)$ in Section 4.2) for the future states and learning actions from the inverse dynamic models implicitly with lower bound on $I(z_{k+1},a_k;z_k)$, which acts as a regularizer.

Reviewer comment (Theory):
Much of the story in the paper is a bit sloppy and even problematic in places: the introduction seems to claim without much support that RL only works well when the observations do not contain noise variables that aren't related to the task at hand. There is a skimming over the real problem, particularly in POMDP settings, where reward-relevant signal might be sparse in observations, which can hurt the sample complexity of algorithms in many problems. The problem is multifaceted, as the SNR (with respect to task-relevant factors, e.g., whether the task variables themselves are noised), reward sparsity, and the degree of partial observability in the environment all factor into the success of various algorithms. It is unclear why, given the introduction's framing, this has anything to do specifically with task relevance in representations: learning all of the transition function, even the noisy part, should still be decodable to a policy that can maximize returns. The real concern is that such representations might not generalize well (e.g., when the noise correlates with a subset of task-relevant variables).

Author response:
**[A1]** It is crucial to clarify the narrative presented in the introduction. The assertion that reinforcement learning (RL) faces challenges in noisy environments is not intended to imply that RL is only effective in the absence of noise. Rather, the point is to highlight the increased complexity and potential for sub-optimal performance when RL algorithms encounter observations cluttered with irrelevant information. Regarding POMDPs and reward sparsity, the reviewer's point is well-taken. Our focus on noise and irrelevant information in the introduction is meant to complement, not overlook, these challenges. The introduction's mention of task relevance in representations is to underline the importance of extracting meaningful information from noisy data, not to suggest that the other variables are insignificant fot the success of a task. By emphasizing task relevance in representations, we underline the importance of extracting meaningful information from noisy data, crucial for ensuring generalizable and robust model performance. The introduction thus sets the stage for presenting Denoised Predictive Imagination (DPI) as a method, designed to tackle these challenges in RL through an information-theoretic approach.

Reviewer comment (Experiments):
There is a missing baseline that is considered important and should be tested against: Self-Predictive Representations (SPR).

Author response:
**[A4]** Thank you for pointing out the omission of Self-Predictive Representations (SPR) [1] as a baseline in our study. Following your suggestion, we have now included SPR as a baseline in our experiments. In the Figure 4, we show that DPI significantly outerpforms SPR in Cheetah run, Walker walk and Cartpole swingup. In Figure 2 of the Supplementary Material, we have also shown that DPI outperforms SPR significantly in the cartpole swingup task in random settings. Additionally, we have also added Dreamer-V2 [2] and VSG [3] as a baseline, which significantly contributed to improving the quality of our research.

Reviewer comment (Novelty):
Beyond missing SPR, the paper does not clearly identify which of the loss components in the proposed model are 1) present in existing baselines or 2) not present in existing baselines. While the paper claims its contribution lies in MI-based objectives, similar approaches have been attempted before, albeit not in exactly the same form. More effort is needed to ground the contributions of the loss components in prior work and to better motivate them, as mentioned above. Additionally, it is unclear why having an MI-based objective in the proposed form(s) is necessary, as other approaches to approximate MI maximization/minimization (e.g., regression losses as in SPR) may work better.

Author response:
**[A5]** Thank you for addressing this. We have consistently mentioned many different baselines in each section representing every component. For example, in Section 4.3, while addressing $I(z_{k+1},a_k;z_k)$, we have mentioned that our objective "is closely related to the predictive coding objective [4,5]". On the top of that, while splitting the term, we have mentioned that the inverse dynamic model resembles to [6-8]. Again, in the same section, while working on $I(z^+,o^+)$, we said that we optimise the lower bound on our MI objective similar to [9]. We respectfully disagree from the reviewer regarding not mentioning the prior work in our loss formulations. The primary focus of our paper is to introduce a novel and generalised framework based on Mutual Information, designed to minimize noise and extract task-relevant information from observations. It is indeed true that there are multiple ways to implement MI maximizations or minimizations, such as using regression losses as seen in Self-Predictive Representations (SPR). Our current approach is designed due to its tight bound and ease in implementation, but we are open to exploring other methodologies in future research.

Reviewer comment (Experiments):
The results are underwhelming. While performance is strong in some experiments (e.g., cheetah run), in most cases it is not as clear. This raises the concern that some degree of hyperparameter tuning may have been performed on half cheetah that was not applied to other environments or baselines.

Author response:
**[A6]** We appreciate your observation regarding the performance of DPI in natural background settings. In our rebuttal, we present enhanced results in additional environments such as Walker Walk and Cartpole Swingup, achieved through setting correct weighing coefficients, conducted subsequent to the paper's initial submission. The details of these newly optimized weights are available in Section 4 of the Supplementary Material. With these optimization, DPI now shows significantly better performance in environments such as Cheetah Run, Walker Walk, and Cartpole Swingup, aligning more closely with our initial claims (Figure 4). We also show a dominant performance of DPI w.r.t. SPR and Dreamer-V2 in cartpole swingup task in random environment (Figure 2 Supplementary Material). On the top of that, we also showed the reconstruction results in the Figure 5 of the Supplementary Material.

Reviewer comment (Experiments):
The design of experiments can be improved. Switching background is not convincing enough to demonstrated the effectiveness of the proposed method.

Author response:
**[A1]** Thank you for your feedback. We appreciate your perspective on the experimental design. It's important to emphasize that changing backgrounds presents one of the most challenging problems in this field, and is an indicator of generalization and robustness. Our experiments, designed around this challenge, are intended to specifically address and demonstrate the effectiveness of our proposed method in tackling this complex issue. We also highlight how other models, despite their strong performance in static backgrounds, fail to maintain their effectiveness in these more challenging dynamic background settings. This contrast not only underscores the unique robustness of our proposed method but also reinforces the complexity and significance of the problem we are addressing.

Reviewer comment (Experiments):
Other basline model like DreamerV2 and DreamerV3 can also be included for comparison.

Author response:
**[A2]** Thank you for suggesting these methods. We have included DreamerV2 in our comparative analysis, but due to time and resource constraints, we were unable to incorporate DreamerV3. This decision was made to ensure a thorough and focused evaluation within our available means. Additionally, we have incorporated other relevant methods as suggested by other reviewers, SPR [1] and VSG [2], broadening the scope of our analysis. Notably, our method outperforms all these included models, showcasing its efficacy and robustness in the evaluated scenarios.

**Additional Comments:**
- Following suggestions from other reviewers, we have incorporated these additional baseline methods, SPR [1] and VSG [2], into our analysis. 
- Our updated results now show a significant improvement over the baseline methods in the specified environments. We encourage reviewing the new graphs presented in Figure 4 for a detailed comparison of these enhancements.
- Additional experiments have been conducted on the Cartpole Swingup task within random environments. These new results clearly demonstrate DPI's capability to filter out background noise and effectively extract relevant information (Results in Supplementary Material Figure 2 and Reconstruction in Supplementary Material Figure 5).

**References**\
[1] Schwarzer, Max et al. “Data-Efficient Reinforcement Learning with Self-Predictive Representations.” International Conference on Learning Representations (2020).\
[2] Jain, Arnav Kumar, et al. "Learning robust dynamics through variational sparse gating." Advances in Neural Information Processing Systems 35 (2022): 1612-1626.

Reviewer comment (Evaluation):
The major weakness is in the experiment setups, results, and discussion around them. In the paper, the authors refer to the results of the first experiment with the static backgrounds as "competitive with all other methods". This is perhaps a bit generous, when DPI has below median performance in four of the six settings. For the natural background settings, the claim is that "DPI achieves better rewards in nearly all environments", which may be true in terms of the blue line spinning above the others, but the performance of DPI only outperforms the baselines in two of the six tasks (Cheetah Run and Walker Walk). Finally, for the random background setting, results are only given for DPI's best performing environment (Cheetah Run). Cheetah Run is also the only environment in the visualized examples, which further supports that it is DPI's best performing environment. Beyond the text perhaps overstating the performance of DPI, there's also no discussion of why the performance of DPI should improve so much in this one environment compared to the others or why it's appropriate to only showcase visualizations and random background performance in it. These issues currently lead to me, as a reader, wondering if the complexity of DPI is justified by the relatively small number of environments it improves in, compared to baselines. Clarification on this point would be greatly appreciated.

Author response:
**[W1]  In the paper, the authors refer to the results of the first experiment with the static backgrounds as "competitive with all other methods". This is perhaps a bit generous, when DPI as below median performance in four of the six settings.**\
**[A1]** Thank you for highlighting this aspect. We have revised the statement to reflect a more nuanced understanding: "DPI exhibits a degree of effectiveness in certain scenarios involving static backgrounds, although it does not consistently outperform all other methods."

Reviewer comment (Presentation):
The minor weakness is that the technical approach section is very dense, making it difficult to get a clear sense of the high level structure of the approach. Figure 2 exists, but is a model of the problem without a strong connection to the implemented solution. Figure 3 also exists but only covers two of the models in the approach. A complete pipeline diagram or text overview would be greatly appreciated.

Author response:
**[W2] For the natural background settings, the claim is that "DPI achieves better rewards in nearly all environments", which may be true in terms of the blue line spinning above the others, but the performance of DPI only outperforms the baselines in two of the six tasks (Cheetah Run and Walker Walk).**\
**[A2]** We appreciate your observation regarding the performance of DPI in natural background settings. In our rebuttal, we present enhanced results in additional environments such as Walker Walk and Cartpole Swingup, achieved through hyperparameter tuning conducted subsequent to the paper's initial submission. The details of these newly optimized hyperparameters are available in Section 4 of the Supplementary Material. With these optimizations, DPI now shows significantly better performance in environments such as Cheetah Run, Walker Walk, and Cartpole Swingup, aligning more closely with our initial claims (Figure 4).

Reviewer comment (Evaluation):
Why does DPI perform so much better in Cheetah Run?

Author response:
**[Q1] Why does DPI perform so much better in Cheetah Run?**\
**[A1]** We have conducted hyperparameter tuning across all environments. Please refer to the updated results for further information (Figure 4, Figure 2 in Supplementary Material).

Reviewer comment (Evaluation):
Why does DPI's performance in the static environments perform below the median in so many cases?

Author response:
**[Q2] Why does DPI's performance in the static environments perform below the median in so many cases?**\
**[A2]** DPI, although not outperforming Dreamer in static environments, is designed for effective generalization across natural settings. As illustrated in the graph, other methods in the same category as DPI also face challenges in these environments. Nevertheless, DPI maintains competitive performance compared to other methods (excluding Dreamer) in four out of the six tested environments.