Summary
The paper proposes an auxiliary task based on minimising mutual information in the learned latent representation. While several approaches have adopted similar loss functions for learning representations, the difference here is that the mutual information is conditioned on the history of past (latent) states and actions. This has the effect of disentangling correlated features, which improves performance when the agent interacts with a new task where those correlations no longer exist. Experiments on continuous control tasks illustrate that the method better generalises to tasks where these correlations are not present, whereas other approaches (also based on mutual information) fail to zero-shot transfer.
Strengths
The paper is well-written and easy to follow, while the experimental results clearly show the advantage of the method over other auxiliary tasks that have previously been proposed.
Although I am not overly familiar with the literature in the space, I believe that the proposed method, while a variation of existing approaches, is novel.
The auxiliary task is well motivated, and the implementation of it within a deep RL framework is elegant, which could likely make it easy for others to incorporate the idea into their own learning algorithms.
Weaknesses
Too much of the experiment section is in the appendix. As it stands, the main paper is not readable without reference to the appendix. I realise that space is an issue, but it would be best if the main text served as a standalone document. For example, there is no explanation of what reversed vs no correlation means (although I can guess). In the main text, only the pendulum task is described - the rest are not even briefly mentioned. To free up space, I recommend moving the discussion of history length and coefficient into the appendix instead. These can be briefly mentioned, but they serve as additional ablations and are not key to the story.
Given Figure 13, the saliency maps in the main text seem a bit cherry-picked, since many of the other ones are not as clean. While the RL performance clearly shows a win, it would be nice to see more evidence that the auxiliary task promotes the kind of disentanglement that is expected. Would it not be possible to investigate this in the non-RL setting? For example, collect data and then simply apply the representation learning part alone to that data. The resulting representations could also then be evaluated using existing metrics from the representation learning literature (e.g. MIG score, D/C/I, FactorVAE, etc.)
While the tasks are engineered to demonstrate the point of the paper (by synthetically causing entanglement), it would have been nice to see the approach tested in a task that wasn't quite as synthetic and hadn't been formulated precisely to showcase the strengths of the method.
Minor:
The intro contains many different examples; it would be clearer to stick with one running one consistently.
The discussion in Section 4.1 may be inaccessible to someone in RL who is not familiar with graphical models, causal learning, etc. Given that space is an issue, it might be helpful to expand on some of these ideas in the appendix (e.g. explaining the concept of a collider). Perhaps Figure 1 could serve as a diagram where various concepts could be illustrated.
Questions
I'm a bit confused about how the method actually works. Clearly, from the empirical results, it works very well, but the thing that confuses me is that the conditioning set is based on the history of latent representations. But these latent representations change over time as the agent learns. So why does the kNN procedure actually work? If the representation changes over time, then won't the buffer contain "old" representations? For example, if the agent visited a state s early on, it would have some representation that's stored in the buffer. But later, if that same state is visited, the representation will be different (because the encoder has been updated since). I can see how things work if z is fixed, but I'm a bit perplexed that it works even though z is continuously changing. Any clarification here would help.
Learning representations conditioned on history is reminiscent of recent work that seeks to learn a Markov representation (Allen et al., 2021). I'm wondering if there's a link there or if the ideas here are related to the Markov property in some way?
Allen, Cameron, et al. "Learning Markov state abstractions for deep reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 8229-8241.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
The limitations of the method with regard to simple, synthetic correlations are discussed well. I wonder if one further limitation could be the reliance on the Euclidean distance metric. While it will likely be fine in the case of the state representation and continuous actions, I'm wondering specifically about the case of discrete actions, since the IDs attached to each action are arbitrary and so Euclidean distance between actions is meaningless.