A Consciousness-Inspired Planning Agent for Model-Based Reinforcement Learning
Authors: Harry Zhao
Zhen Liu
Sitao Luan
Shuyuan Zhang
Doina Precup
Yoshua Bengio
Venue: NeurIPS 2021 Poster (2021)
Keywords: consciousness
planning
reinforcement learning
deep learning
model-based reinforcement learning
neuro-inspired AI
artificial intelligence
brain-inspired AI

Abstract: We present an end-to-end, model-based deep reinforcement learning agent which dynamically attends to relevant parts of its state during planning. The agent uses a bottleneck mechanism over a set-based representation to force the number of entities to which the agent attends at each planning step to be small. In experiments, we investigate the bottleneck mechanism with several sets of customized environments featuring different challenges. We consistently observe that the design allows the planning agents to generalize their learned task-solving abilities in compatible unseen environments by attending to the relevant objects, leading to better out-of-distribution generalization performance.

A Consciousness-Insprred Planning Agent for Model-Based Reinforcement Learning
Mingde Zhao |A,, Zhen Liu Sitao Luan IA, Shuyuan Zhang!/4
Doina Precup 1,3,4,5t Yoshua Bengio?.4.1
McGill University; Université de Montréal; 'DeepMind; "Mila; CIFAR Al Chair Equal Contribution, t: Equal Supervision
Abstract
We present end-to-end, model-based deep reinforcement learning agent which dynamically attends to relevant parts of its state during planning. The agent bottleneck mechanism over set-based representation to force the number of entities to which the agent attends during planning to be small. In experiments, we investigate the bottleneck mechanism with several sets of customized environments featuring different challenges. We consistently observe that the design allows the planning agents generalize their learned task-solving abilities in compatible unseen environments by attending to the relevant objects, leading to better out-of-distributio generalization performance. Check project page https: thub. com/PonerHarrry/Cr
1 Introduction
Whether when planning our paths home from the office or from a hotel to an airport in an unfamiliar city, we typically focus on small subset of relevant variables, e.g. the change in position or the presence of traffic. An interesting hypothesis of how this path planning skill generalizes across scenarios is that it is due to computation associated with the conscious processing of information 2 3 14 Conscious attention focuses on a few necessary environment elements, with the help of an internal abstract representation of the world 433 14]. This pattern, also known as consciousness in the first sense (C1) [1], has been theorized to enable humans' exceptional adaptability and learning efficiency 2 14 43 15). A central characterization of conscious processing is that it involves bottleneck, which forces one handle dependencies between very few environmental characteristics at a time [1 15).
may seem limiting, facilitates Though focusing on subset of the available information (OOD) and systematic generalization to other situations where the ignored variables are different and yet still irrelevant 7 15
In this paper, we encode some of these ideas into reinforcement learning agents. Rein- forcement learning (RL) is an approach for learning behaviors from agent-environment interactions [1 However, most of the big successes of RL have been obtained by deep, model-free agents 30 37 38]. While Model-Based RL (MBRL) has generated significant research due to the potentials of using an extra model [31] its empirical performance has
typically lagged behind, with some recent notable exceptions 36 24
Our proposal is to take inspiration from human consciousness to build an architecture which learns a useful state space and in which attention can be focused on a small set of variables at any time, where the aspect of "partial planning" " enabled by modern deep
Partial planning is interpreted in different ways. For example, concurrent work [66] focuses on modelling "affordable" temporally extended actions, s.f. "intent" could be achieved more efficiently.
35th Conference Neural Information Processing Systems (NeurIPS 2021).

---

