An Attentive Approach for Building Partial Reasoning Agents from Pixels
Authors: Safa Alver
Doina Precup
Venue: Accepted by TMLR ()

Abstract: We study the problem of building reasoning agents that are able to generalize in an effective manner. Towards this goal, we propose an end-to-end approach for building model-based reinforcement learning agents that dynamically focus their reasoning to the relevant aspects of the environment: after automatically identifying the distinct aspects of the environment, these agents dynamically filter out the relevant ones and then pass them to their simulator to perform partial reasoning. Unlike existing approaches, our approach works with pixel-based inputs and it allows for interpreting the focal points of the agent. Our quantitative analyses show that the proposed approach allows for effective generalization in high-dimensional domains with raw observational inputs. We also perform ablation analyses to validate of design choices. Finally, we demonstrate through qualitative analyses that our approach actually allows for building agents that focus their reasoning on the relevant aspects of the environment.

Published in Transactions on Machine Learning Research (08/2024)
An Attentive Approach for Building Partial Reasoning Agents from Pixels
Safa Alver Mila, McGill University
safo. refo.olver@mmii.mmglll
Doina Precup
dprecup@cs.megillccm
Mila, McGill Universiby and Google DeepMind
Reviewed on OpenReview: https: // openreview. net/ pormfo id- S3FUKFMR
Abstract
We study the problem of building reasoning agents that are able to generalize an effective manner. Towards this goal, we propose an end-to-end approach for building model-based reinforcement learning agents that dynamically focus their reasoning the relevant aspects of the environment: after automatically identifying the distinct aspects of the environment, these agents dynamically filter out the relevant ones and then pass them to their simulator to perform partial reasoning. Unlike exisling ing approaches, our approach works with pixel-based imputs and it allows for interpreting the focal points of the agent. Our quantitative analyses high-dimensional
show that the proposed approach allows for effective generalization domains with raw observational inputs. We also perform ablation analyses to validate our design choices. Finally, we demonstrate through qualitative analyses that our approach actually allows for building agents that focus their reasoning the relevant aspects of the environment.
1 Introduction
The goal in model-based reinforcement learning (RL) is to build reasoning agents that maximize long- term cumulative reward through trial and error interaction with the environment. In the last decade, the employment of deep neural networks as function approximators have allowed model-based RL agents to achiove significant successes across a wide range of areas, ranging from challenging board and video games (Sileer al., 2017b 2018 Ha Schmidluber] 2018 Halner et al.] 2020 2021 2023 Schriltwieser el. al. 2020) to discovering efficient algorithms (Fawzi et al. 2022 Mandhaan et al. 2022 Mankowitz al.) [2023). However, despite these successes, these agents typically designed to improve sample efficiency in regular RL settings and thus lack the inductive biases that are necessary for achieving effective performance in generalization settings.
One of the promising lines of research for tackling this shortcoming considers the use of partiol models in the
reasoning process of the agent (see e.g. Talvitie & Singh 2008 Khetarpal et al. 2020 2021 Zhao et al. 2021 Alver & Precup 2023). Inspired by studies cognitive science, these approaches aim for building agents that mimie the reasoning process in humans, in which planning hypothesized
be performed over few abstract aspects of the environment (Bengio 2017 Goyal & Bengio 2022 e.g. when planning route from point A to point B humans typically focus their reasoning to the points in between and discard the rest. More specifically, rather than building plain reasoning agents that plan over every aspect of the environment, these approaches advocate for building partial reasoning agents that focus the planning on the relevant aspects of the environment. Among these studies, the study of Zhao et al. (2021) has empirically demonstrated that this type of planning allows for effective generalization scenarios where the details of the environment keep changing but the overall task remains the same. However, this study has two main shortcomings: in providing approach for building partial reasoning agents that works with pixel-based inputs, i.e.
(i) it fails

---

Published in Transactions on Machine Learning Research (08/2024)
Aspect Identification Module
S36tt Attention Nodule
Aspect Jderthhing
Stots
Beeraive Attentionn
Encoder hetwork
Decoder
GRJ
Netwwrkk
X
Inpuc Observation 9t
Recon. Observation
Permutation Envariant Flltration Module
3
Valle-Equlvalent Simulation Module
Fooe Pontl
Encoding
Amtoniinn
Figure 1: Our overall architecture for building partial reasoning agents. This architecture consists of three main modules: (Box 1) the aspect identification module, (Box 2) the commutation-ivvaiian filtration module, and (Box 3) the value-equivalent simulation module. The details of these modules are provided in the main text. The q, & and blocks represent a linear mapping to obtain the query, key and value matrices, respectively, and the attention blocks represent a soft attention operation. The GRU block represents a GRU cell. Lastly, the black arrows indicate the direction of information flow. More details on this architecture can be found in App. A
approach that works beyond low-dimensional environments in which the distinct aspects of the environment is hand-provided via symbolic inputs, and (ii) it lacks detailed qualitative analyses demonstrating that their proposed approach actually allows for building agents that focus their reasoning on the relevant aspects of the environment.
In this paper, we overcome these shortcomings by proposing an approach for building partial reasoning agents that works with pixel-based inputs and allows for interpreting the focal points of the agent. Unlike prior work, our end-to-end approach allows for building reasoning agents that automatically identify the distinct aspects of the environment and then dynamically altend to the relevant ones, which makes possible for them to work on high-dimensional domains with raw observational inputs. Importantly, the built agents also allow for further inspection what they actually incorporate into their model throughout the course of interaction with the environment. After laying out its details (Sec. B). we demonstrate through quantitative analyses that our approach allows for effective generalization in high-dimensional domains (Sec. 4.3). Tn addition to the quantitative analyses, we also perform ablation analyses to validate our design choices (Sec. 4.4). Finally, we perform several qualitative analyses with our approach and show that actually allows for building agents that dynamically focus their reasoning on the relevant aspects of the environment, (Sec. 4.5)
We hope that our study will bring the RL community step closer to building scalable and interpretable reasoning agents that are able to effectively generalize to novel situations throughout their interaction with the environment.
Key Contributions. The key contributions of this study are as follows: (i) We propose an approach for building partial reasoning agents that works with pixel-based inputs and allows for interpreting the focal points of the agent. (ii) We demonstrate, through quantitative analyses, that our end-to-end approach allows for effective generalization high-dimensional domains with raw observational inputs. (iii) We demonstrate, through qualitative analyses, that our approach allows for building agents that actually focus their reasoning the relevant aspects of the environment.
2

---

Published in Transactions on Machine Learning Research (08/2024)
2 Background
Reinforcement Learning. In RL (Sutton & Barto 2018]; an agent interacts with its environment through sequence of actions to maximiie its long-termn cutnulative reward. The environment is usually modeled as a Markov decision process (MDP) M = (S, A, R, do, 7) where S and A are the (finite) set of states and actions, P S x A + Dist(S) is the transition distribution, R : S x A S + IR is the reward function, dg S + (S) is the initial state distribution and 7 E [0,1) is the discount factor. At each time step 1, after taking an action at E A, the environment's state transitions from St E S to 81+1 E S; and the agent receives an observation 01+1 O (reflecting S(+1) and an immediate reward rr. The goal of the agent is to learn a policy : O + (ist(AA) that maximizes E. 00 EE 7 R(S,, AG, 11+1)S5+ so dol, where Ex,pl denotes the expectation over trajectories induced by and P.
Model-Based RL. One way of achieving this goal is through model-based RL (Moerland et al. 2023) in which there are two alternating phases: (i) the learning and (ii) planning phases. In the learning phase, the gathered experience is mainly used in learning a model m, and in the planning phase, the learned model m is then used for simulating experience either to be used alongside real experience improving the value predictions or just to be used in selecting actions at decision time dAlver & Precup 2024).
The Attention Mechanism. In its general form, the soft attention mechanism (Vaswani et al. 2017 is described as y (QK JV where Q E Rmdd, K E Raxd, and V E Raxd, are the query, key and value matrices, respectively, and a(-) is the softmax function. Here, m and n denote the number of query and key / value vectors, respectively; and, dg and d, denote the dimensionality of the query / key and value vectors. More explicitly, the attention output y is computed in three steps: (i) first, each of the m do- dimensional query vectors are dotted with a database of n dy-dimensional key vectors to compute the score values, (fi) then these values are passed through a softmax function o.),, along the axis of the key vectors, to compute the attention weights, and (iii finally, the output y is computed by taking a weighted sum of the n dimdiimnsil value vectors with the weights
Self- Attention and Permutation-Ivvaiiance. A particular form of attention that is of interest in this study is self-attention. In self-attention, the Q, K, V matrices are a function the input x, e.g., Q = q(z), K = k(x) and V = =(x) where q(-); k(-) and v(-) are linear mappings. By default, the self-altention operation is not cermutation-invariin
i.e., permuting the input also results in a permutation in the output. However, clever way to enable cetuuttttioonivarinnee in self-attention is to make Q a function of set of fixed
positional embeddings, rather than the input (see Lee et al. 2019 Tang & Ha. 2021
3 An Attentive Approach for Building Partial Reasoning Agents from Pixels
In this section, we present our approach for building partial reasoning agents that works with pixel- inputs and allows for interpreting the focal points of the agent. Our approach consists of an end-to-end architecture and an accompanying training procedure. The architecture, depicted in Fig. D consists of three separate modules: (i) the aspect identification module, (ii) the permutation-inwee invariint filtration module and (iii) the value-equivalent simulation module. The details of these modules and the procedure for collectively training them are presented in the following sections.
3.1 Aspect Identification Module
At a high level, at each time step t in the environment, the aspect identification module (see Box in Fig. I takes in the raw observational input o, and binds its distinct aspects into a set of n distinct slots ( The main function of this module is to automatically identify the distinct aspects of the environment, c.g. the distinct objects or the distinct regions in the observational input. Internally, this module resembles an autoencoder architecture [Goodfellow et al. 2016) and it consists of (i) an encoder network, (ii) a slot attention module and (iii) a decoder network.
We prefer slot attention over other tochniques it allows for better generalization [ocatello al. 2020 We also note that even though Locatello et 2020) introduced slot attention and used it for the purposes of object discovery and prediction in the context of computer vision, in this study, we are using slot attention for the purpose building partial reasoning agents.
3

---

Published in Transactions on Machine Learning Research (08/2024)
The encoder network (see Box La in Fig. 1 takes as imput the raw observation or and encodes it into vectorized representations Se that are augmented with positional embeddings, i.e. ft = e(ot). Tt follows a simple structure: a convolutional backbone followed by a positional embedding layer.
The vectorized representations fi are then passed through the slot attention module (Locatello et al. 2020 to obtain a set of aspect identifying slots (st)==, which form a set with permutation symmetry. Each slot in the set can, for example, describe an object or certain region in the observational input. This module uses iterative attention mechanism to produce the set of slots. More specifically, slots are initialized random and thereafter refined for T iterations with Gated Recurrent Unit (GRU, Cho et al. 2014) to bind to particular grouping of the vectorized reprosentations ft (here each iteration refers to single recurrence of the GRU). At each iteration, slots "compete" for the reconstructing parts of the input via a soft attention mechanism (Vaswani et al.) 2017). The internal details of this module are depicted in Box 1.b of Fig. I For
more details, we refer the reader to the study of Locatello et al. (2020)
Finally, the aspect identifying slots {s are then passed through the decoder network d (see Box 1.c in Fig. [1 to obtain the observation reconstruction oz, i.e. 6, = d({s)]= The decoder network is spatial broadcast decoder (Watters et al. 2019), as in (Greff et al., 2019): each of the aspect identifying slots are broadcasted into a 2-dimensional grid and augmented with positional embeddings. Each grid is decoded using a convolutional network to produce an output of size W x H x 4, where W and H are the height and width of the image, respectively. The output channels encode the RGB color channels and an unnormalized alpha mask. The alpha masks are then normalized using a softmax across the slots and used as mixture weights to combine the individual reconstructions into a single RGB image, which we refer to as ò
For more on the details of aspect identification module, we refer the reader to App. A.1
3.2 Permutation-Invarian Filtration Module
The permutation-invariie (PI) filtration module (see Box 2 in Fig. 1 takes as input the aspect identifying slots {s{)==,, i.c. the output of the slot attention module, and performs "soft" filtration operation over these slots to distill out the relevant information. The main purpose of this module is to filter out the relevant information from the slots that are likely to be useful in pocision-making which is later to be used for partial reasoning. Here, it is important to note that the filtration is dynamic, iee. throughout the agent-environment interaction, the PI filtration module can filter out different slots at different time steps. Internally, the slots are weighted by a PT soft attention mechanism (Lee et al. 2019). Here, the use of a PT attention mechanism is critical, as the aspect identifying slots form a set with permutation symmetry, i.e. the order of them can change at each time step. More on the details of this module can be found in App. A.2
It is also important to note that the PT filtration module has some analogues to the "consciousiees-inspieee bottleneck" that was proposed in Zhao et al. (2021). ). However, there is an important difference: rather than using a top-k semi-hard attention mechanism Ke et al.] 2018), we use a relatively straightforward soft attention mechanism (Vaswani et al., 2017]. This difference makes our approach much simpler and it also mitigates the burden of tuning for the optimal k value. Experiments in the following sections also demonstrate that the use of top-k semi-hard attention does not bring any performance advantages.
3.3 Value-Equivalent Simulation Module
The value-equivalent simulation module (sce Box 3 Fig. 1 takes the filtered aspect identifying slots as input and uses them as the root state of for performing value-equivalent simulation (Grimm al.] 2020; 2021 Schrittwieser et al., 2020 The main function of this module is to perform partial reasoning, i.e. to perform planning with the relevant aspects of the environment. More specifically, at each time step t, for cach of the hypothetical & = 1,..., K steps, the simulator takes in the root state xt and predicts three future quantities: (i) the policy pt R 101, Ot, O++1,..., Eluc+++1 + uc+k+2 + ... ....oo...., ot, 0.+1. dc+k], and (iii) the immediate reward of ze Ui+k where u is the true reward, æ is the policy used to select real actions, and y is the discount factor. Internally, the simulator is represented by combination of a dynamics and prediction network. The dynamics network recurrently predicts, at each hypothetical step k, an internal state of and an immediate reward r ie.
(ii) the value function of v
4

---

Published in Transactions on Machine Learning Research (08/2024)
of,rī m(ata ,af). (at). The policy and value function then computed from the internal state ot by prediction network i.e. pt. or = g We refer the reader to App. A.3 for more details on this module.
Given such a simulator, it is possible to do planning over the internal state space and reward induced by the dynamics network. Specifically, in this module, we use the Monte-Carlo tree search (MCTS, Coulom 2006) algorithm to perform the planning- At each time step I, this algorithm outputs a recommended policy Rt from which an action 0r+1 ~ M. is then selected.
Finally, it is important to note that, in this module, have preferred value-equivalent simulation (see e.g. Tamar et al. 2016 Oh et al. 2017 Silver et al., 2017a Schrittwieser et al. 2020) over other reconstruction-
based simulation methods (see e.g. Watter et al, 2015 Wahlström et al. 2015 Kaiser et al. 2019 Hafner
et al. 2020 2021
2023) as former simulation hots directly oplimize for learning representation ations that are useful for achieving the task of interest, which also aligns with the main objective of this study.
3.4 Training Procedure
Using batches of data sampled from the replay buffer, the overall architecture is jointly trained, in an end-to- end manner, with a loss function that is weighted sum of the reconstruction loss Lrecon and the simulation loss Cxim'
Liotal = a.rrecon (O; 01) + B Σ (p(4Ho1) + Loo + C,
(1)
Lerm
where Cp. C.. and C, are the loss functions for the policy, value function and reward, respectively, constituting the simulation loss Csim- While Lreeon optimizes for minimizing the error between input observation and the reconstructed observation ô,, Cp. L, and L, optimize for minimizing (i) the error between the search policy Tt+* and predicted policy pt, (ii) the error between the improved value target 2++A = #4+1 + yu4+2 + ... and predicted reward r respectively. In E,, the improved value targets 2444 are generated by collecting trajectories in the environment, and in Cp, the improved policy targets Te+* are generated by performing MCTS. More details individual losses and the weighting between them be found in App. B
+ +-"--uu4 + Ut4n and the predicted value uf, and (iii) the error between the observed reward We+s
4 Experimental Results
4.1 Environment Descriptions and Evaluation Setting
Environment Descriptions. We perform our experiments on the (i) MiniGrid (Chevalier-Boisvert. et al., 2018) and (ii) Procgen domains Cobbe al. 2020). We choose these domains as the former al- lows for designing controlled experiments and the latter is helpful in demonstrating the scalability of our approach to more complex do- mains. Here, it is important to note that these domains the agent only receives raw observational inputs, i.e. 64x64 RGB images, and not symbolic inputs as in Zhao et al. 22021
DynObs-Pass
DynObs-Collen
CoinRun
BigFish
In the MiniGrid domain, consider two customized environments: (i) DynamicObsttcclssPaas and (ii) DynamicObstacles- Collect (abbre- viated as BynObs-Pass and DynObs- DynObs-Collect, respectively, see the top row of Fig. [ In both of these environments, the agent (red square) is randomly spawned at the leftmost column the grid and has to pass through a region with randomly oscillating obstacles to reach the goal (green square), which is randomly placod the rightmost column of the grid. At each time step the agent is able to move to one of its
Figure Sample frames from the (top row) DynObs-Pass DynObs- Collect, (bottom row) CoinRun and BigFish environments.
Zhao al. 2021 perform experiments on customized MiniGrid environment with symbolic inputs. The symbolic input the form triple (xt, idz), where and y are the coordinates on the grid, and ids the id of the object
5

---

Published in Transactions on Machine Learning Research (08/2024)
neighboring cells. After every episode, a new game level is generated by randomly sampling obstacles among the orange, cyan, yellow, pink and purple obstacles and then by randomly placing them inside the region between the agent and the goal. In DynnOs--Pass, the agent has to complete the level by dodging the existing doadly orange, cyan, pink and purple obstacles (colliding with the existing yellow obstacle has no effect). In the agent has to complete the level by collecting the existing yellow obstacle and dodging the existing deadly orange, cyan, pink and purple ones. In both of these environments, if the agent suc- cessfully completes the level, it receives a reward of +1 and the episode terminates; otherwise the episode terminates with total reward of 0.
In the Procgen domain, we consider two commonly-used environments: (i) CoinRun and (ii) BigFish (sce the bottom row of Fig. 2). In CoinRun, the agent (an astronaut) must collect the coin that is at the far right of the level. In BigFish, the agent (a green fish) starts as a mid-sized fish and has to grow by eating the fish that are smaller than itself. In both of these environments, a new game level procedurally generated after every episode. For more details on these environments, we refer the reader to the study of Cobbe et al. 2020
consider the generalization setting, in which the agent's 's performance is evaluated by how it performs on a set of test levels (test set) as it gets trained on a set of training levels (training set) of the same environment. Importantly, in this setting, the agent does not get to interact with the test levels, ie. evaluation is performed in a zero-shot manner.
Evaluation Setting. Following Zhao et al. 2021
4.2 Agents
In our experiments, we compare the following agents:
Partial Reasoning (PR) Agent. A partial reasoning agent that was built by our proposed approach in Sec. 8
MuZero (MZ) Agent. A MuZero (Schrittwieser et al., 2020) agent. Unlike the PR agent, the MZ agent performs regular reasoning. The main difference between the MZ agent and the PR agent is in the procodure for constructing the root state xy: while the PR agent constructs it by the use of an aspect identification and PI filtration module (Sec. 3), the MZ agent constructs it by the use of a plain convolutional and residual architecture (He et al. 2016). We use this agent as a baseline to demonstrate the impact of partial reasoning in achicving better generalization.
Model-Free (MF) Agent. A version of the MZ agent that does not perform any planning. This agent takes actions by skipping the search process and just following its value predictions (which makes it a model-free agent). We use this agent as a baseline to demonstrate the impact of planning in achieving better generalization.
PR-NoPIF Agent. A version of the PR agent with no PT filtration module.
PR-SH Agent. A version of the PR agent in which the PI filtration module is implemented with top-k semi-hard attention. The purposes of this agent and the PR-NoPIF agent will be unfolded in the ablation analyses section (Sec. 4.4
We note that the PR agent and the baseline agents share architectures as much as possible ensure fair comparisons. More details on all of the agents can be found in App. C
4.3 Quantitative Analyses
Our approach allows for effective generalization in high-dimensional domains. The study of Zhao et al. (2021) demonstrated that partial reasoning allow for effective generalization over regular reasoning. It also demonstrated that planning plays important role for better generalization. However, these demonstrations were mainly performed on low-dimensional gridworlds with symbolic inputs. One of the main goals of our approach is lo generalize these successes to high-dimensionall domains. Towards this goal, we compare the generalization performances of the PR, MZ and MF' agents across two image-based MiniGrid and Procgen domains (Sec. 4.1). Performance curves in Fig. consistently show that, the PR agent displays a better generalization performance than the MZ agent, both (i) highlighting the role of partial reasoning in
achieving effective generalization and (ii) corroborating the findings of Zhao et al. (0021. I. We can also see
6

---

Published in Transactions on Machine Learning Research (08/2024)
1.0
1.0
1.0
1.0
0.8
0.8
0.8
0.8
0.6
0.6
MF
NZ
0.4
Mi-kofff
Mu-SH
Timestops
Timestops
1e6
Timestops
166
Timestaps
(a) DynObs PasPass (Train)
(b) DynObs-Paar Pass (Test)
(c) DynObs-Collect (Train) (d) DynObs-Collect (Test)
15
5oore
Score
2-
0
5
2
3
2
3
4
1
2
165
106
Timastaps
166
166
(0) CoinRum (Train)
(I) CoinRun (Test)
(4) BigFish (Train)
(h) BigFiih (Test)
Figure 3: The generalization performances of the PR. and baseline agents on the (a, b) DynObs-Pass, (c, d) DynObs-Collect, (c, f) CoinRun and (g. h) BigFish environments. Here, the plots are obtained by evaluating the agent on a set of test levels (referred to as Test) as it gets trained on a set training levels (referred as Train) of the same environment. In MiniGrid and Procgen experiments, the training sets consist of 16 and 500 randomly-sampled game levels, respectively. In both experiments, the test sets consist of all of the possible game levels. The mcans and the confidence intervals are computed over 15 independent runs.
that the MF agent underperforms the PR and MZ agents, which validates the role of planning in achieving better generalization.
4.4 Ablation Analyses
The PI filtration module plays a critical role in the performance of our approach. One of the crucial components of our proposed agent architecture is the Pl filtration module (Sec. 3.2). This module is particularly important as it acts as a "soft" filter to what the agent would be modeling. To quantify its importance, we compare the PR agent with the PR-NoPIF agent (Sec. 4.2]. Recall that in the PR-NoPIF agent the aspect identifying slots are directly passed, without any filtration, to the value-equivalent simulation module. Results in Fig. 3 show that the removal of the PI filtration module significantly deteriorates the generalization performance of the PR agent, making it on par with the MZ agent. These results indicate that the PT filtration module indeed plays critical role in the performance of our approach.
The use of top-k semi-hard attention in the PI filtration module does not bring any per- formance advantages. As expressed previously, in this study, we have implemented the PR agent's PI filtration module with soft attention, in which the attention weights can take any value between zero and onc. However, another way to implement this module could have been with top-k semi-hard attention, which was used in the "consciousness iessspppu bottleneck" of Zhao et al. (2021) (Sec. 3.2) In top-k semi-hard attention, only the top-k attention weights kept and all others this type of attention would allow for a hard filtration of only the k aspect identifying slots with the largest atlention weights. To investigate the impact of this type of attention, we compare the PR agent with the PR-SH agent (Sec. [.2). The comparison results, shown Fig. 3 indicate that the use of top-k semi-hard attention does not bring any advantages in the generalization performance of the PR agent; fact, the use of slightly decreases the performance. These results suggest that the use of the more simple soft attention mechanism be better choice in the implementation of the PI filtration module.
set to zero. If used in the PI filtration module,
7

---

Published in Transactions on Machine Learning Research (08/2024)
Observation
Slot 1
Slot 2
Slot
Slot
Slot 5
Slat 6
Slot 7
Slot
Recon.
DynObs
CoinRun
BigFish
Figure 4: (First column) The observational inputs to the aspect identification module across different envi- ronments. (Middle columns) The correspouding masked reconstructions of each of the aspect identifying slots. The visualization are obtained by multiplying the alpha masks and reconstructions of each slot, which are both obtained by passing the slots through the decoder network. We note that while the aspect iden- tifying slots bind to different objects in the DynObs-Pann bind to different regions the CoinRun environment. (Last column) The corresponding reconstructions of the aspect identificatinn module.
JynObs-Collees and BigFish environments, they
4.5 Qualitative Analyses
The study of ZZhao et al. (222)) claimed the learning of partial reasoning agents, but lacked in providing detailed qualitative analyses that demonstrates that the learned agents actually focus their reasoning on the relevant aspects of the environment. In this section, we perform qualitative analyses with the PR agent and demonstrate that our approach actually allows for building reasoning agents that focus on the relevant aspects of the environment.
Our approach allows for building reasoning agents that automatically identify the distinct aspects of the environment. Before the dynamic attention illustrations, we first start by presenting visualizations to demonstrate that the PR agent is actually able to automatically identify the distinct aspects of the environment. Recall from Sec. 3.1 that the automatic identification of the distinct aspects falls under the role of the aspect identification module. For this purpose, Fig. f we visualize the masked reconstructions of each of the individual aspect identifying slots across different environments. Inspecting that the slots indeed bind to the distinct aspects of the environment, corresponding either to different objects or to different regions in the observational input. This demonstrates that our proposed approach is indeed succossful in automatically identifying the distinct aspects of the environment.
this figure,
Our approach allows for building reasoning agents that dynamically attend to the relevant aspects of the environment. We now present illustrations to demonstrate that the PR agent actually able to dynamically attend to the relevant aspects of the environment. Rocall that the dynamic attention the relevant aspects falls under the role of the PI filtration module (Sec. 3.2 To see the PI filtration module is actually successful in this task, in Fig. 5 present the attention maps of the PR agent throughout its course of interaction with different enviromments. Examining this figure, module indeed dynamically attends to the relevant aspects of different environments (see the caption of Fig. for the interpretations), which demonstrates that our proposed approach indeed allows for building reasoning agents that dynamically focus on the relevant aspects of the environment.
see that the Pl filtration
5 Related Work
Partial Reasoning Agents. The use partial models for building partial reasoning agents has previously been investigated in the model-based RL literature TTavvitie & Singh| 2008 Khetarpal et al.] 2020 2021
Zhao al. 2021 Alver Precup. 2023 Alver al. (024).. Among these studies, the study of Zhao et al. (2021) is closest. to work in that it proposes approach for building partial reasoning agents and

---

Published in Transactions on Machine Learning Research (08/2024)
DynObs Callect
ComnRun
BigFith
Figure Example illustrations of how the attention maps of the PR agent evolve interacts with different environments. The attention maps are obtained by overlapping the weighted alpha masks of the aspect identifying slots, which are obtained by weighting the alpha masks with the attention weights from the PI filtration module. In the attention maps, the yellow and dark blue regions indicate the aspects with the highest and lowest attention, respectively. (DynObs-Collect) In the DynObs-Collect illustrations, we can see that the agent initially focuses on the yellow obstacle, while also taking into account the deadly obstacles around it, and after the collection of this obstacle it completely focuses on the goal. (CoinRun) In CoinRun illustrations, we can see that the agent constantly focuses on where will land at the end of its jump while moving to the far right of the level. (BigFish) In BigFish illustrations, we can see that the agent initially focuses the red fish next to it, and then after cating this fish it focuses the pink fish in front of it.
demonstrates through quantitative analyses that their approach allows for effective generalization. However, our work differs in that we propose an approach in which the agent automatically identifies, in an end-to-end manner, the distinct aspects the environment and then dynamically attends to the relevant oncs, which makes it possible to work on domains beyond low-dimensional gridworlds which the distinct aspects of the environment hand-provided via symbolic inputs, i.e. high-dimensional domains with raw observational inputs. Another limitation of Zhao et al. (2021) is that they lack detailed analyses demonstrating that their proposed approach actually allows for building agents that focus the relevant aspects of the environment. Our work overcomes this by presenting qualitative analyses in which the built. partial reasoning agents dynamically attend to the relevant aspects of the environment.
Our work also related to the study Alver & Precup (023) in that our partial reasoning agent is also built towards achieving effective generalization. However, rather than proposing an approach for building black-box agents that only act as partial reasoning agents the behavior level, our study proposes an approach for actually building partial reasoning agents, i.e. as opposed to Alver Precup (2023). in our study propose interpretable architecture that allows for actually knowing that the agent is dynamically attending to the relevant aspects of the environment.
Attention in RL. Attention mechanisms have found many uses in RL (see e.g., Sorokin et al. 2015 Zambaldi et al. 2018 Mott et al. 2019 Tang et al.) 2020: Liu et al.] 2020 Tang & Ha 2021). Different from these studics, which use attention the context of model-free RI, our work here specifically uses attention build partial reasoning agents. Specifically, we utilize attention for the purposes of aspect identification and (soft.) filtration in high-dimensional domains.
9

---

Published in Transactions on Machine Learning Research (08/2024)
SaalueEquiivleet Planning. A recent trend in model-based RL is to learn models that directly optimize for value-based planning, which has been studied under the name of value-equivalent planning (see e.g.. Silver et al.] 2017a Farquhar et al.) 2017] Oh et al. 2017; Schrittwieser et al. 2020: Grimm et al. 2020 2021j. Our work also advocates the idea that models should optimize for value-based planning and uses a value-equivalent simulation module. However, our work differs that we are also explicit in semantics of the representations, i.c. we restrict them to be the different aspects of the environment, which allows for better representation learning. We also demonstrate that this explicitness can provide signifieant advantages in generalization.
Cognitive Science. Performing partial reasoning also has connections to the conscious planning literature in cognitive science (see e.g. Baars 1993 2002: Dehaene et al.) 2021 Van Gulick, 2022 Bengio 2017 Goyal
& Bengio 2022p. In conscious planning, it is hypothesized that humans perform planning with a few abstract
entities that are obtained by attending to the relevant aspects of the environment. This type of planning has
been theorized to allow for efficient generalization (Bengio) 2017 Goyal & Bengio 2022). Similar to Zhao
jt al. 2021), , our work can also be viewed study that. incorporates some of the ideas from this literature to the model-based RL literature.
Interpretable RL. Our work is also related to the line of research that aims for building interpretable RL agents (Jiang & Luo 2019 Druce et al.] 2021 Glanois et al. 2021 Di Langosco et al. 2022 Delsosse et al. 2023 2024a b Bastani et al. 2018 Kohler et al. 2024 Luo et al. 2024]. However, rather than leveraging
interpretable architectures in salety-critical applications and for the main purpose of understanding the behavior of the agent, our study leverages them in a complementary fashion: for the purpose of shedding light the internal working mechanisms of our partial reasoning agent and in understanding that it actually focuses the relevant aspects of the environment.
6 Conclusion
To summarize, in this study, we have presented a novel approach for building partial reasoning agents which are agents that dynamically focus their reasoning on the relevant aspects of the environment. Unlike existing approaches, our approach works with pixel-based inputs and allows for interpreting the focal points of the agent. Our empirical results (i) demonstrate that the proposed approach allows for effective generalization in high-dimensional domains with raw observational inputs, and (ii) show that it allows for building reasoning agents that indeed dynamically focus their reasoning the relevant aspects of the environment throughout the agent-cnvironmen interaction. We believe that our approach can be an important step towards building scalable and interpretable reasoning agents that are able to effectively generalize to novel sations. Before closing, it is important to note that even though we have made use of particular attentive architectures in the modules, our overall architecture, presented in Fig. [ is a general architecture and its modules can be implemented by using different attentive architectures. Finally, we note that in this study, similar to Zhao et al. (2021,, we have only considered environments in which the details keep changing but the overall task remains the same. In future work, hope to extend our work and tackle environments in which the overall task also changes over time.
Acknowledgments
This project has been partly funded by an NSERC Discovery grant and the Canada CIFAR AI Chair program. We would like to thank the anonymous reviewers for providing critical and constructive feedback.
References
Safa Alver and Doina Precup. Minimal value-equivalent partial models for scalable and robust planning in lifelong reinforcement learning. In Sarath Chandar, Razvan Pascamu, Hanic Sodghi, and Doina Procup (eds.); Proceedingss of The 2nd Conference on Lifelong Learning Agents, volume 232 of Proceedings of Machine Learning Research, pp. 548 567. PMLR, 22 25 Aug 2023.
Safa Alver and Doina Precup. look at value- based decision-time vs. background planning methods across different settings. arXiv preprint ar wrx.i:2200008/.2 2024.
10

---

Published in Transactions on Machine Learning Research (08/2024)
Safa Alver, Ali Rahimi-Kalahroudi, and Doina Precup. Partial models for building adaptive model-based reinforcement learning agents. arXiv preprinl or orXim:2455.68890 2024.
Bernard J Baars. A cognitive theory of consciousness. Cambridge University Press, 1993.
Bernard J Baars. The conscious access hypothesis: origins and recent evidence. Trends in cognitive sciences, 6(1):47-52, 2002.
Osbert Bastani, Yewen Pu, and Armando Solar-Lezama. Verifiable reinforcement learning via policy extrac- tion. Advances in neural information processing systems, 31, 2018.
Yoshua Bengio. The consciousness prior. arXiv preprint arXiv:1 709.0096...88 2017.
Maxime Chevalier-Boisvertt, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for openai gym. https://githuu: com/ htmm/maximecb///emmomm wwwpaaawiineiiiiaaaiieeeeeeee sym-miiiggid 2018.
Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rmn encoder-decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empérical Methods in Natural Language Processing (EMNLP), pp. 1724. Association for Computational Linguistics, 2014.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pp. 2048 405-- PMLR, 2020.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International con- ference computers and games, pp. 72-83. Springer, 2006.
Stanislas Dehacne, Hakwan Lau, and Sid Kouider. What is consciousness, and could machines have it? Robolics, AI, and Humanilyy: Science, Ethics, and Policy, pp. 43 56, 2021.
Quentin Delfossc, Jannis Blüml, Bjarne Gregori, Sebastian Sztwiertnia, and Kristian Kersting. Ocatari: Object-centric atari 2600 reinforcement learning environments. arXiv preprint orXiv:2306. 08649, 2023.
Quentin Delfosse, Hikaru Shindo, Devendra. Dhami, and Kristian Kersting. Interpretable and explainable logical policies via neurally guided symbolic abstraction. Advances in Neural Information Processing Systems, 36, 2024a.
Quentin Delfosse, Sebastian Sztwiertnia, Wolfgang Stammer, Mark Rothermel, and Kristian Kersting, In- terpretable concept bottlenecks to align reinforcement learning agents. arXiv preprinl orXiv2240100882 2024b.
Lauro Langosco Di Langosco, Jack Koch, Lee D Sharkey, Jacob Pfau, and David Krueger. Goal misgeneral- ization deep reinforcement learning. In International Conference on Machine Learning, pp. 12004-12019. PMLR, 2022.
Jeff Druce, Michael Harradon, and James Tittle. Explainable artificial intelligence (xai) for increasing user trust in deep reinforcement learning driven autonomous systems. orXiv preprint arXiv:2106.033 03777, 2021.
Gregory Farquhar, Tim Rocktäschel, Maximilian Tgl, and Shimon Whiteson. Treequ and atreec: Differen- tiable tree-structured models for deep reinforcement learning. arXie preprint arXiv:1710.11 11417, 2017.
Alhusscin Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J R Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, al. Dis- covering faster matrix multiplication algorithms with reinforcement learning. Nature, 610(7930):47-53, 2022.
Claire Glanois, Paul Weng, Matthieu Zimmer, Dong Li, Jianye Hao, Tianpei Yang, and Wulong Liu. A survey on interpretable reinforcement learning. IEEE Transactions on Neural Networks and Learning Systems, 2021.
11

---

Published in Transactions on Machine Learning Research (08/2024)
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. deeplearningbook. org.
MIT Press, 2016. http://www
Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. Proceedings 2022.
of the Royal Society A,
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Damiel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning, pp. 2424-2433 PMLR, 2019.
Christopher Grimm, Andre Barreto, Satinder Singh, and David Silver. The value equivalence princi- ple for model-based reinforcement learning. In Н. Larochelle, M. Ranzato, R. Hadsell, M.F. Bal- can, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 5541 5552. Curran Associates, Inc.. 2020. URL 3bbBBaa0001400eeeeeeeeeddda66333
pdf
Christopher Grimm, Andre Barreto, Gregory Farquhar, David Silver, and Satinder Singh. Proper value equivalence. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Newral Information Processing Systems, 2021. URL
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Proceedings of the 32nd International Conference Newral Information Processing Systems, pp. 2455-2466 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. Tn International Conference on Learning Representations, 2020. URL
https://openreviers mett/forumiid=s1100005
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis-
crete world models. In International Conference on Learning Representations, 2021. URL https:
//openreview.o www.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXi::2301.04 104, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pp. 630 645. Springer, 2016.
Zhengyao Jiang and Shan Luo. Neural logic reinforcement learning, In International conference on machine learning, pp. 3110-3119. PMLR, 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Rodel-based reinforcement learning for atari. XiXiv preprint Xiv: 1103..900.000 09.0000774........ 2019.
Nan Rosemary Ke, Anirudh Goyal, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse allentive backtracking: Temporal credit assignment through reminding. Advances in neural information processing systems, 31, 2018.
Khimya Khetarpal, Zafarali Ahmed, Gheorghe Comanici, David Abel, and Doina Precup. What can i do here? a theory affordances in reinforcement learning. In International Conference on Machine Learning, pp.
PMLR, 2020.
Khimya Khetarpal, Zafarali Ahmed, Gheorghe Comanici, and Doina Precup. Temporally abstract partial models. Advances in Neural Information Processing Systems, 34:1979-1991, 2021.
Hector Kohler, Quentin Delfosse, Riad Akrour, Kristian Kersting, and Philippe Preux. Interpretable and editable programmatic tree policies for reinforcement learning. arXiv preprint arXiv:2405. 14956, 2024.
12

---

Published in Transactions on Machine Learning Research (08/2024)
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-basel permutations ion-invariant neural networks. In International conference on machine learning, pp. 3744 3753. PMLR, 2019.
Tou-Jen Liu, Raymond A Ych, and Alexander Schwing. Pic: permutation invariant critic for multi-agent deep reinforcement learning. In Conference on Robot Learning, pp. 590 602. PMLR, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot allention. Advances in Neural Information Processing Systems, 33:1152; 11538,, 2020.
Lirui Luo, Guoxi Zhang, Hongming Xu, Yaodong Yang, Cong Fang, and Qing Li. Insight: End-to-end neuro- symbolic visual reinforcement learning with language explanations. arXiv preprint ar Xiv: 12451, 2024.
Amol Mandhane, Anton Zhernov, Maribeth Rauh, Chenjic Gu, Miaosen Wang, Flora Xue, Wendy Shang, Derek Pang, Rene Claus, Ching-Han Chiang, al. Muzero with self-competition for rate control vp9 video compression. arXie preprint Xix::2200 06626, 2022.
Daniel J Mankowitz, Andrea Michi, Anton Zhernov, Marco Gelmi, Marco Selvi, Cosmin Paduraru, Edouard Leurent, Shariq Iqbal, Jean-Baptiste Lespiau, Alex Ahern, et al. Faster sorting algorithms discovered using deep reinforcement. learning. Nature, 618(7964)2257 263, 2023.
Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based reinforcement learning: A survey. Foundations and Trendo0 in Machine Learning, 16(1):1-118, 2023.
Alexander Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, and Danilo Jimenez Rezende. Towards interpretable reinforcement. learning using attention augmented agents. Advances in neural information processing systems, 32, 2019.
Yazhe Niu, Yuan Pu, Zhenjic Yang, Xucyan Li, Tong Zhou, Jiyuan Ren, Shuai Hu, Hongsheng Li, and Yu Liu. Lightzero: A unified benchmark for monte carlo tree search in general sequential decision scenarios. Advances in Newral Information PacioonnPProeee Processing Systems, 36, 2024.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in neural information processing systems, 30, 2017.
Julian Schrittwieser, Toannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockbart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604--69 2020.
David Silver, Hado Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Arnola David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-to-end learning and planning. In International Conference on Machine Learning, pp. 3191 3199. PMLR, 2017a.
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676) 359, 2017b.
David Silver, Thomas Hubert, Julian Schrittwieser, Toannis Antonoglou, Malthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(64119:1140--1 2018.
Ivan Sorokin, Alexey Seleznev, Mikhail Pavlov, Aleksandr Fedorov, and Anastasiia Ignateva. Deep attention recurrent q-network. arXiv preprint arXiv:1512.01693, 2015.
Richard S Sutton and Andrew G Barlo. Reinforcement learning: An introduction. MIT press, 2018.
Erik Talvitic and Satinder Singh. Simple local models for complex dynamical systems. Advances in Neural Information Processing Systems, 21, 2008.
13

---

Published in Transactions on Machine Learning Research (08/2024)
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural informalion processing systems, 29, 2016.
Yujin Tang and David Ha. The sensory neuron as transformer: Permutation-imvaiian neural networks for reinforcement learning. Advances in Neural Information Processing Systems, 34:22574-22585 2021.
Yujin Tang, Duong Nguyen, and David Ha. Neuroevolution of self-interpretable agents. In Proceedings of the 2020 Genetic and Evolutionary Computation Conference, pp. 414 424, 2020.
Robert Van Gulick. Consciousness. In Edward N. Zalta and Uri Nodelman (eds.), The Stanford Encyclopedia of Philosophy. Metaphysics Research Lab, Stanford University, Winter 2022 edition, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30,
2017.
Niklas rählström, Thomas B Schön, and Marc Peter Deisenroth. From pixels torques: Policy learning with deep dynamical models. arXiv preprint arXiv:102.0025 022011 2015.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
Nicholas Watters, Loic Matthey, Christopher P Burgess, and Alexander Lerchner. Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes. aiv preprint 0rXis:1901.07017 2019.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al. Deep reinforcement learning with relational inductive biases. In International conference on learning representations, 2018.
Mingde Zhao, Zhen Liu, Sitao Luan, Shuyuan Zhang, Doina Prccup, and Yoshua Bengio. A consciousness- inspired plamning agent for model-based reinforcement learning. Advances in Neural Information Process- ing Systems, 34, 2021.
14

---

Published in Transactions on Machine Learning Research (08/2024)
A Internal Details of the Modules
A.1 Aspect Identification Module
As explained in the main text, the aspect identification module consists of (i) an encoder network, (ii) a slot
attention module and (i) decoder network. The architectural details of the encoder and decoder networks are summarizod in Fig.
respectively.
Conv2d
(ritters peers-16,
16,
kernel size=5,
stride 1,
nadding=2)
Conv2d (fittrrs 16, 16, kerel size_szzs S, strie= 1, padding2))
Conv2d (rliters - 16, kernel kernel size=5, strid= 1, paoding=2)
ReLu
ReLU
Soft (hidden -2.--22 resal resalution=66)
Conv2d (hiters 32, kemel strie 1, padding 22
ReLU
ReLU
Flatten
Figure 6: The architecture of the encoder network. The black arrows indicate the direction of information flow.
Conv2dtranperr (fiiterstff kerel derneLsne -pc--, str00-2 padding-2)
Conv2drancpos (fitter kernel Narnllsaa six-5, str00mm, aaddng-2)
Conv2TTaanpos (fiters kernel size strdem., padding-2)
sontrosste
rerolmmn-a)
(99%
ReLU
ReLU
Comfastrnngee (fittrs-16 kernel kerceLaze-S s0c-5, strid= paccing22)
Convasttenged (fiters: 16, kernll Spo--,
stride pooding22)
CanadTaangge (fibers -4,
herrel sizem 2 poodin-22
ReLL
Ret.
ReLU
Figure 7: The architecture of the decoder network. The black arrows indicate the direction of information flow.
For the implementation of the slot attettoon module, we have followed the pseudocode in Locatello et al. (2020), presented in Algorithm [ and have used the hyperparameters in Table In our experiments, we have experimented with 4, 6, 8, 10, 16 slots and all of them resulted around the same performance (score).
Table 1: The hyperparameters of the slot attention module.
Parameters
Values
Batch size
256
Resolution
64
Number of slots (n)
8
Number of iterations (T)
3
Warmup steps
1o4
Finally, we note that we have used the same aspect identification module across the domains (MiniGrid and Procgen) considered in this study.
15

---

Published in Transactions on Machine Learning Research (08/2024)
The pscudocode of the slot attention module (Locatello et al. 2020 Note that use the same variable names with Locatello et al. (2020).
Algorithm 1:
Input:
inputs, slots (u, diag(a))
Layer Parameters: q, k, u: linear mappings for obtaining the query, key and value vectors; GRU, MLP
for i= 0....,, T do
_lots_prev = slots
5
attention sof tmax (g(slots) k(inputs)", axis='slots')
6
updates = weightedsum (veights=attention, values=u(inputs))
slots GRU (state=slots_pree _pree, inputs-updates)
slots MLP(slots)
end
10 return slots
A.2 Permutation-Invant Filtration Module
Internally, the PI filtration module PI soft attention mechanism (Lee al. 2019). in which the key and value matrices K, V obtained by linear transformations of the slots and the query matrix is lincar transformation of fixed positional embeddings, see Algorithm 2
Algorithm 2: The pscudocode of the PI soft attention mechanism.
Input:
inputs, pos_ encodingg
Layer Parameters: 9. k, u: linear mappings for obtaining the query, key and value vectors
score = g(pos_ (ppos_encodings) k(inputs))
att stt_weights = softmax (score, axis='inputs')
output weightedsum brights-att_weighte values= v(inputs))
return output
A.3 Value-Equivalen Simulation Module
The value-equivalent simulation module consists of dynamics and prediction network. The architectural
details of these networks are summarized in Fig, S & subscripts.
For the sake of presentation clarity omit the time
BatchNormzd feature=0o))
charnels-Gn actiovation - ReLL
norm= norm= maschmarm))
State Action Encoder
Conv2d (riicers- kemel size-3, stide- padding=1)
RaLU
BatchMorr22 Knum, features-6s 64,
aod muoeen wotts 22
(fiter=:t kernel stzem,, stride=1, padding-1)
2-
100
ReLU
Rattenn
MLP
wew.anaoooo mnntss seoone
Figure 8: The architecture of the dynamics network. The black arrows indicate the direction information flow.
16

---

Published in Transactions on Machine Learning Research (08/2024)
66. actioration comeeBeeee acttora rorm- mecchnommm
HLL Whiioe as Nu -2. cnmmer
out untte- actvesionn won--alll rorm bottee-ttiui batchiomm
-11)
BeccMorrrr (uum. featuras
Corvze (nters striee pacding
RcW
Fatten
on
Ratchhoorm2e (rum, foaturas- 15
Fedar Lur 22 oer unite set--aae art0 sopooli pallin nomee
Convae (ftees stncem pedding 13
RoLu
Fltten
Figure 9: The architecture of the prediction network. flow.
The black arrows indicate the direction of information
B Details of the Training Procedure
As explained in the main text, we train our proposed architecture (sce Fig. 11 in an end-to-end manner, with the loss function Lrotal depicted in Eq. [ which is weighted sum of the reconstruction loss Lrron and the simulation loss Lsim = Cp + C. + C, In our experiments, we ended up in using an equal weighting for the losses as the different combinations of e [0, 1] and B E [0, 1] did not lead to significant increase or decrease in the generalization performances, i.e. set and 3 to 0.5.
In Crotal, the reconstruction loss Crecon is simply the squared error between the input observations and the reconstructed obscrvations ò:
Lrecon (0, 6) -0- 0)
(2)
And, the policy, value function and reward losses are simply cross-entropy losses:
Lp(x,p) = x T =Tlog
(3)
Co(2,0) =((2] logv,
(4)
(,(u,)) ( logr,
(5)
where o(-) refers to transformation from scalar representation to categorical one, and the p, and indicate the categorical outputs of the dynamics and prediction networks. Note that, in all of the loss functions, the loss is over batches of data and the time subscripts are omitted for simple presentation.
Finally, we also note that the total loss Ciotal in Eq. [ also contains an weight decay term. However, for the sake of simplicity, we have omitted it the presentation.
C Implementation Details of the Agents
PR Agent. The PR agent is a partial reasoning agent that was built by our proposed approach in Sec. This agent implemented by using the MCTS simulation framework of Niu al. (2024) More specifically, we have built upon the available MuZero implementation by replacing the internals of the Representaationleeewook class in lzero/modee/comoeeeec common, py with the internals of the aspect identification
SSee httpp/p://gtttl.cccc////aaaaaaaa
for the publicly available cade.
17

---

Published in Transactions on Machine Learning Research (08/2024)
and PI filtration modules (see App. A.1 & A.2
The hyperparameters of this agent presented in Table
Table 2: The hyperparameters of the PR agent.
Parameters
Values
Weight of policy loss
1
Weight of value loss
1
Weight of reward loss
1
10 (MiniGrid)
Number MCTS simulations
25 (Proogen)
Reanalyze ratio
0
Number of frames stacked
1
1 (MiniGrid)
Number of frames skip
4 (Procgen)
Length of game segment
400
Replay buffer size (in transitions)
1e6
TD steps
5
Number unroll steps
5
Batch size
256
Model update ratio
0.25
Reward clipping
True
Optimizer type
Adam
Learning rate
le-4
Discount factor
0.99
Frequency target network update
100
Weight decay coefficient
le-4
Max gradient norm
10
Discrete action encoding type
True
Priority exponent coefficient
0.6
Priority correction coefficient
0.4
Dirichlet noise weight.
0.25
MZ Agent. The MZ agent is a MuZero Schrittwieser et al. 2020 agent. For this agent, we have used the already available MuZero implementation in Niu et al. [2024 which directly follows the implementation of
Schrittwieser et al. (2020) The architecture of its representatiog ion network is depicted in Fig. 10 For a fair comparison, we use the same hyperparameters with the PR agent, iee. the hyperparameters Table
ResBlock channels= actiovation--an ReLL, norm -batchmorm))
ResBlock (in channels-33 32, out chwnnchaansscce actiovatio-Relu ReLL,
norm batchnormm type downsample)
(I- Buip '6 C-apule aziss raune 26 aianni PeALo
mnum netunne
ReLU
Barce
ResBlock (in channels 661 actiovation ReLL, norm- - batchnomyy
ReaBlock (in channels =4 64, actiovation-RaL. ReLL, battbbatchomm)
Avggool2d (kernel = stride-2, padding=1) 1)
AvgPool2d (keree stride-2. 2 padding=1) 1)
Figure 10: The architecture of the representation network of the MuZero (MZ) agent. indicate the direction of information flow.
The black arrows
18

---

Published in Transactions on Machine Learning Research (08/2024)
MF Agent. The MF agent is model-free counterpart of the MZ agent, i.. version of the MZ agent in which no MCTS is performed: action selection is done solely with the value predictions from the root state. This agent shares the same architecture and hyperparameters with the MZ agent.
PR-NoPIF Agent. The PR-NoPIF agent is a version of the PR agent with no PI filtration module, i.e. a version in which the aspect identifying slots are directly passed to the value-equivalent simulation module. This agent shares the same hyperparameters with the PR agent.
PR-SH Agent. The PR-SH agent is version of the PR agent which makes use of top-k semi- hard attention (rather than soft attention) in the implementation of its PI filtration module (see Algorithm 3) For this agent, we have experimented with different values of k E Zt, ranging between 2 to 7, and ended up with
setting it to 3 as it gave the best results (nole that this also what was done in Zhao el. al. 2021 This agent shares the same hyperparameters with the PR agent.
Algorithm 3: The pseudocode of the PI top-k semi-hard hard attention mechanism (Ke ct al. 2018 Zhao et al. 2021). The (opk_mask() function keeps the top k scores and sets all the others to -inf.
inputs, mos_encodings, k Layer Parameters: g, k, u: lincar mappings for obtaining the query, key and value vectors
Input:
score q(pos_encoooeelo encodings) k(inputs)
masked score renormalize mask(score, k-k))
att stt_weights softmax mcore, axis='inputs)
output veightedsum (weights-att_welle values=v(innuts))
return output
Additional Experimental Results
4-
3-
0-
0-
2
Timesteps
106
Timesteps
105
(a) Heist (Train)
(b) Heist (Test)
The gencralization performances of the PR. and bascline agents on the Heist environment. The plots are obtained again by evaluating the agent on set of test levels (referred to Test) as it gets trained set of training levels (referred to as Train) of the same environment. The training set consists of 500 randomly-sampled game levels. The test set consists of all of the possible game levels. The means and the confidence intervals are again computed over 15 independent runs.
Figure 11:
To strengthen quantitative and ablation results, also compare the generalization performances of the PR and bascline agents on the Heist environment. In this environment, the agent (a stealer) must steal the gem that is behind a network of locks. Similar lo every Proegen enviromment, a new game level is procedurally generated after every episode. For more details on the Heist environment, we refer the reader to the study Cobbe et al. 2020). We see that the performance curves in Fig. I1 corroborate the findings of Sec.
19


Reviews:
[{"review_id": "m2rvMYjcfs", "reviewer": "Reviewer_rKkT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I am familiar with model-based RL, less with attention training. My understanding is that the authors combined existing architectures to achieve generalization, claiming that their agent is able to perform attention-based reasoning, focusing on important aspects of the environment. \nThe authors also claim that, unlike prior work, their architecture 1) works from pixels, 2) allows to interpret what the agent focuses on. \n\nThe paper is well written and easy to read, but I am not sure it has enough novelty and if experiments are sufficient for the paper to be accepted. \n\n1) Experiments are very limited (only 2 pairs of environment/tasks), even though the authors write \"we compare the generalization performances [...] across a **variety** of image-based MiniGrid and Procgen domains\". \n2) I am not sure how to interpret Figure 3. I would expect that the agent is trained on one environment (for example CoinRun) and then tested on another (e.g., BigFish). But the caption only say where the agent was trained OR tested. For example, in (a) the agent was trained on DynObs-Pass, but where was it tested? \n3) I believe that the main contribution of the paper is in Sec 4.5, but it's very short and I would have expected more. This is probably due to the lack of overall experiments (few environments). Figure 4 and 5 are really interesting, but it would be much better to really have a variety of domains where the architecture is evaluated. \n4) I think comparison with other approaches are also necessary. I am not very familiar with attention-based architecture, but how would this compare to classic VAE? Wouldn't a VAE be able to identify important pixels? In the end, the background is always black and it should be easy for a VAE to give importance to non-black pixels denoting objects and the agent. \n5) I am a bit confused by the kind of \"generalization\" the authors aim for. It seems to be \"task generalization\" (even though not state explicitly, it is written that the agent needs to learn on relevant aspects of the environment, and then apply its knowledge to new task). However, \"generalization\" could also be \"environment generalization\" (same task, different environments). But isn't this what also happens in Procgen? CoinRun and BigFish are two different environment with two different tasks. Does it mean that the architecture can generalize across both tasks and environments? \n6) While there is (limited) empirical evidence about why the PI module is important, this is never explained formally. Instead, the authors always refer to Zhao et al. (2021). It would be better (and make the paper easier to understand) if the authors would add a minimal explanation of why the PI module works. \n6) There is no mention to ever releasing any code. This is an empirical paper and I believe it is important to provide code to replicate results. \n\nA few more notes:\n- Figure 1 should be large, the text within it is very hard to read. \n- About MiniGrid, the authors write \"otherwise the episode terminates without any reward\". In MiniGrid, this is not true (unless the authors have changed the implementation). The agent always receive 0 rewards, which are still rewards. The sentence should be changed to \"otherwise the episode terminates with only 0 rewards\" (or something like that).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a novel architecture for model-based RL. Compared to other architectures, the proposed one uses an attention module that allows the agent to identify important components of the environment and to generalize to different tasks. \nThe authors validate their claim on a few environments.\n\nEDIT: \nThe authors have addressed my concerns and are planning to add more experiments and release the code for the camera-ready version.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See \"strenghts and weaknesses\". To summarize:\n- Discuss more about generalization\n- More experiments (more MiniGrid or Procgen, or other environments like MiniWorld would be really nice)\n- More baselines\n- Better explanation of Figure 3\n- Release code (or at least plan to)", "broader_impact_concerns": "No concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HBT51PPUOe", "reviewer": "Reviewer_dVqb", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is clearly written and well organized. The experiments are well designed and the results are conclusive. Although it appears that both the Aspect Identification Module and Permutation-Invariant Filtration Module have been introduced in previous work for other purposes, introducing the current overall architecture along with stable training and evaluations and visualizations is important in my opinion. The final decision will depend on discussions with other reviewers. Meanwhile I ask the authors to clarify these points to help with understanding the algorithm:\n\n1. What makes the Slot Attention Module bind to objects? Although the visualizations make it clear that this part of the Aspect Identification Module discovers slots that correspond to the objects in the environment, the current text in Section 3.1 does not describe why this behavior emerges from this structure. Have Locatello et al. provided an intuitive explanation? If so, bringing that explanation here is important.\n\n2. Is the use of slot attention module for this purpose novel? As the submitted paper states, this module was introduced before by Locatello et al.. My question is whether this module was originally proposed to help with model learning and to extract objects in environments similar to the current experiments or was it originally proposed and evaluated in a completely different context. If the latter is the case, showcasing the efficacy of this module in these environments is an extra contribution of the submission. While in my opinion the submission is impactful enough either way, mentioning the original motivation for this module is helpful for situating the submission in the literature.\n\n3. What is the target for the alpha mask? The decoder outputs 4 channels which are RGB and the alpha mask and is trained to reconstruct the input. Does this mean that the input has an alpha mask? If not, how is the alpha mask channel trained?\n\n4. It appears that the architecture is based on vision transformer where the sequence is made of patches of pixels. If this is the case, it should be mentioned in the main text as currently it is not clear what the sequence is.\n\nMinor comment: There are quite a few typos through the text that I can list in the later stages of the review process. The most important one is the definition of d_0 in the last line of page 2 as a function from the state space to distributions. The function d_0 is a itself a distribution, not a function to distributions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Recent work has proposed partial reasoning models that aim to use a subset of the environment's objects in predictions. Such models have had the limitation of requiring information about the objects. The paper proposes an architecture for end-to-end training with partial reasoning models. Visualizations demonstrate that, during training, one component of the model (Aspect Identification Module) extracts objects from raw images and another component (Permutation-Invariant Filtration Module) subselects the objects in a soft manner to be used in prediction. Experiments are conducted in the generalization setting where the training and evaluation environments are slightly different. The results confirm that the algorithm with the proposed architecture outperforms a model-free or a simpler model-based algorithm and ablation studies show that partial reasoning is crucial to the performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Clarifying the points in the previous section.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ys28nCWdPL", "reviewer": "Reviewer_pXif", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Outperforms baseline methods and ablations on the chosen example\n- End-to-end trainable architecture \n- The presented approach can directly use pixels as inputs\n\nWeaknesses:\n- Only a handful of examples are provided. Are there any failure cases? If yes, what are they? I encourage the authors to provide some additional examples in the appendix. \n- Ablations on algorithm parameters are missing, e.g., what is the impact of the number of slots (on performance and the computational time)? Does this number have to be chosen according to the number of objects in the scene? What happens if there are fewer or more slots available? \n- Why were the networks for the BigFish or CoinRun examples not trained until convergence? \n-The environment variability is not fully specified, e.g., how many games are possible? Is it possible to transfer aspect identifiers, filters, or planners between environments?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an approach for building partial reasoning RL agents that attend to relevant parts of the high-dimensional input, e.g., pixels. \n\nTo achieve this, the authors propose three modules:\n1. Aspect identification: to find the 'masks' that highlight the relevant inputs\n2. Permutation-invariant filtration: filters the identified objects/regions for downstream decision-making\n3. Value equivalent simulation: plans using the filtered objects as the initial state", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "-Address the weakness (critical for recommendation)\n-The code for reproducibility is missing (which would strengthen the work). \n- Could you comment on how this scales to higher dimensional environments or different observational inputs? \nClarify the meaning of aspect, iteration, and time step (critical for recommendation) early in the manuscript\n\n-", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers generally agree that, following the discussion period and revision of the manuscript, the paper meets the TMLR criteria for acceptance. The authors have stated that they intend to add experimental results in a new environment. Based on the discussion, I think this additional experiment should be added but the reviewers have already approved the experimental setup and I do not think the outcome of the experiment would make or break the paper, so my instinct is to treat this as a minor revision.
Decision: Accept (poster)