Structured Predictive Representations in Reinforcement Learning

The idea of combining GNNs with a forward prediction model for learning predictive embeddings is interesting. But its contribution is limited, as it has been well investigated to learn representations by predicting raw future observations[1] or GNNs[2] in RL.

• The main baseline, min-AIS presented in (Ni et al., 2024), was evaluated in 20 tasks from the MiniGrid benchmarks and several MuJoco tasks. I don't understand why the experiments only consider 4 MiniGrid tasks. I think drawing a conclusion based on the results on the picked 4 tasks is unsuitable.

• According to Figure 4, min-OP is comparable to its counterpart using GNNs on 3 of 4 task, namely DoorKey, KeyCorridorS3R2, UnlockPickup tasks. Therefore, the Graph-based representation learning methods DO NOT outperform the MLP-based methods in all the cases. Therefore the conclusion in line 413 is inaccurate.

• In Figure 5, the interval of min-OP is overlapped with Graph-OP in the presence of distraction changes (a) or size changes (b). So it is wrong to say that Graph-OP outperforms min-OP in ALL scenarios.

[1] Improving sample efficiency in model-free reinforcement learning from images, AAAI 2020.
[2] Contrastive learning of structured world models, ICLR 2020.

1. The paper is largely influenced by the reference Ni et al. (2024), from the problem formulation, method, to experiments. However, some parts inherited from the reference are not related to this paper, which seems irrelevant, as the Q-irrelevant abstractions, model-irrelevant abstraction. The idea is incremental and intuitive. And the contribution compared the reference is minimal.
2. Although the Proposition 3.1 is mentioned in the experimental section, no specific analysis is offered for clear description.
3. Experiments on only one benchmark of MiniGrid is not sufficient, and the performance is also not stably convinced.
4. Beside using the graph neural network as the observation representation learner, transformer is also widely used with better performance, as M-CURL [TPAMI 2023] to learn the state representation, TACO [NeurIPS 2024] to learn the state and action representations, and so on.

There are also some minor writing problems, as

1. “{S2, U, S4, U, S5} (shown in red)” is not consistent with Fig. 2.
2. The same paragraph as the above, the first character of "for instance" should be in capital.
3. Why and how choosing the number of neighbors as m=4?
4. The comparison method as AIS should be clearly stated in this paper to make the paper self contained.

To sum up, the authors are suggested to investigate the literature thoroughly to focus on a more important and clearer problem.

The main problem is the efficiency of the approach compared to other model-based versions of RL: For square grids a set of multi-resolution look-up table would be similarly useful as a maps (or model) in small or medium size grid-world problems. For larger problems, it needs to be be checked whether the GNN does not introduce any errors that could actually be misleading. In other words, the relation of any training errors of the GNN (or a degree of non-stationarity of the task) and the returns would be needed to reasonably assess the benefits of this particular scheme, beyond the general insight that models are useful in well-described (and stationary) environments.

A related problem is the question of applicability to continuous problems, which could have been considered at least in the theoretical part.

Abbreviations (GNN, MLP) need to be be explained in the abstract, because abstracts are to be readable also by non-specialists. The “O” in POMDP is usually “observable” (rather than “observed”).

The role of discount factor in Eq. 1 needs to be reconsidered or explained.

Sect. 2.2: Put a full stop before “This”.

The paper is not particularly clear about the details of the experiments, although these can be inferred from other papers. E.g.: “We periodically increase the size”: If the size is always increased, it is not strictly “periodical”. Also, it is necessary to state the size range (8x8 and above?). To what level this range is explored? Is the size change even visible in the observations? Likewise, the number of "distractors" is only indirectly stated and whether there are particular arrangements of distractors and how they are affecting the agents perception or action, remains open. Also, changes of topology are mentioned, but there is no further information about this.

There's a few areas of improvements that are warranted by the paper.

• While the overall clarity is great, it was not immediately clear in the abstract and introduction that the authors were focused on environments with partial observability (and hence the selection of MiniGrid as their testbed). It's clear that as you read the paper that their work focuses on the POMDP model (history encoder, extending the objectives to the POMDP case, etc.). Additional clarity would make this paper slightly stronger.
  • Specifically, line 17, 53, etc. says in RL, but should be more specific.
• The work only tests their method on a gridworld domain (MiniGrid) and it's not clear if this work can generalize to more complex environments, where perhaps the graph structure is more complicated and similarities between trajectories are less clear.
• Further experiments to understand what the learned structure in the latent space would be very useful to provide further insight into why the GNN-based method works in various settings. Work similar to Figure 1 for other environments enhance the contributions of the paper. Such results would bolster the intuition described in Section 3.1. It would also help strengthen the significance of the paper and it's contributions, because it would shed light on how such structures can apply to other
• The baselines comparison could be substantially stronger. The authors only compare to a version their method without the GNN, but not against other variants or other methods in the literature (such as the ones compared to in Ni et al. or other cited work). Concretely, the authors say "the MLP predicts the next observation for each latent state in a batch and does not use relational reasoning for the whole batch." I understand this to mean that the MLP predicts each state in the batch independently. However, an ablation study where the authors compare to a variant that is an MLP allowed to predict over the entire latent batch, like the GNN, but with fully connected layers instead of message passing, would be helpful in isolating what element of the method makes the largest contribution (or if you need both).
  • Also, if I misunderstood the MLP baseline, and it is actually a fully connected layer that operates over the states in a batch (i.e. is provided the trajectories instead of individual states), then I would make the same argument about comparing to the single-state version. More baselines would be helpful to fully understand the nature of the improvement.

Nits (did not affect score but just flagging for the authors):

• There's a grammatical error in the sentence construction on line 155.
• There's a typo on line 291, extra observation.
• Line 423 has a typo on tMiniGrid, should be MiniGrid.
• Line 537 should be incorporate more algorithms. (into it is not quite grammatical)

Flaws in Notation/Clarity:

• minor notational mistakes (for instance Equations (1), (RP))
• General typos (L189, 228, 292, 294, 330+331,…)
• Usually, “value based methods” refer to methods, that don’t learn a policy (L.119)
• Term “Q-Function” is not explained
• subscript of h seems to have different meanings (L. 217 & L. 124)

Method:

• Proposition 3.1 is not clear, as not all spaces and metrics are formally defined. Moreover, there appears to be no proof for this proposition.
• It is not clear how this Theorem provides the foundation for the hypothesis “Training a model to minimize the loss L by reasoning over both these trajectories ensures that the model generalizes between these subtasks, capturing the similarities between these histories.”, especially how this generalises between subtasks.
• It is not clear how the construction of the graphs (L 274 -279) enforces the intuition behind the example in Figure 2.

Experiments:

• The results on some environments look substantially worse than the results presented in the baseline paper (Comparing Figure 4 here with Figure 6 in (Ni et al., 2024)).
• I assume that the shaded regions in figure 4 and 5 respectively are 95%-CIs. Due to (partially) high overlap, it is not so clear, that the Graph Neural Network based method outperform the MLP counterpart.