RL techniques 42 26
Specifically, we propose an end-to-end latent-space MBRL agent
which does not require reconstructing the observations, as in most existing works, and uses Model Predictive Control (MPC) framework for decision-time planning + 35 From an observation, the agent encodes a set of objects as a state, with a selective attention bottleneck mechanism to plan over selected subsets of the state (Sec. 4). Our experiments show that the inductive biases improve a specific form of OOD generalization, where consistent dynamics are preserved across seemingly different environment settings (Sec. 5
2 Background & Context
We consider an agent interacting with its environment at discrete timesteps. At time f, the agent receives observation and takes action 4t, receiving a reward fri1 and new observation 01+1- The interaction episodic. The agent also building a latent-space transition model, M, which can be used to sample a next state, $1+1, a reward fi+1 and binary signal ii+1 which indicates if the model predicts termination after the transition. We will now compare and contrast our approach with some existing methods from the MBRL literature, explaining the rationale for our design choices.
Observation Level Planning and Reconstruction as Latent Space Planning
Many MBRL methods plan in the observation space or rely on reconstruucioo--bseed losses to obtain state representations [24 36 17 48) Appropriate as these methods may be for some robotic tasks with few sensory inputs, e.g. continuous control with joint states, they are arguably difficult with high-dimensional inputs like images, since they may focus on predictable yet useless aspects of the raw observations 31 Besides suffering from the need to reconstruct noise or irrelevant parts of the signal, it is not clear if representations built by a reconstruction loss (e.g. L2 the observation space) are effective for an MBRL agent to plan or predict the desired signals (39 17 18 egg. values (in the RL sense), rewards, elc.. In this work, we use an approach similar to those 33 36 17 building a latent space representation that is jointly shaped by all the relevant RL signals (to serve value estimation and planning) without using reconstruction.
Staged Training us End-to-End Training
use two explicit stages of training: (1) an inner representation of the world is trained using exploration (usually with random trajectories); (2) the representation is fixed and used for planning and MBRL. Despite the advantages of being more stable and easier to train, this procedure relies on having an environment where the initial exploration provides transitions that are sufficiently similar to those observed under improved policies, which is not the case in many environments. Furthermore, the learned representation may not be effective for value estimation, if these transitions do not contain reward information that can be used to update the input-to- representation encoder. End-to-end MBRL agents, egg. [39 (3), are able to learn the representation online, simultaneously with the value function, hence adapting better to non-stationarity in the transition distribution and rewards.
Some MBRL agents based on world model [6 24
Type of planning
MBRL agents can use the model in different ways. Dyna 40 learns model to generate "imaginary" transitions, which contribute to the training of the value estimator 40), in addition to the real observations, thus boosting sample efficiency. However, if the model is inaccurate, the transitions it generates may be "delusional", which may alter the value estimator and negatively impact performance. Moreover, Dyna is typically used to generate extra transitions from the states visited in trajectory, and updates the model based on the observed transitions as well. This means Dyna is focused on the data distribution encountered by the agent and may have trouble generalizing OOD. In contrast, simulation- based model-predicte
control (MPC) and its variants 34 35 18 only update the value estimator based on real data, using the model simply to perform lookahead at decision-tima Hence, model inaccuracies impact less, with more favorable OOD generalization capabilities. Hence, MPC is adopted in our approach.
Vectorized us Set Representations for RL Most Deep Reinforcement Learning (DRL) work focus on learning vectorized state rep-
2

---

extract
vectorized representation
features
state vector
extract
set-based representation
observation
features
to set unorder
feature
position
Figure 1: Set-based state encoder compared to classical vectorized state encoders: the feature map extracted by some feature extractor, egg. CNN, "chopped" into feature vectors and concatenated with positional information. the features of observed entities. The permutation-invariance of set computations forces the learner to be robust to small changes in the set (e.g. one of the elements being different missing).
the resulting concatenations are treated as objects a set, capturing
resentations, where the agents' observation is transformed into feature vector of fixed dimensionality 19]. Instead, set-based encoders, a.k.a. object-orienter architectures, are designed to extract a set of unordered vectors from which predict the desired signals via permutation-invarian have shown the promise of set-based representations capturing environmental states, in terms of generalization, as well as their similarities to human perception [13 47
computations [00,, as illustrated in Fig. [ Recent works in RL
32
46 29
Additionally in this work, we utilize the compositionality of set representations to enable the discovery of sparse interactions among objects, i.c. underlying dynamics, well as to facilitate the bottleneck mechanism, analogous C1 selection. The set-based representatin coupled with the bottleneck provides an inductive bias consistent with selecting only the relevant aspects of a situation on-the-fly through an attention mechanism. The small size of the working memory bottleneck also enforces sparsity of the dependencies [Z] [5] captured by the learned dynamics model: each transition can only relate a few objects together, no more than the size of the bottleneck.
3 MBRL with Set Representations
We present an end-to-end baseline MBRL agent that uses a set-based representation and carries out latent space planning, but without consciousmess-inspirees small bottleneck. This agent serves as baseline to investigate the OOD generalization capabilities brought by the bottleneck, which is to be introduced later in Sec.
The mapping from observations to values is a combination of an encoder and a value estimator. The encoder maps observation vector to set of objects, which constitutes the latent state. set-to -eeetor architecture that maps the latent
The value estimator a permutation-invant state to value estimate. Note that the same state set is used for all the agents' predictions, including future states, rewards etc., as we will discuss later.
For image-based observations, we use the features at each position of the CNN output feature map to characterize the feature of an object, similar [9) as shown in Figure [
Encoder.
To recover positional information lost during the process, we concatenate each object feature vector with a positional embedding to form a complete object embedding. Such approach is different from the common practice of mixing positional information by addition [4]. This is for the compatibility with our dynamics model training procedure, discussed below.
(State-Action) Value Estimator takes the form S where S the learned state space by the set-based encoder (hoping to capture the real underlying state space of the MDP) and
discrete action set. We use an improved architecture upon DeepSets 50 depicted in Figure?? 2 The architecture performs reasoning on set of encoded objects, resembling pervasive usage in natural language processing, where the objects are typically word tokens 33

---

Permutation- Invariant Ops
mean
feature
FC
SA
FC
t
position
intermediate set
summary
Q values
Figure 2: Value estimator Q and a generic set-to-vector architecture: we modify the design of DeepSets 50 by replacing the MLP before pooling with transformer layers (multi-head Self-Attention (SA) + object-wise Fully Connected (FC)) 45). We found this change to be helpful for performance. After applying the transformer layers, the intermediate set (colored gray) entangles features and positions. Please check the Appendix for more details on the self-attention operations involved.
Transition Model. The transition model maps from St, ai to Š1-1, fr and @1+1- We separate this into: 1) the dynamics model, in charge of simulating how the state would change with the input of a, and 2) the sewardterminationn estimator which maps $,, aj to 9, and @
While designing
estimator is straightforward (a two-headed augmented architecture similar to the value estimator), the dynamics model requires regression unordered sets of objects (set-to-set). A common approach is to use matching methods, e.g. Hausdorff distance, However, they are computationally demanding and subject to local optima 15, 8 28]. Targeting this, our feature-positinn computations position-aware, but also
Chamfer matching
separated set
encoding not only makes the permutation-invarian allows simple end-to-end training over the dynamics. By forcing the positional tails to be immutable during the computational pass, we can use them to solve the matching trivially: objects "labeled" " with the same positional tail in the prediction S1+1 (output of the dynamics model) and the training sample $ (state obtained from the next observation) are aligned, forming pairs of objects with changes only in the feature, as shown in Figure 3
Tree Search MPC. The agent employs tree-search based behavior policy (with e-greedy exploration). During planning, each tree search call maintains priority queue of branches to simulate with the model. When designated budget (e.g. number of steps of simulation) is spent, the agent greedily picks the immediate action that leads to the most promising path. We present the pseudocode of the Q-value based prioritized tree-search MPC in Appendix.
Equivalence could be drawn from this planning approach to Monte-Carlo Tree Search (MCTS) 37 38). While this method is far more simplistic and require fewer simulations for each planning call (see example in Appendix).
Training. The proposed agent is trained from sampled transitions with the following losses:
Temporal Difference (TD) LTD: regresses the current value estimate to the update target, e.g. calculated according to DQN Double DQN (DDQN) 30 44). In experiments, distributional output is used for both value and reward estimation, making this loss KL-divergence [6].
Dynamics Consistency Layn' A L2 penalty established between the aligned 6 and St+1, where $111 is the imagined next (latent) state given 0,, a, and 6441 is the true next (latent) state encoded from 0++1-
Reward Estimation L,: the KL-divergence between the imagined reward f/+1 predicted by the model and the true reward Fr+1 of the observed transition.
Termination Estimation La: the binary cross-entropy loss from the imagined termination @i1 to the ground truth WTi1, obtained from environment feedback.

---

action
Nx
downscale FC
imagined feature
feature
SA
+
FC
position
position
St
copy original tails-
$1+1
Figure 3: Dynamics model: for FC sub-layers of the transformer layers, we inject action embedding s.f. the transformer computations are now action conditioned. After getting the intermediate set, we downscale each of the objects, leaving the positions untouched and directly copied from the input St- FC downscale is a linear transformation which downscales the dimensionality of the intermediate objects to that of the features part of objects (before the layernorm). In this way, after concatenating the positional tails the objects have consistent dimensionality. Intuitively, each object slot recovers its positional tail at the output. Though the objects in the sets input-intermediite outputt are aligned, within each set they are still unordered, i.e. permutation-invarian..
The resulting total loss for end-to-end training of this set-based MBRL agent is thus
L LTD Layn + L La
Jointly shaping the states avoids the representation collapsing trivial solutions and makes the representation useful for all signal predictions of interest.
4 Consciousness-Ispprre Bottleneck
In this section, we introduce an inductive bias which facilitates Cl-capable planning. In nutshell, the planning is expected to focus on the parts of the world that matter for the plan. Simulations and predictions are all expected to be performed on (small) bottleneck set, which contains all the important transition-related information. As illustrated in Figure the model performs 1) selection of the bottleneck set from the full state-set, 2) dynamics simulation on the bottleneck set and 3) integration of predicted bottleneck set to form the predicted next state.
Conditional State Selection We select a bottleneck set c; of n objects from the potentially large state set s, of m > objects. Then we only model the transition for the selected objects in Ct. To make this selection, we use a key-query-value attention mechanism, where the key and the value for each object in 9; are obtained from that object, and the query is a function of some learned dedicated set of vectors and of the action considered (see Appendix for details). Inspired by the work on self-attention for memory access [25,, we use semi-hard top-k attention mechanism to facilitate the selection of the bottleneck set. That is, after the query, the top-k attention weights are kept, all others are set 0, and then the attention weights are renormalized. This semi-hard attention technique limits the influence of the ill-matched objects on the bottleneck set c; while allowing for a gradient to propagate on the assignment of relative weight to different objects. With purely soft attention, weights for irrelevant objects are never 0 and learning to disentangle objects may be more difficult.
Dynamics/ Neward-Terminatie tecture as described in Sec. 3 but taking the bottleneck objects as input rather than the full state set. Details of the architecture are in the Appendix.
Prediction on Bottleneck Sets. We use the same archi-
Change Integration. An integration operation, intuitively the inverse operation of selection, is implemented to 'soft paste-back' the changes of the bottleneck state onto the state set st,
IIn our experiments, no re-weighting is used for each term the total loss. This is possible for the fact that they are in similar magnitudes. In our experimental implementation, no recurrent mechanism used however the same training procedure is naturally extendable.

---

rrward-term estimator
at
St
dynamiics modde
no bottleneck with botfleneck
reward term estimator
at
St
s/+1
dymanmics model
select
integrate
Ct
E1+1
Figure 4: Bottleneck stages (operations colored in purple are conditioned on a chosen action): 1) bottleneck set C soft-selected from the whole state (object set) sr through semi-hard multi-head attention; 2) dynamics
applied the bottleneck set e to form 1; the reward and termination signals are predicted from Ct, c and ar- Then, the changes introduced in ē are integrated with 8t to obtain 81+1, the imagined next state, with the help of attention. Note that the two computational flows in stage 3 are naturally parallelizable.
yielding the imagined next state set 8e+1. This is also achieved by attention operations, more specifically querying Ćt-1 with St, conditioned on the action a;. Please check the Appendix for more details.
Discussion. The bottleneck described in this section natural complement to the MBRL model with set representations discussed previously. In particular, planning and training are carried out the same way as discussed in Sec. 3
We expect the Conscious Planning (CP) agent to demonstrate the following advantages:
Higher Quality Representation: the interplay between the set representation and the selection integration forces the representation to be more disentangled and more capable of capturing the locally sparse dynamics.
More Effective Generalization: only essential objects for the purpose of planning participate in the transition, thus generalization should be improved both in-distribution and OOD, because the transition does not depend on the parts of the state ignored by the bottleneck.
Lower Computational Complexity: directly employing transformers to simulate the full state dynamics results in a complexity of ((Is,/d) where is the length the objects, due to the use of Self-Attention (SA), while the bottleneck lowers it to ((sr)lec)dd)
(a) In-dist, diff 0.35 (b) OOD, diff 0.25 (c) OOD, diff 0.35 (d) OOD, diff 0.45 OOD, diff 0.55
Figure 5: Non-Static RL Setting, with in-distribution and OOD tasks: (a) example of training environments
examples of OOD environments (rotated 90 degrees, changing the distribution of grid elements). For OOD testing, we evaluate different levels difficulty (b e). The agent (red triangle) points in the forward movement direction. The goal is marked in green. For each episode (training or OOD), we randomly generate new world from sampling distribution. Note that the training environments and the OOD testing environments have intersecting observations.

---

Experiments
We present our experimental settings and ablation studies of our CP agent against baselines to investigate the OOD generalization capabilities enabled by the Cl-inspired bottleneck mechanism. To clarify, the OOD generalization we refer to specifically is the agents' ability to generalize its learned task skills across seemingly different tasks with common underlying dymamics. Take the set of experiments in this section for example, we want the agent to be able to generalize its navigation skills in unseen environments.
5.1 Environment/ Task Description
We use environments based on the MiniGrid-Baby-Al framework [1 10 21 which can be customized for generating OOD generalization tests with varying difficulties. To make sure we assess the agents as clearly as possible, the customized environments feature clear object definitions, with well-understood underlying dynamics based on object interactions. Furthermore, the environments are solvable by Dynamic Programming (DP) and can be easily tuned to generate OOD evaluation tasks. These characteristics are crucial for the experimental insights we are seeking.
In this section, the experiments are carried out on 8 x 8 gridworld as shown in Figure 5 The agent (red triangle) needs to navigate (by turning left, right or stepping forward) to the goal while dodging the lava cells along the way If the agent steps into lava (orange square), the episode terminates immediately with no reward. If the agent successfully reaches the goal (green square), receives a reward of +1 and the episode terminates. For better generalization, the agent needs to understand how to avoid lava in general (and not at specific locations, since their placement changes) and reach the goal as quickly as possibleF) The environments provide grid-based observations that are ready to be interpreted as set representations: each cell of the observation array is an object, thus resulting in a set of 64 objects in s; for each observation.
For the agent be able to understand the environment dynamics instead of memorizing specific task layouts, we generate a new environment for each training or evaluation episode. In each training episode, the agent starts at random position on the leftmost or rightmost edge and the goal is placed randomly somewhere along the opposite edge. In between the two edges, the lava cells are randomly generated according which controls the probability of placing lava cell at each valid position. The difficulty parameter controls partially how seemingly different the OOD evaluation tasks are to the in-distribution training tasks, though we know the underlying dynamics of all these tasks are the same. For training episodes, the difficulty is fixed to 0.35. We note that most usual RL benchmarks contain fixed environments, where the agent is expected to acquire specific optimal policy. These environments are ill-suited for our purpose.
a difficulty parameter
For OOD evaluation, the agent is expected to adapt in new tasks with the same underlying dynamics in a 0-shot fashion, i.e. with the agent's parameters fixed. The OOD tasks are crafted to include changes both in the support (orientation) and in the distribution (difficulty): the agent is deployed in transposed layouts with varying levels of difficulty ([0.25, 0.35, 0.45, 0.55)). The differences of in-distribution (training) and OOD (evaluation) environments are illustrated in Figure
5.2 Agent Setting
We build all the set-based MBRL agents included in the evaluation on a common model-free baseline: set-based variant of Double-DQN (DDQN) 44 with prioritized replay and distributional outputs. For more details, please check the Appendix.
'W provide additional results for world sizes ranging from 6 x 6 to 10 x 10 in the Appendix. 8 8 chosen as the demonstrative case.
IIn the Appendix, we provide additional test settings with different dynamics, which also demon-
strates the agents' ability to work well despite cluttering distractions.
?Please check the Appendix for extra sets of tasks with different agent actions and task objectives.
"The agent starts at the top or bottom edge and the goal is respectively on the bottom or top edge, whereas training environment has the agent and goal on the left right edges
7

---

We compare the proposed approach, labelled CP in the figures (for Conscious Planning) against the following methods:
methods:
UP (for Unconscious Planning): the agent proposed Section lacking the bottleneck. model-free: the model-free set-based agent is the basis for the set-based model- agents. It consists of only the encoder and the value estimator, sharing their architectures with CP and UP.
Dyna: the set-based MBRL agent which includes a model-free agent and an observation- level transition model, i.e. transition generator. For the model, we use the CP transition model (with the same hyperparameters as the best performing CP agent) on the original environment features without an encoder. We also use the same hyperparameters as in the CP model training. The agent essentially doubles the batch size of the model-free baseline by augmenting training batches with an equal number of generated transitions.
Dyna": A Dyna baseline that uses the true environment model for transition generation. This is expected to demonstrate Dyna's performance limit.
WM-CP: A world model CP variant that differs by following 2-stage training procedure
[16]. First, the model (together with the encoder) is trained with 10 random transitions. After this, the encoder and the model are fixed and RL begins.
NOSET: A UP-counterpart with vectorized representations and no bottleneck mechanism.
Particularly, for CP and UP agents, we also test the following variants:
CP-noplam: A CP agent that trains normally but does not plan in OOD evaluations, ie. carrying out model-free behavior. This baseline aims to demonstrate the impact of planning in the training process on the OOD capability of the value estimator.
UP-noplan: UP counterpart of CP-noplan.
Note that the compared methods share architectures as much as possible to ensure fair comparisons. Details of the compared methods, their design and hyperparameters are provided in the Appendix.
5.3
Performance Evaluation
5.3.1 In-Distribution
In Figure 6 we present the in-distribution evaluation curves for the different agents. For UP, CP and the corresponding model-free baselines, the performance curves show no significant difference, which demonstrates that these agents are effective in learning to solve the in-distribution tasks. During the "warm-up" period of the WM baseline, the model learns a representation that captures the underlying dynamics. After the warm-up, the encoder and the model parameters are fixed and only the value estimator learns to predict the state-action values based on the given representation. The increase in performance not only delayed due to the warm-up phase (during which rewards are not taken into account) but also harmed, presumably because the value estimator has no ability to shape the representation
better suit its needs. The Dyna baseline performs badly while the Dyna" baselines perform relatively well. This is likely due to the delusional transitions generated by the model at the early stages of training, from which the value estimator never recovers. However, the Dyna" baseline does not achieve satisfactory OOD performance (Figure 7, presumably because its planning only focuses on observed data, and hence only improves the in-distribution performance, due to insufficiently strong generalization. The NÓSET baseline performs very badly even in-distribution, per Figure In the Appendix, show that the NOSET baseline seems only able to perform well in more classical, static RL setting, which may indicate that relies on memorization. We provide more results regarding the model accuracy in the Appendix.
5.3.2 OOD Task-Solving Performance
The OOD evaluation focuses on testing the agents' performance in a set of environments forming gradient of task difficulty. In Figure 7 we present the performance error bars of
8

---

1-
0.8
CP(8)
Figure 6: In-distribution task perfor- mance: the x-axis shows the training progress (2.5 x 10 agent-environment in- teractions). The y-axis values gener- ated by agent snapshots times corre- sponding to x-axis values. CP, UP, model-free and Dyna* agents all learn to solve the in-distribution tasks quickly. All bars obtained from 20 indepen- dent runs.
success rale
UP
0.6
modelfree
WM
0.4
Dyna
0.2
Dyna*
NOSET
1.5
2.5
agent-env interactions x10
CP(8)
modelfrree WM
Dyna Dyna
NOSET
UP
0.6
02
0.5
1.5
2
25
0.5
1.5
2.5
0.5
1.5
2.5
0.5
1.5
Interactions
agentt-ctti interactions
agend Insoractions
(a) OOD, difficulty 0.25 (b) OOD, difficulty 0.35 (c) OOD, difficulty 0.45 (d) OOD, difficulty 0.55
Figure 7: OOD performance under gradient of difficulty. The figures show consistent pattern: the MPC-based end-to-end agent equipped with bottleneck (CP) performs the best. All error bars are obtained from independent runs.
the compared methods under different OOD difficulty levels. CP(8), CP with bottleneck size n 8, shows clear performance advantage over UP, validating the OOD generalization capability. The Dyna* baseline, essentially the performance upper bound of Dyna-basee planning methods, shows no sigmificant performance gain in OÒD tests compared to model- free methods. WM may have the potential to reach similar performance to warm up the encoder with large portion of the agent-enivirnmeet interaction budget, free unsupervised phase is provided. We dive into this matter in the Appendix.
CP, yet it needs
5.3.3 Ablation
We validate design choices with ablation. Figure g visualizes two of these experiments. For more ablation results, which include validation of the effectiveness of different model choices, and further quantitative measurements, egg. of OOD ability as function of behavior optimality and model accuracy, please check the Appendix.
5.4 Summary of Experimental Results
With the scope limited to our experiments, the results allow us draw these conclusions:
Set-based representations enable at least in-distribution generalization across different environment instances in our non-static setting, where the agents are forced to discover dynamics that are preserved across environments;
Model-free methods seem to face more difficulties in solving our OOD evaluation tasks which preserved the same environment dynamics to the corresponding in-distribution training settings;
MPC exhibits better performance than Dyna in the tested OOD generalization settings;
Online joint training of the representation with all the relevant signals could bring benefits to RL, as suggested in 22]. Please check Appendix E for more discussions of this matter;
In accordance with our intuition, transition models with bottlenecks tend to learn dynamics better our tests. This likely for they prioritize learning the relevant aspects, while models without bottleneck may have to waste capacity on irrelevance;

---

CP(8)
CP(16)-best
0.6
n8)
CP(8)-best
CP(s)-best
0.5
Hoplan(LIP)
stccess rate
success rate
CP(16)
modieffrce
0.4
CP(8)
0.4
CP(4)
0.2
0.1
0
0
0.5
1.5
2.5
0.5
1.5
2.5
agent-env interactions
x10
agent-env interactions
x10
(a) Bottleneck benefits OOD capability: no- plan(8) and noplan(UPP) correspond to the CP(8) and UP variants with planning disabled during OOD tests. Comparing noplan against model- free,
(b) Value estimators do generalize well in our OOD tests: random heuristic significantly outperforms best-first heuristic OOD.
that planning during training is beneficial for both value estimation and repre- sentation learning.
Figure 8: Key ablation results: With diff 0.35, each error bar obtained from 20 independent runs.
From further experiments provided in the Appendix we observe that bottleneck- equipped agents may also be less affected by larger environmental scales, possibly due to their prioritized learning of interesting entities.
6 Conclusion & Limitations
We introduced a conscious bottleneck mechanism into MBRL, facilitated by set-based representations, end-to-end learning and tree search MPC. In the non-static RL settings, the bottleneck allows selecting the relevant objects for planning and hence enables significant OOD performance.
One limitation of our work is the experimental focus on only Minigrid environments, due to the need to validate carefully our approach. For future works, we would also like to extend these ideas to temporally extended models, which could simplify the planning task, and are also better suited conceptual model of C1. Finally, we note that the architectures we use are involved and can require careful tuning for new types of environments.
Acknowledgements
Mingde is grateful for the financial support from the Fonds de Recherche du Québec Nature et Technologies (FRQNT). Yoshua acknowledges the financial support from Samsung Electronics and IBM.
We acknowledge the computational power provided by Compute Canada. We are also thankful for the helpful discussions with Xiru Zhu (about the design of the environment generation procedure), David Yu-Tung Hui (about the bag-of-word representations, insights BabyAA AI well about the writing of the introduction section), Min Lin (about the design of the dynamics model as well as the early stage brainstorming) and lan Porada (for consistently supporting the student authors).
10

---

References
[1] L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXio, 1607 1607.06450, 06440, 2016. www.iii..rgggbs////o 1607.06450
[2] B. Baars. cognitive theory of consciousness. Cambridge University Press, 1993.
[3] B. Baars. The conscious access hypothesis: origins and recent evidence. Trends in cognitive sciences, 6(1):47-52, 2002.
[4] D. Bahdanau, K. Cho, and Y. Bengio. Neural machine translation by jointly learning to align and translate. arxiv, 1409.0473, 2014. https///ara....og/aab///// 0473
[5] H. G. Barrow, J. M. Ienenbaum, R. C. Bolles, and H. C. Wolf. Parametric correspondence and Chamfer matching: new techniques for image matching. Technical report, International Menlo Park CA Artificial Intelligence Center, 1977.
repoor, SRI
[6] M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective inforcement learning. International Conference on Machine Learning, 2017. https
[7] Y. Bengio. The consciousness prior. ar Xiv, 1709.08568, 2017. http://arxiv.org/
[8] G. Borgefors. Hierarchical chamfer matching: A parametric edge matching algorithm. TEEE Transactions on pattern analysis and machine intelligence, 10(6):849 88008855,, 1988.
[9] N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko. End-to- end object detection with transformers. European Conference on Computer Vision, 2020.
https://arxiv.orgg absp://bs///00// 12872
[10] M. Chevalier-Boisvert, D. Bahdanau, S. Lahlou, L. Willems, C. Saharia, T. H. Nguyen, and Y. Bengio. Babyai: A platform study the sample efficiency of grounded language learning. International Conference Learning Representations, 2018. http:
//arxiv.org/alte hatggg/bs///1111 1ess/1810.000 08272
[11] M. Chevalier-Boisvert L. Willems, and S. Pal. Minimalistic ginim environment for openai gym. GitHub repository, 2018.
gym sym-minigrid
[12] R. C. Conant and W. Ross Ashby. Every good regulator of system must be model of that system. International Journal of Systems Science, 1(2):89-97) 1970.
[13] G. Davidson and B. M. Lake. Investigating simple object representatiions in model-free deep reinforcement learning. arxiv, 2002.06703, 2020. https://arxiv.orgs 2002 06703
[14] S. Dehaene, H. Lau, and S. Kouider. What is consciousness, and could machines have it? Science, 358, 2020.
[15] A. Goyal and Y. Bengio. Inductive biases for deep learning of higher-level cognition.. ar Xiv, 2011 2011.15091, 2020. https://araiv.org/////000..111
[16] D. Ha and J. Schmidhuber. Recurrent world models facilitate policy evolution. In Conference on Neural Information Processsing Systems, volume 31, 2018 htttss://papee.
Abstract.ttttttt
[17] D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mastering Atari with discrete world models. In International Conference Learning Representations, 2021. https
//arxiv. orw.vv.rggaaaa aass//bss00000 02193
[18] J. B.B. Hamrick, A. L. Friesen, F. Behbahani, A. Guez, F. Viola, S. Witherspoon, T. Anthony,
I.. Buesing, Velickovic, and T. Weber. On the role of planning in model-based deep reinforcement learning. Xiv, 2011.04021 2020.
04021
11

---

[19] M. Hessel, J. Modayil, H. van Hasselt, T. Schaul, G. Ostrovski, W. Dabney, D. Horgan, B. Piot, M. G. Azar, and D. Silver. Rainbow: Combining improvements in deep reinforcement learning. AAAI Conference on Artificial Intelligence, 2017. http://arxiv org/abs/1710. 02298
[20] D. Horgan, J. Quan, D. Budden, G. Barth-Maron, M. Hessel, H. van Hasselt, and D. Silver. Distributed prioritized experience replay. International Conference on Learning Representations, 2018. http://arxiv. wwwpaaawww.ii.....aaaaaaaaa.aaaa /tt.//aabb//10.003333333
[21] D. Y.-T. Hui, M. Chevalier-Boisvert, D. Bahdanau, and Y. Bengio. Babyai 1.1. arviv, 2007.12770, 2020. http://arxiv.org/abs//
12770
[22] M. Jaderberg, V. Mnih, W. M. Czarnecki, T. Schaul, J. Z. Leibo, D. Silver, and K. Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. Interna- tional Conference on Representation Learning, 2017.
[23] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. arxiv, 1906.08253, 2019. http://arxiv.oog/abbb/////
08253
[24] L. Kaiser, M. Babaeizadeh, P. Milos, B. Osinski, R. H. Campbell, K. Czechowski, D. Erhan, C. Finn, Kozakowski, Levine, et al. Model-based reinforcement learning for atari. abs 1ttps//1130000014444
arXio, 1903.00374, 2019. http://arxiv.oo
[25] N. R. Ke, A. Goyal, O. Bilaniuk, J. Binas, M. C. Mozer, Pal, and Y. Bengio. Sparse atten- tive backtracking: Temporal credit assignment through reminding. arXiv, 1809.03702, 2018. http://arxix> wrg/ abs/ 03702
[26] K. Khetarpal, Z. Ahmed, G. Comanici, and D. Precup. Temporally abstract partial models. In Conference Neural Information Processing Systems, 2021. http://arxiv org/abs/2108. 03213
[27] D. P. Kingma and J. Ba. Adam: A method for optiiiaaii International Conference on Learning Representations, 2014. http://arraiv.or//abs///222 6980
[28] A. R. Kosiorek, H. Kim, and D. Rezende. set with transformers. arv, 2006.16841, 2020.
[29] Löwe, K. Greff, R. Jonschkowski, A. Dosovitskiy, and Kipf. Learning object-centric video models by contrasting sets. arXiv, 2011.10287, 2020. https://arxiv.. wwwpexivvv.org//ao / 2011.10287
[30] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015. https:/ http://www..aurreeeeeeeee com/ articles/naturel.22
[31] M. Moerland, J. Broekens, and C. M. Jonker. Model-based reinforcement learning: A 16712
survey, arXiv, 2006.16712, 2020.
[32] T. Mu, J. Gu, Z. Jia, H. Tang, and H. Su. Refactoring policorrrrooooyyyyooooo for compositional gener- alizability using self-supervised object proposals. In Conference on Neural Information Processing Systems, volume 33, 2020.
[33] I. Porada, K. Suleman, A. Trischler, and J. C. K. Cheung. Modeling event plausibility with consistent conceptual abstraction. In Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1732-1743, 2021. http:ss////zt..oorgb////l....
[34] A. V. Rao. A survey of numerical methods for optimal control. Advances in the Astronautical Sciences, 135(1):497-528,
2009. 2009.
[35] A. G. Richards. Robust constrained model predictive control. PhD thesis, Massachusetts www. wttp::tt.ttttaaade////...../////
Institute of Technology, 2005.
12

---

[36] J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, I. Graepel, al. Mastering Atari, Go, Chess and Shogi by planning with a learned model. Nature, 588(7839):604 609, 2020. https://axxiv.orr abs/1911. 08265
[37] D. Silver, A. Huang, C. Maddison, A. Guez, L.. Sifre, G. Van Den Driessche, J. Schrit- twieser, I. Antonoglou, V. P'anneershllaam, M. Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484 2016. hetps://wwwnnttee aawaaurre.cc//rttulll///aatuttaa
[38] D. Silver, J. Schrittwieser, K. Simonyan, I. Antonoglou, A. Huang, A. Guez, T. Hubert, L. Baker, M. Lai, A. Bolton, et al. Mastering the game of go without human knowledge. Nature, 550(7676):354-33333 359, 2017. https://wwat //www. nature. wwatuuee.om/arteeateeeeet comm//rrriiceess/mtut2227777
[39] D. Silver, H. van Hasselt, M. Hessel, T. Schaul, A. Guez, T. Harley, G. Dulac-Amold, D. P. Reichert, N. C. Rabinowitz, A. Barreto, and T. Degris. The predictron: End- to-end learning and planning. International Conference on Machine Learning, 2016.
http://arxiv.ore org/ httgg//bs/1612. 08811.8.1100
[40] R. S. Sutton. Dyna, an integrated architecture for learning, planning, and react- ing. SIGART Bulletin, 2(4):160-163, 1991. https://dl.aet httts:///d..ccmorr//66////.. orgp:///1.//////o/m /t/.m///o//////11111111144444 122377
1145/122344
[41] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018. http://inncompleeeddeas....eee www.seeeeetoootttbeeeeeeee
html
[42] E. Talvitie and S. Singh. Simple local models for complex dynamical sys- tems. In Conference on Neural Information Processing Systems, volume 21, pages 1617-1624. Citeseer, 2008, https://paperrs http:://pppprs..is.ccc/ccc
[43] R. van Gulick. Consciousness. In E. N. Zalta, editor, Stanford Encyclopediaa of Philosophy. Stanford: Metaphysics Research Lab, 2004.
[44] H. van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. AAAI Conference on Artificial Intelligence, 2015. http://arxiv. org/ / 1 1509.0665 .6461
[45] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. International Conference on Neural Information Processing Systems, 2017. htttps://axxi..org/abaa
03762
[46] O. Vinyals, I. Babuschkin, W. M. Czarnecki, M. Mathieu, A. Dudzik, J. Chung, D. H. Choi, R. Powell, T. Ewalds, P. Georgiev, et al. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575(7782)3300-35,, 2019. https:
//www.nature hotpe////mmmtiiles///155566 010-866-119-2222 -19-17
[47] T. Wang, R. Liao, J. Ba, and S. Fidler. Nervenet: graph neural networks. In International Conference on Learning Representations, 2018.
Learning structured policy with
https://openrevier httpi/rrwwiiieeeeorrmummmmmi
[48] X. Wang, W. Xiong, H. Wang, and W. Y. Wang. Look before you leap: Bridging model- free and model-based reinforcement learning for planned- vision-aand-languuae navigation. In European Conference Computer Vision, 2018. http://arxiv.org/ 0803.07729
[49] M. Xie and S. Ermon. Differentiable subset sampling. In International Joint Conference on Artificial Intelligence, 2019. https://arxiv.orrg wwt.orggg/bbb...ooe///// 10517
[50] M. Zaheer, S. Kottur, S. Ravanbhakhsh, B. Póczos, R. Salakhutdinov, and A. J. Smola. Deep sets. Conference Neural Information Processing Systems, 2017. https://arxiv org/abs/1703. 111.b/1/71000111 06114
[51] A. Zakharov, M. Crosby, and Z. Fountas. Episodic memory for learning subjective- timescale models. arv, 2010.01430, 2020.
00130
13


Reviews:
[{"review_id": "QUp8U-kf3R", "reviewer": "Reviewer_sQ3E", "summary": "The authors introduce a model-based deep reinforcement learning algorithm that uses tree search based MPC over a learned latent space. In this work, the latent space does not use a vectorized representation but rather a set-based representation with a small bottleneck. The authors then compare their methods to baselines and ablations on a toy environment.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I want to thank the authors for their work. The paper is rather clear, well structured and easy to follow. However, there are many abbreviations in the experimental part that I believe may be unusual for most readers, which makes it hard to follow.\nI think that this research domain is very promising. Indeed, MBRL methods recently demonstrated impressive results and studying how to better structure latent spaces is an exciting research direction. However, this paper exhibits many limitations: \n\nThis method falls in the category of deep MBRL methods, however it lacks many references to modern state-of-the-art methods in the field. Notably, the authors mention Line 30 that except for MuZero, MBRL methods show poor performance. They also argue that only Predictron and MuZero plan over a latent space while the others methods construct a model over the true state space. I think that the authors missed many recent and important references. Notably, methods such as PETS, MBPO, Dreamer or LEAP demonstrated impressive results in terms of both sample efficiency and asymptotic performance on challenging benchmarks. Among these methods both LEAP and Dreamer plan over a latent space. Furthermore, Dreamer also introduced set-based representation. Thus, it would be key in this paper to cite Dreamer and to carefully position the proposed method with respect to it.\n\nAs this paper does not provide theoretical results or analysis, the experimental section is key to demonstrate the usefulness of the method. However, this section is in my sense too weak to achieve this goal. Indeed, the authors considered a single toy environment that does not exhibit any of the challenges that the deep RL and deep MBRL have been designed to address. I would highly recommend the authors to assess the performance of their method on standard benchmarks such as the Atari suite or the Gym Mujoco environments. In addition, the authors do not compare their methods to classical SOTA methods such as the one aforementioned. The authors should compare their algorithm to at least Dreamer.\n\nRegarding the paragraph between lines 72 and 83, the authors missed the reference to Hamrick et al. (on the role of planning in model-based deep reinforcement learning) which I believe is important to mention here.\n\nThe conclusions in Section 5.5 are too broad and in my opinion, not enough justified. For instance, the authors claim that \"Model-free methods face difficulties in OOD generalization;\" where Model-free methods refer to double DQN and this method has been tested on a single toy environment. There have been tremendous efforts to study and improve the generalization of model free methods and I would recommend the authors either moderate their conclusions or to better inform them with more experiments and relevant citations.\n\nRegarding the notations, it is misleading to use $\\mathcal{S}$ for the latent state space and $\\hat{s}$ for latent states. Most of the time, $\\hat{s}$ is used for reconstructed states. I would recommend another notation such as $z$ for latent variables. Also, when the authors introduce the action space line 44, I would recommend to mention whether these actions are assumed to be discrete or continuous.\nTo conclude, I think that this research direction is promising, however, this paper lacks many important references and exhibits too weak an experimental study. I would recommend the authors strengthen their positioning with respect to the existing literature and either bring theoretical elements or consider more challenging environments and baselines in their study.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This paper does not add limitations or potential negative societal impact to existing reinforcement learning methods.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "Pgh3Gbl6rFb", "reviewer": "Reviewer_dyBZ", "summary": "This work presents a model-based reinforcement learning agent that makes use of a bottleneck attention mechanism for the planning module.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposes the use of attention-based bottleneck to improve the performance of model-based RL (MBRL) agents, which is intuitive and well explained. This mechanism seems similar to the top-k attention mechanism presented with BRIMs [1] also targeting OOD generalisation. But to the best of my knowledge, this is the first work that assess the impact of this mechanism in MBRL. Authors also study the impact of using a vector embedding vs am array of vectors (set embedding) which is also interesting. I think these results would be of interest for the RL community.\n\nClarity and quality is good in general, authors should be commended for the abundant graphs and visual explanations. My comments on this side focus on a few points easily fixable that I raise below. My biggest concern, and the reason why I think the paper cannot be accepted in its current form, are the claims that authors do of the conclusions extracted from their work, which are too generic and broad for the work presented and do not reflect correctly what has been presented:\n\n*Line 304 says that authors have drawn the conclusion that set-based representations are better in multitask environments. I think that if authors want to state this they should use more than one benchmark, currently all the experiments concern the same task which is \"reach  green cell while avoiding lava cells\"\n\n*Line 307 \"model-free methods face difficulties in OOD generalisation\". Again too broad when tested on a grid world with a single type of task\n\n*Line 309 states that \"online joint learning is good for RL\" has the same problem and, moreover this one is a controversial statement since there are evidences in the opposite direction (e.g., [2]), again I think authors cannot do such broad claims with the experiments included.\n\nTo clarify, I am not asking the authors to do more experiments to consider accepting this paper, I believe the proposed approach and the experiments included are enough to show the potential of the proposed mechanisms, I am ok if authors just adjust their claims to do more realistic statements according to what they present.\n\nAdditional comments:\n* Checklist: multiple answers should be expanded, the answer should include brief explanations, most are only \"yes\"\n* Figure 2 caption, it would be helpful to say which appendix section to go for.\n* The FC downscale operation, is barely explained, how much they downscale?\n*The FC layer in the permutation invariant receives as input an array of vectors, but the FC layer must operate with vectors, I understand that each vector is passed through the FC separately, isn't it?\n* Sections 5.3 and 5.4 should be merged.  Section 5.3 is supposed to be about in-distribution evaluation but the last lines are about OOD evaluation, which is section 5.4. I think it is good to have the text as it is now, thus, I would only put everything in the same subsection.\n* Lines 111-115. You cite twice \"[10]\" but seems tha tehy should be different references (one to refer what are you following and the other to refer to other families of approaches).\n*Caption Figure 10, you say \" Note that the concatenation does not happen outside the residual pass in case the dimensions do not match\" but from the picture is  looks like in the residual connection is where there is no concatenation of actions.\n\n--After Rebuttal---\nAuthors have correctly addressed my concerns about the claims and updated them accordingly. This, together with the additional experiments and updates makes me confident to raise my score and recommend the acceptance of this work.\n\n[1] Mittal, Sarthak, et al. \"Learning to combine top-down and bottom-up signals in recurrent neural networks with attention over modules.\" International Conference on Machine Learning. PMLR, 2020.\n[2] Lehuger, Auguste, and Matthew Crosby. \"Fixed $\\beta $-VAE Encoding for Curious Exploration in Complex 3D Environments.\" arXiv preprint arXiv:2105.08568 (2021).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Authors state limitations of their study, but make too broad claims for the experiments included as explained above.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "Orb0td1tzyb", "reviewer": "Reviewer_hG23", "summary": "The authors take inspiration from theories of human consciousness to construct a model-based architecture for RL, with the goal of generalizing more effectively.  They conduct experiments on a gridworld to show that their proposed architecture generalizes more effectively.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "After reading the authors' responses, the new results, and the other reviews, I believe that my concerns have been mostly addressed.  I am raising my score from a 5 to a 7.  (I am tempted to give an 8 or 9, as this paper has some fascinating and exceptionally well-motivated ideas. However, I cannot give an extremely high score to an empirical-only paper which tests only on gridworlds, despite the fact that there are good reasons why more complex environments are beyond the scope of this work and the fact that the authors did add a new gridworld variation.)  I hope the authors follow up this work by showing how to apply similar ideas to more complex environments.\n\n-----------------------\n\nSections 1-4 (relatively small concerns):\n\nThe encoder description in Section 3 (“We use the features...dynamics model, discussed below”) was insufficient and vague; I read it several times, but have no real idea how the encoder works.  A more precise definition of the encoder would strengthen this work (in the appendix if there’s no room in the main paper). Edit: while re-reading, I noticed that Figure 1 is describing this.  I think simply adding a reference to Figure 1 in the text (I don’t think there is any existing reference to Figure 1) might be sufficient to address most of this concern.  However, I’m still confused about how “This approach is different from the practice of adding positional information onto the features”, so clarifying some more in the appendix might still be helpful.\n\nTD loss: “In experiments, a distributional output is used for both value and reward estimation, making this loss a KL-divergence.” This is a little vague and confusing; I’d suggest defining it more precisely.\n\nSection 5 (significant concerns):\n\n-Low number of trials (5 runs): This is too low to make the claims that are made in Section 5.  RL algorithms are notoriously unstable between trials, and the differences shown could easily not appear in the true curves (that is, given an infinite number of trials).  I suggest 1) running MANY more trials, and plotting the standard error instead of standard deviation, to show that the difference is significant, and/or better yet, 2) do proper statistical hypothesis testing to show your conclusions are valid given the data.  For either approach, you will certainly need many more than 5 trials.  This concern is somewhat mitigated in 5.4.1 by having effectively 20 trials.  Still, it would be nice to see the 4 5-trial plots combined as a single plot with 20 trials and standard error bars, and/or see some statistical hypothesis testing (and far more than 20 trials would be better).\n\n-Gridworld: for an entirely empirical paper, only running gridworld experiments is not ideal.  The claim is “CP allows better generalization”, but at best (ignoring the concern above) what is shown is “CP allows for better generalization on 8x8 Gridworlds”, which is not a huge contribution.  More environments (including more difficult/complex ones) would help alleviate this problem.\n\nMinor edits:\n\n-Figure 1 caption, add an article adjective before CNN: “e.g., a CNN”\n\n-Too informal: “Luckily, our feature-position...”\n\nStrengths:\n\n-The plots of different bottleneck sizes are a great idea.  However, I’m not 100% satisfied; more sizes and more trials (runs) per size would be nice.\n\n-Good limitations and failed experiments sections in the appendix.\n\n-Many of the hyperparameter selections are justified in the appendix via references.  This is great and more papers should do this (otherwise the conclusions drawn may be based on the effects of the hyperparameter search results rather than the effects of the differences between the algorithm and baseline).\n\nQuestion/suggestion:\nAre the losses weighted when summed to produce the total loss?  If so, this information should be included (that is to say, the total loss equation is wrong).  If not, this is surprising; I’d suggest giving some intuition about why these losses are simply summed without any weighting and why that works.  Perhaps some empirical evidence in the ablation showing that the summed losses all tend to be a similar order of magnitude, and all help achieve good performance.  In other words, if you remove any of the losses, do the results get significantly worse?  Or are there one or two losses that are not affecting much?\n\nSummary:\n\nWhile the ideas the authors propose are fascinating and well-motivated, this paper has no theoretical contribution (which is fine for a strong empirical paper), and a somewhat weak empirical contribution.  Nonetheless, it is an interesting paper and has its strengths.  I recommend continuing to refine this work by strengthening the experimental results, as discussed in more detail above.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "~4.5"}, "extra_text": {}}, {"review_id": "5gX5UaV4JL", "reviewer": "Reviewer_orSE", "summary": "The authors introduce an architecture for focusing agents in model-based reinforcement learning on just the relevant subset of the environment representation for making decisions (reward, termination, and successor state prediction). The approach leans heavily on recent advances in set-based representations and the approach is tested in the MiniGrid-BabyAI framework. Results show promising behaviour, particularly in the ability to generalize across different settings.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Originality\nThe ideas in this paper are novel and interesting. The contrast to related work in the background section is one of the best I've ever read. It is immediately clear what novel aspects are part of this work and the motivation behind them.\n\n## Quality\nThe submission is technically sound, and the core hypothesis is adequately tested in the evaluation environment. The set of baselines are well-thought-out, and the evaluation seems complete.\n\nThe authors are quite honest and clear about the limitations, including the complexity of the hyperparameter space (which is a concern for models such as this).\n\nOne concern I have is in the presentation of some of the results. For Fig 7, it looks like the World Models approach is still warming up -- further training would quite possibly change the \"best approach\". What happens when you allow for a longer budget? I understand that there is a claim to data efficiency, but that is a weaker result than what is currently presented.\n\n\n## Clarity\nI found the paper to be extremely clear and filled with interesting insights -- both on the analysis in the evaluations and on the design decisions compared to related work. Nothing to suggest.\n\n\n## Significance\nThis is perhaps the biggest weakness of the paper. I find the ideas compelling, and the paper rides on that + the clarity alone, but the approach is tested in an arguably very simple setting. It is designed to exhibit the phenomena of interest, but there is little indication that the results will generalize to more complex settings. From the difficulty of jointly learning the representation+selection in the bottleneck, to the argument for ignoring reconstruction, there are many decisions that may only serve useful in this one particular environment.\n\n\n## Questions\nI am left with a few open questions after reading the work at various levels:\n\n1. The bottleneck still looks at everything in the memory in order to get the compressed latent state representation. This is arguably different from how we might do it cognitively (going over every fact to find those we might want to change). How might this be improved? It's not only about reasoning on a small set of facts, but avoiding the full $O(n)$ computation of what those $m << n$ facts would be.\n\n2. I am interested in the claim included in the first Background & Context section as to why reconstruction isn't useful. Would this still be the case if given the opportunity for unsupervised pre-training for a decent reconstruction?\n\n3. The \"integration\" phase appears to not be an inversion of the \"selection\" phase? If so, why? If they are indeed a mirror (i.e., the integration places the new representation back in the appropriate position), then the paper should clarify this.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The biggest limitation is the scope of the evaluation. The authors have identified this, along with other key limitations of the approach. Generally, the limitations of the work are carefully considered and discussed.", "review_text": "", "extra_scores": {"ethical_concerns": "I see no major ethical concerns with this body of work.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agreed that the paper contains compelling ideas and should be accepted.  Some initial concerns about the generality of the claims were well addressed by the authors.
Decision: Accept (Poster)