Feature-Attending Recurrent Modules for Generalization in Reinforcement Learning
Authors: Wilka Torrico Carvalho
Andrew Kyle Lampinen
Kyriacos Nikiforou
Felix Hill
Murray Shanahan
Venue: Accepted by TMLR ()

Abstract: Many important tasks are defined in terms of object. To generalize across these tasks, a reinforcement learning (RL) agent needs to exploit the structure that the objects induce. Prior work has either hard-coded object-centric features, used complex object-centric generative models, or updated state using local spatial features. However, these approaches have had limited success in enabling general RL agents. Motivated by this, we introduce “Feature- Attending Recurrent Modules” (FARM), an architecture for learning state representations that relies on simple, broadly applicable inductive biases for capturing spatial and temporal regularities. FARM learns a state representation that is distributed across multiple modules that each attend to spatiotemporal features with an expressive feature attention mechanism. We show that this improves an RL agent’s ability to generalize across object-centric tasks. We study task suites in both 2D and 3D environments and find that FARM better generalizes compared to competing architectures that leverage attention or multiple modules.

Published in Transactions on Machine Learning Research (11/2023)
Feature-Attending Recurrent Modules for Generalization in Reinforcement Learning
Wilka Carvalho"
wcarvalho@g. hww.ard..aa..ooo.. edu
Kempner Institute for the Study of Natural and Artificial Intelligence
Horvard University
Andrew K. Lampinen
lampinen heghileeeeggggge...oo com knikfforou @google. com com mshonahan @google. com
Kyriacos Nikiforou
Felix Hill
Murray Shanahan
Google DeepMind
Reviewed on OpenReview: /tttps://openrewwe formm? id= 34y3gN7VEW
Abstract
Many important lasks are defined in terms of objects. To generalize across these tasks, reinforcement learning (RT) agent needs exploit the structure that the objects induce. Prior work has either hard-coded object-centric features, used complex object-centric generative models, updated state using local spatial features. However, these approaches have had limiled success in enabling general RL agents. Motivated by this, we introduce "Feature- Attending Recurrent Modules" (FARM), an architecture for learning state representations that relies on simple, broadly applicable inductive biases for capturing spatial and temporal regularities. FARM learns a state representation that is distributed across multiple modules that each attend to spatiotemporal features with an expressive feature attention mechanism. We show that this improves an RL agent's 's ability generalize across object-centric tasks. We study task suites in both 2D and 3D enviromments and find that FARM better generalizes compared to competing architectures that leverage attention multiple modules.
1 Introduction
Objects are key to real-world lasks. For example, a self-driving car needs to represent the movements of other cars, and household robot neods to recognize and use kitchen items. In order to generalize across tasks with objects, reinforcement learning (RL) agent should capture and exploit object-induced structure present across the tasks.
One way to capture this structure is in an agent's state representation. Unfortunately, flexibly capturing objects in a state representation is challenging because an objects have many dimensions that can vary. Consider a household robot tasked with cooking. Completing the task might require memory about objects that range size, shape, and color (e.g. stove vs. tomato). Additionally, objects motion might require that the agent represent temporal information about the objects. is unclear how best incorporate objects into stale representations to enable generalization.
Prior work has attempted to capture object-induced task structure by hand-designing object-centric state features (Diuk et al. 2008 Carvalho et al. 2021 Borsa et al. 2018 Marom & Rosman 2018). The "COBRA" agent Watters et. al. 2019) avoids hand-designing features by learning an object-centric generative model. However, these methods are limited their generality because they rely on relatively strong inductive biases. For example, COBRA relies on environments being fully-observable and objects having regular shapes
Work done during internship. Codehase: Cocphheww////////hhaaaaoc/////rr

---

Published in Transactions on Machine Learning Research (11/2023)
environment: observe remember object-motions
(b) 3D environment
regularities: object-motiins
Sources of structural regularities: 3D objects
Legend
Objects to pick
Objects to place
Agent
up
agent 2 motions
agent
4 motions
Train: smaller room
Test: larger room
Language instruction: "Go object that moved clockwise"
Language instruction: "Go to object that updown"
Language instruction: "Place toothbrush on table"
Language instruction:
"Place toothbrush table"
(Language instruction given after motions finish)
(c) KeyBox environment: retrieve key matching goal box color from across the hallway.
(d) Challenges of environments
Ballet
Place
structural regularities: object configurations
Moving objects
Language
3D
Task:
objects
Begin
Train
?
Object- Centric
2. Box determines goal color
Navigate across hallway
Memory
Obstacles
3. Obtain goal key f Bring goal key to goal box
Keybox
Test
Object
configurations
Figure 1: Three environments with different structural regularities induced by objects. In the Ballet environment, tasks share regularities induced by object motions; in the KeyBox environment, they share regularities induces by object configurations; and in the Place environment, tasks share regularities induces by 3D objects. The Ballet and KeyBox environments pose learning challenges for long-horizon memory and require generalizing to more objects. The KeyBox and Place environments pose learning challenges obstacle navigation and requires generalizing to a larger map. Videos our agent performing these tasks: See $3 for description of the problem sctting.
learn representations by predicting object segmentations. We focus on weak inductive biases in order to maximize an architecture's flexibility.
Objects can be described by subsets of features over space and time. We conjecture that weak inductive biases for capturing subsets of features over space and time may enable agents that can flexibly incorporate objects into state across wide range of environments.
We propose Feature Attending Recurrent Modules (FARM), simple but flexible architecture for learning state representations when tasks share object-induced structural regularities. FARM learns state representations that are distributed across multiple, smaller recurrent modules. To help motivate this, consider word embeddings. A word embedding can represent more information than a one-hot encoding of the same dimension because subsets of dimensions can coordinate activity to represent different patterns of word usage. Analogously, lcarning multiple modules enables FARM coordinate subsets of modules to represent different. temporal segments in agent's experiences. To capture general object-induced patterns, modules select. observation information to update with by applying mask the channels of spatiotemporal observation features.
We study FARM across three diverse object-centric environments, each with their own suite of tasks that share object-indee induced structural regularities. Tasks in the Ballet environment share regularities induced by object motions; tasks the Place environment share regularities induced by navigating towards and around 3D objects; and tasks in the KeyBox environment share regularities induced by object configurations. These environments present number of challenges. First, their state-space grows exponentially with the number of
2

---

Published in Transactions on Machine Learning Research (11/2023)
objects. The more distractor objects an environment has, the larger the chance an object will obstruct an agent's path. This requires learning policy that can navigate around distractor-based obstacles. When task objects appear in sequence, this can require long-horizon memory of object information (e.g. of goal information). Finally, tasks defined by language can require an agent learn a complex mapping (e.g. to object motions and to irregular shapes in our lasks). Across these environments, we study an agent's 's ability to recombine object-oriented memory, obstaclo-avoidance, and navigation to longer tasks with more objects.
We compare against methods with weak inductive biases for enabling objects to emerge in state representation. Recent work has shown that spatial attention is simple inductive bias for strong performance on object- centric vision tasks because it enables attending to individual objects (Greff ct al. 2020 Locatello et al.) 2020 Goyal et al. 2020a). Thus, we compare against recent RL agents that leverage spatial attention for object-centric state-update functions (Goyal et al. 2020b Mott et al. 2019).
Our core contribution is to show that we can improves an RL agent's ability to generalize to out-of- distribution (asks by having multiple modules allend to spatiotemporal features with feature attention. We expand on this below:
1. FARM leverages multiple modules that each apply featuro-wise attention to spatiotemporal features. This enables generalizing (a) memory to longer combinations of object motions ($5.1); Db; (b) navigation to 3D objects in larger environments (5.2); ; and (c) memory of goal information to longer tasks with more distractors ($5.3).
2. Competing methods have modules which leverage spatial attention, which has been shown to enable object-centric state updates. Across diverse object-centric RL tasks, we find that spatial attention has mixed benefits and can interfore with the benefits of Icarning multiple modules.
3. We hypothesize that FARM enables an RL agent to generalize to combinations of its experience by representing different temporal segments across subsets of modules (see Figure 2 In $5.31 we analyze FARM and provide evidence that object-induced temporal regularities are indeed represented across subsets of modules.
2 Related work on generalization in deep RL
The key question for generalization is how to capture structure in the problem in flexible way. How much structure do you build in? How much do you let the agent discover? Some work takes a data-driven
approach (Tobin ct al. 2017 Packer et al. 2018 Hill ct al., 2020 Justosen et al. 2019). Others have a policy
that captures task structure with either hierarchical RL (Oh et al. 2017: Zhang et al. 2018 Sohn el. al.
2018 2021 Brooks et al. 2021) or successor features (Borsa et al. 2018 Barreto et al. 2020). A final strand
focuses on lcarning invariant representations (Higgins ct al. 2017 Chaplot et al. 2018 Lee ct al. 2020 Zhang et al. 2021 or building in inductive biases (Molt et al.] 2019: Goyal et al. 2020b). In this work we focus on weak inductive biases for capturing structure. Below we review approaches most closely related to ours.
Generalizing across object-centric tasks dates back at Icast to object-oriented MDPs (Džeroski et al. 2001
Diuk et al.) 2008) which enabled generalization by representing dynamics with logical object attributes Kansky
et al.] 2017 Marom & Rosman 2018). Successor features have also leveraged objects for general generalization lization
by formulating rewards as linear with object-centric features (Borsa et al., 2018 Barreto et al. 2020). A
common thread among these directions is that they relied on hand-designed object features. Watterss et al.
(2019) avoided hand-designing features by learning an object-centric generative model (Burgess et al., 2019). However, they focused on fully-observable top-down environments with regular shapes, which allowed them to predict future object masks. This is incompatible with our environments. While research on object-centric models (Kabra et al.) 2021 Zoran et al. 2021) has progressed, these methods commonly add training complexity (more objective terms, extra modules, etc.) and make stronger assumptions (e.g. on the number of objects). We differ from this work because we focus on simple, broadly applicable inductive biases for capturing object-induced task regularities.
Generalizing with feature attention has also been studied by Chaplot et al. (0018) They showed that mapping language instructions to masks over the channels of observation features enabled generalization to language instructions with new feature combinations. While FARM also learns mask over observation
3

---

Published in Transactions on Machine Learning Research (11/2023)
features, work has two important differences. First, we develop multi-head version where different recurrent modules produce their own masks. This enables FARM to leverage this form of attention in settings where language instructions don't indicate what to attend to (this is true in 2/3 of our tasks). Second, we are the first to show that feature attention enables generalizing memory of object motions and of goal information to longer lasks (Figure and Figure f respectively).
Generalizing with top-down spatial attention. Most similar to FARM are the Attention Augmented Agent (AAA) (Mott et al., 2019) and Recurrent Independent Mechanisms (RIMs) (Goyal et al., 2020b). Both are recurrent architectures that leverage spatial attention to learn an object-centric state-update function. Both showed generalization to novel distractors. The major difference between AAA, RIMs, and FARM is that FARM attends to an observation with feature attention as opposed to spatial attention. Our experiments indicate that spatial attention has limited utility in updating state during reinforcement learning of tasks defined by object motions (Figure 1 or 3D objects (Figure 5). In terms of modularity, we also show different. results from RIMs who showed that their modules "specialize.. Our experiments suggest that in FARM, modular state instead leads subset of modules to joinlly represent. regularities in an agent's 's experience (.3.1).
3 Problem setting
We study generalization across tasks within deterministic, partially-observablee, pixel-based environments. Within environment, a task is defined by Partially Observable Markov decision processes (POMDP): M A, R. T, .²). S corresponds to environment states, A corresponds to actions that agent can take, O corresponds to the agent's observations, r = R(s, a) E IR is the reward function, s' = T(x, a) E S is the environment transition function, and o = (s) E O is an observation function that maps the underlying environment state to an RGB image.
We seek an RL agent that learns to perform tasks by finding a policy и that maximizes the expected discounted sum of rewards it obtains starting at state V(s) E - A.)] A,)]-also known the value of state. In a POMDP, the agent doesn't 't have access to the environment state. A common strategy is to instead learn an "agent state" representation, st', that compresses the full history (01, a1.7........ Gc-1, or) into a sufficient statistic suitable for selecting actions. The agent state is commonly learned with a recursive function s² = Ty(,,, 7t-1,7=-1,
Object-induced structural regularities. We study object-centric environments, where objects induce structural regularities across tasks in the reward functions R, transition functions T,, and observation functions Φ. For example, consider the KeyBox environment in Figure D (c). First, R(s, a) always specifies the goal key based on a goal box. Second, whenever the agent has to navigate around an obstacle (see Figure 2 b), the agent always sees the sprite it controls move closer to an object and then around it. This is true regardless of where in the hallway the agent observes the obstacle because of regularities in the transition function T(s, a) and observation function (s)). We want an agent that captures these regularities in its representation for state to enable zero-shot generalization to new tasks.
4 Architecture: FARM
We propose a new architecture, "Feature Attending Recurrent Modules" (FARM) for learning an agent's state representations when an environment has object-inducod structural regularities. We provide an overview of the architecture in Figure 2 Instead of representing agent state with a single recurrent function, FARM learns state representation that is distributed across recurrent functions {n²}=_ Jk=1, which we call modules (Figure 2 a). Distributing state across modules allows subsets of modules to jointly represent different regularities in the agent's 's experience (Figure 2 b). We hypothesize that having subsets of modules represent different regularities in the agent's experience enables the agent to flexibly recombine its experience for generalization.
At each time-step t, each module updates with both observation features and information from other modules. First, the agent computes observation features with a recurrent observation encoder, Z, -(o,, Z(-1). A Afterward, each module creates query vector by combining its previous module-state with the previous
4

---

Published in Transactions on Machine Learning Research (11/2023)
(a) Agent's State Representation
(b) Active Modules as Agent Completes Task
Goal key
Pickup goel key
Module 1
Module
Obstacle
Agent
Ot
Legend
Recument
Go around obstacle
connection
dz
Active Modules
Events
Figure 2: Overview of FARM. (a) FARM learns an agent state representation that is distributed across recurrent modules. (b) By distributing agent state across multiple modules, FARM is able to represent. different object-centric task regularities, such navigating around obstacles picking up goal keys, across subsets of modules. We hepothesize that this enables deep RL agent to flexibly recombine its experience for generalization. See 4 for details on the architecture and 5.3.1 for supporting analysis.
The query is used to attend to observation features via a dynamic feature attention mechanism uf f (Z,, 1-1). The query is also used to relrieve information from other modules with transformer-stylee attention mechanism va = fahare r (We explain both attention mechanisms in more detail below). Each module updates with both attention outputs to produce the next module-state h = (ut vt: g/-1). If task additionally has language description Orang (as our experiments do), the module update also updates with an embedding of this description, 21ang = fiang (Piang) Agent state is then defined by the combination of these module-states sf "e.....,h We illustrate this in Figure 16 and summarize the computations below:
action and reward, oi 1N
2. = Mor, Zz-1)
obs features
(1)
d 1 = [h/_1, a1 r1-1
query
(2)
8 S (Z, (
obs attention
V/
b = pk
Sahrro (52-1, 96-1)
share info
(4)
h = 1f (ut ; v s(-
module update
(5)
s7 [ 1........... hT]
agent state
where []] is an operation that concatenates input vectors into a long vector. We now describe each computation in more detail.
Structured spatiotemporal observation features. Our first insight is that modules can attend features describing object's motion if an agent learns observation features that describe both spatial and temporal regularities. An agent can accomplish this by learning structured spatiotemporal features with recurrent observation encoder Z, = δ(x, Z,-1) ) maxd, that share d, features across m spatial positions At each spatial position, these features both describe what is there visually along with temporal information about the recent dynamics of these features. We show example toy computations in Figure 3 (b).
Dynamic feature attention. Our second insight that feature attention an expressive attention function that can focus on desired information present across all spatial positions in observation features. An agent accomplishes this by having module predict feature coefficients that it applics to uniformly across all spatial positions in Z, (erre al. 2018 Chaplot et al.] 2018). We show example toy compulations in Figure 3
We found it useful to linearly project the features before and after using shared parameters as in Andreas
et al. 20166; Hu et al. (2018) The operations are summarized below:
fht (Z,, (Z,,g/_)) (Z.W1 oWWE o 1)) W₂
1 IOne convert. height by width ohservation features as fallows: Rhxwxds Xd.
5

---

Published in Transactions on Machine Learning Research (11/2023)
(a) Module Update
(b) Spatiotemporal Features
(c) Dynamic Feature Attention
2
Previous
module
Z
p=1
p=m
fk
att
Flatten
f sshare
17
Legend
Recurrnnt
Recument
Recureent
U2
(Wa g
featune
Position p-1
Coefficients
fkc att
Static
Static
Static
p=m
positions 771
9--1
Gc-1 Tt-1
Ot
Ot
Ot
Figure 3: Computations of FARM. (a) Schematic of updates. See 2nd paragraph in $ for details. (b) order to update with features that describe both visual and temporal regularities, the agent learns structured spatiotemporal features Z, mxd that share d, spatio-temporal features across m spatial positions. Here we show toy computations where static observations features (blue) the top and bottom row move to spatial positions to the right. The resultant spatio-tempont features (red) also include temporal information about the features (here, that the features came from leftward spatial positions). (c) fat computes coefficients for features and applies them uniformly across all spatial positions. This allows the agent to attend to all spatial position that possess desired features.
where @ denotes an element-wise product over the feature dimension and ar a sigmoid non-linearity. Since our features capture dynamics information, this allows module to attend to object motion ($5.1) When updating, we flatten the attention output. Flattening leads all spatial positions be treated uniquely and allows module represent aspects of the observation that span multiple positions, such as 3D objects ($5.2) and spatial arrangements of objects ($5.3). Since the fenture-coeffeiients for the next timc-step are produced with observation features from the current time-step, modules can dynamically shift their attention when task-relevant events occur (see Figure 7 b for an example).
Sharing information. Similar to RIMs GGyyal et al. 2020b). before updating, each module retrieves information from other modules using transformer-style atlention (Vaswani el al. 2017). We illustrate this
Figure 16 (c). We define the collection of previous module-states as H h : (m) R(++1)xd,, where 0 is null-vector used to retrieve information. A module computes "retrieval query" to search for information as dhh W/2 q E Rd, That module computes "retrieval keys and values" as K* H W R==1xxd and y He-1WraL E ga+lxds respectively. Each module then retrieves information follows:
& T
I (5 (()1) = softmax
Vt.
(8)
Vdz
Intuitively, the dot-product inside the softmax is computing 1 scores (one for each "key"), which then form probabilities. The outter dot-product multiplies each "value" by its probability and sums them to perform soft-sclection.
5 Experiments
In this section, we study the following questions:
1. Can FARM generalize memory to longer spatiotemporal combinations of object motions?
2. Can FARM generalize navigation towards and avoidance of 3D objects to larger environments?
3. Can FARM generalize memory of goal-information to larger maps with more distractor- based obstacles?

---

Published in Transactions on Machine Learning Research (11/2023)
Table 1: Baselines.
Baselines. Our first bascline is a common choice for learn- ing state-representaiions, a Long Short-term Memory
Observation
Modular
Method
(LSTM) (Hochreiter& & Schmidhuber 1997). We study two other basclines that also attend to obscrvation features: Atten- tion Augmented Agent (AAA) (Mott et al. 2019) and Re- current Independent Mechanisms (RIMs) (Goyal ct al.
Attention
State
LSTM
X
AAA
Spatial Spatial
RIMs
FARM (Ours)
Feature
2020b). Both employ transformer-style attention (Locatello
et al. 2020; Vaswani et al., 2017) to attend to individual spatial positions by reducing observation features to a weighted average over spatial positions. We instead attend to features shared across all spatial positions. RIMs, like FARM, represents state with a set of recurrent modules. We expand on the differences between baselines in
Implementation details. We implement our recurrent observation encoder, Φ, as a ResNet (He et al. 2016 followed by a Convolutional LSTM (ConvLSTM) (Shi et al. 2015). We implement the update function of each module with an LSTM. We used mulliheaa-atttenioon /Vaswani al., 2017) for / Sabrre We trained the architecture with the IMPALA algorithm (Espeholt et al.) 2018) and an Adam optimizer (Kingma & Ba 2015). We tune hyperparameters for all architectures with the "Place X next to Y" task from the Baby
environment (Chevalier- Boiisvrr- et al. 2019) ($ B.2
We expand on implementation details in $D For
details on hyperparameters, see S
5.1 Generalizing memory to more object motions
We study this with the "Ballet" grid-world (Lampinen et al. 2021 shown in Figure controls a while square that begins in the middle of the grid. There are m other "ballet-dancer" objects that move with one of 15 distinct object motions. The dances move in sequence for 16 time-steps with a 48-time-step delay in between. After all dancers finish, the agent is given a language instruction indicating the correct ballet dancer to navigate towards. All shapes and colors are randomized making motion the only feature indicating the goal object. Observations. The agent observes a top-down RBG image of the environment. Actions. The agent can move left, right, up, and down. Reward is dancer and otherwise. Tasks split. Training tasks always consists of sceing m {2, 4] dancers; testing tasks always consists of seeing m {8) dancers. All agents learn with sample budget of 2 billion frames. A poorly performing agent will obtain chance performance, 1/m.
1 (a). Tasks. The agent
touches the correct
(a) Main Results
(b) Ablation
2 dancers
2 Dancers (train)
Dancers (train)
8 Dancers (test)
1.0
0.5
Success ROT
Success Rate
Soccess
Success
FARM (4 modules)
modules No Altention Recurrent
module Atteetion Recurrent &
modules Attention Recurrent
X
module No Attention Recurrent #
X
Frames (In billions
189
Frames (In Billoon)
349
Frames
billo1s1
169
FARM LSTM
RIMe
Allention Augmanted Agent
Frames (in billions)
125
200 lell
Figure 4: FARM enables generalizing memory to longer spatiotemporal combinations of object motions. We present the success rate means and standard errors computed using 5 seeds. (a) Only FARM is able to go above chance performance for each setting. (b) Given spatiotemporal features, we find that either using multiple modules or feature attention enables learning memory of object motions. These results suggest that spatial attention removes the benefits of using multiple modules for learning to remember object motions. Encouragingly, feature attention by itself can support it.
We present the training and generalization success rates in Figure ] We learned spatiotemporal observation features with RIMs and AAA for a fair comparison. We found that only FARM is able to obtain above chance performance for training and testing. In order to understand the source of our performance, we ablate
7

---

Published in Transactions on Machine Learning Research (11/2023)
using recurrent observation encoder, using multiple modules, and using feature-attention. We confirm that a recurrent encoder is required. Interestingly, we find that either using multiple modules or using our feature-attention enables task-learning, with our feature-attention mechanism being slightly more stable.
5.2 Generalizing navigation with more 3D objects
Here, we study the 3D Unity environment from Hill et al. (2020) shown in Figure (b). Tasks. The agent is an embodied avatar in a room filled with task objects and distractor objects. The agent receives a language instruction of the form "X X on Y -eg-- "toothbrush on bed". We partition objects into two sels as follows: pickup-able objects O₁ = A U B and objects to place them on O2 = C D. Obscrvations. The agent receives first-person egocentric RGB images. Actions. The agent has 46 actions that allow to navigate, pickup and place objects. Reward is 1 if it completes the task and 0 otherwise. Tasks split. During training the agent sees A x D and B x C' in 4m x 4m room with 4 distractors, along with A x C and B x D in 3m x 3m room with O distractors. We test the agent on A x C and B x D in a 4m x 4m room with 4 distractors. We also train with "Go to X" and "Lift X
(a) Main Results
(b) Ablation
Test Setting A on C
Test Setting B on D
Test Setting A on
Test Setting B on D
1.:
Success hare
Success Palt
SUCCOOO Taatee
022
13
3.00 50
0.2
0.50
0.75
1.0
1.35
2.50
3.75
0.00
825
3.50
0.75
1,00
125
150
1.35
18
9.60
850
3.00
125
1.50
0.00
0.5
0.50
0.75
1.00
1,25
1.5 175
Romes (billons)
105
Frames (billonn)
169
Fromiss (Milioos)
289
Romes (bilions)
169
FARM
1STM
RIMs
Attention Augmenter Agent
module + wttertion
Emoditess attection
modduees attertion
Figure 5: FARM enables generalizing navigation towards and avoidance of 3D objects to a larger environment. We present the success rate means and standard errors computed using 3 seeds. (a) FARM generalizes best. (b) Our performance benefits mainly come from learning multiple modules, though feature attention slightly improves performance and lowers variance. These results suggest that spatial allention interferes with reinforcement learning of 3D objects.
We present the generalization success rate in Figure 5 We find that baselines which used spatial attention learn more slowly than an LSTM or FARM. Additionally, both models that use spatial attention have poor performance until the end of training where AAA begins to improve. FARM achicves relatively good performance, achieving a success rate of 60% and 80% on the two test settings, respectively.
5.3 Generalizing to larger maps with more objects
To study this, we create the "keyBax" environment depicted in Figure I (c). Tasks are defined with levels. Each level is a hallway with a single box and key of the same color that the agent must retrieve. The agent and the box always starts in the left-most end and the goal key always starts in the right-most end. The agent always begins in the first level. It is teleported to the next level after placing the goal key next to the goal box. The hallway for level n consists of a length-n sequence of w x w environment subsections. Fach subsoction contains d distractor objects. Obscrvations. The agent observes egocentric top-down images over short segments of the hallway. Actions. The agent can move forward, turn left, turn right, pick up objects and drop them. Rewards. When completing level, the agent gets a reward of n/nuex where Raa is the maximum level. Tasks split. Learning tasks include levels 1 to 7max = 10. Test tasks only use levels 2n and 3 We study two generalization settings: a densely populated setting with subsections of area w" = 9 and d= distractors, and a sporsely populated selling with subsections of area w² = 25 and d = 4 distractors.
We present the generalization success rates in Figure i In the dense sctting, we see an LSTM quickly overfits in both settings. All architectures with attention continue to improve in generalization performance as they continue training. In the dense setting, we find that FARM generalizes better (by about 20% for AAA
8

---

Published in Transactions on Machine Learning Research (11/2023)
(a) Generalization in densely populated setting
(b) Generalization in sparsely populated setting
(c) Ablation in sparse setting
Evel: Level-30
Level-20
Level-30 Length Distractors-1120
Level
D1stractors--00
Longth Distractors-60
Length 120. Distractoss-00
1.2 Lenrgtt to 320, Distrectonnn0
LITH
0.0
0.5
Mos
Asszntion Augeiiciiool Agoit.
0.6
2.6
Dacaen Sare
0.3
02
0.2
02
0.
000
0.3
lnc
1nd
Framss 18 billonn
3c5
Pormss 24
Figure 6: FARM enables generalizing memory of goal-information and avoidance of obstacles to larger maps with more objects. We show the the success rate mean and standard error computed using 10 seeds. (a) In the densely populated setting, FARM better generalizes to longer hallways with more distractors. (b) In the sparsely populated setting, FARM has slightly better performance than AAA for Mxxx but comparable performance for max (c) Using multiple modules and feature attention both improve generalization. These results suggest that spatial attention interferes with generalization benefits of learning multiple modules. Learning feature attention and multiple modules, instead, act synergistically.
and about 30% for RIMs). In the sparse setting, both RIMs and an LSTM fail to generalize above 30%. FARM generalizes better than AAA for level 20 but gets comparable performance for level 30. In some ways, this is our most surprising result since it is not obvious that either learning multiple modules or using feature attention should help with this task. In the next section we study possible sources of our generalization performance.
5.3.1 Analysis of state representations
We study the state reprosentations FARM learns for categorics of regularly occurring events. We collect 2000 generalization episodes in level 20. We segment these episodes into 6 categories: pickup ball, drop ball, pickup wrong key, drop wrong key, pickup correct key, and drop correct key. We study the time-series of the L2 norm of each module-state and their attention coefficients. For reference, we also show the L2 norm for the entire episode. We note that we observed consistent activity that was not captured by our simple programmatic classification of states; for example, salient activity from module 0 when the agent moved around obstacles. We show an example in Figure 7 a.
Due to space constraints, we present a subset of results in Figure y For all results, please see the $ While some modules are selective for different recurring events such as attending to goal information (Figure 7 b), it seems that subsets of modules jointly represent different aspects of state. We hypothesize that this enables FARM to leverage overlapping sets of modules to store goal-information or to navigate around abstacles in a decoupled way that supports recombination. This is further supported by our ablation where we find that having 4 or 8 modules significantly outperforms using a single large module (all had about 8M params) (Figure (c)). Feature attention consistently improves performance.
6 Discussion and conclusion
We have presented FARM, novel state representation learning architecture for environments that have object-induced structural regularities. Our results show that we can improves an RL agent's ability to generalize to out-of-distribatiio tasks by having multiple modules attend to spatiotemporal features with feature attention. Specifically this enables generalizing (a) memory to longer combinations of object-motions ($5.1p: (b) navigation around 3D objects to larger environments ($5.2; and (c) memory of goal information to longer sequences of obstacles ($5.3). Our ablations suggest that fcature attention mainly helps with long-horion memory. Interestingly, learning multiple modules helped across all conditions (memory, Dostaclo-avoidance and language learning). Our analysis suggests that learning multiple modules enables subsets to represent.
9

---

Published in Transactions on Machine Learning Research (11/2023)
Goal Box
Full map during generalization level 20
Goal Key
Modules D and
t=85
t=86
t=88
t=112
t=117
t=118
Correation actilty over event.
Partial Obssevvations
turn left
go forward
right
Observe goal key
Pickup goal
Take back goal key
Modules 2 and 6
achiity over event
100
Correlation
Active
L2 Norn Module Statue States
Module
Chernal dimersions
-0.5
Module
1.
Module
Module
Modulo
Full eplodde
Module
abject (be
1.0
Module
Module
Taredep
inactivate
piclup wrong sey
drop wrong otject (bel))
wrong keyy
Timesoop
Temporal window
piceup
(c) Average pair-wise correlation between module-states across 1000 episodes
norm of module-states
(b) Attention coefficients for module agent picks up goal key
agent around obstacles
Figure 7: We show evidence that different subsets of modules jointly represent object-induced task regularities. (a) Module commonly exhibits salient activity when the agent moves around an obstacle. (b) Module 6 activates its attention coefficients as the agent picks up the goal key. (c) Modules 2 and 6 correlate for picking up the correct key but anti-correlate for dropping the wrong object. This is similar to when neurons in word embeddings correlate for some words (e.g. man and king), but anti-correlate for other words (o.g. man and woman). In general, we find rich patterns corrclation and anti-corrclation betwoen the modules. These results suggest that FARM representing lask regularities across the modules in complicated and interesting ways. Videos of the state-activity and attention coefficients: littps://bit.t) http://bbbtl//cxaattt Cxatr
object-centric task-relevant events in flexible ways. We hypothesize that this enables a deep RL agent to flexibly recombine its experience for generalization.
We compared FARM other architectures that used spatial attention as weak inductive bias for enabling objects emerge stale representation. We found that spatial attention hindered learning tasks with object motions and 3D objects. In the KeyBox task, spatial attention seemed to help AAA most in the sparse setting with many objects. This makes sense since spatial attention has been shown to help with distractors and the agent mainly needed to ignore objects and move forward. Interestingly, pairing spatial attention with multiple modules (RIMs) removed the benefits of both.
One limitation of FARM is that fcature attention is not spatially invariant since it treat all positions as unique. Future work can look adapt this attention for something that still describes multiple positions but in spatially invariant way. Another limitation of FARM is the length of temporal regularities it can capture. Transformers Vaswani et al. 2017) have shown strong performance for representing long sequences. An interesting next-step might be to integrate FARM with a transformer. We hope that our work contributes to future RT. algo that leverage weak inductive biases for capturing object-centric task regularities.
Acknowledgments
The authors would like to thank David Abel, Rosemary Ke, Christos Kaplanis, Tony Creswell, Agnelos Filos, Satinder Singh, Honglak Kee, and Richard Lewis for feedback on these ideas. The authors would also like thank anonymous reviewers for their helpful feedback in making this work more accessible to broader audience.
10

---

Published in Transactions on Machine Learning Research (11/2023)
References
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE conference on computer vision and pottern recognition, pp. 39 48, 2016.
André Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy of Sciences, 117(49))30009-30087 2020.
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. orXiv preprinl arXiv:1812.0762) 2018
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NNmmPy programs. 2018. URL http::t/tttthhtieee coo.:/oge//////aa..
Ethan A Brooks, Janarthanan Rajendran, Richard L Lewis, and Satinder Singh. Reinforcement learning of implicit and explicit control flow in instructions. ICML, 2021.
Christopher P Burgess, Loie Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. arXio preprint arXiv: 2019.
Wilka Carvalho, Anthony Liang, Kimin Lee, Sungryull Sohn, Honglak Lee, Richard L Lewis, and Satinder Singh. Reinforcement learning for sparso-reward objoct-interaction tasks in first-person simulated 3d environments. IJCAI, 2021.
Devendra Singh Chaplot, Kanthashree Mysore Sathyendra, Rama Kumar Pasumarthi, Dheeraj Rajagopal, and Ruslan Salakhutdinov. Cated-attention architectures for task-oriented language grounding. In Proceedings of the AAAT Conference on Artificial Intelligence, 2018.
Maxime Shevallir-Booiverr, Damitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. BabyAl: First steps towards grounded language learning with human in the loop. In International Conference on Learning Representations, 2019. URL https:///penreview.nrr net/ forum?idsr XidrrJeXooX
Carlos Diuk, Andre Cohen, and Michael L Littman. An object-oriented representation for efficient reinforce- ment learning. In Proceedings of the 25th ICML, pp. 240-247, 2008.
Saão Džcroski, Lue De Racdt, and Kurt Dricssons. Relational reinforcement learning. Machine learning, 43 (1):7-53 52, 2001.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In ICML, pp. 1407-1416. PMLRR 2018.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blundell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv, 2020a.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. ICLR, 2020b.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arxiv, 2020.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, pp. 770 778, 2016.
Trina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In ICML, pp. 1480-1490 PMT.R, 2017.
11

---

Published in Transactions on Machine Learning Research (11/2023)
Felix Hill, Andrew Lampinen, Rosalia Schneider, Stephen Clark, Matthew Botvinick, James L McClelland, and Adam Santoro. Favironmental drivers of systematicity and generalization in situated agent. ICLR, 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 988):1335-1780, 1997.
Jie Hu, Li Shen, and Gang Sun. Squeeze-anddexcitttionn networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 7132-7141, 2018.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukeuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397 2016.
Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sebastian Risi. Illuminating generalization in deep reinforcement learning through procedural level generation. AAAI, 2019.
Rishabh Kabra, Daniel Zoran, Goker Erdogan, Loie Matthey, Antonia Creswell, Matthew Botvinick, Alexander Lerchner, and Christopher P Burgess. Simone: View-invariant, temporally-abstracted object representations via unsupervised video decomposition. arXiv preprint arXiv:2106.033: 2021.
Ken Kansky, Tom Silver, David A Mély, Mohamed Eldawy, Miguel Lézaro-Gredilla, Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema networks: Zero-shot transfer with a generative causal model of intuitive physics. In ICML, pp. 1809-1818. PMLR, 2017.
Dicderik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
Andrew Kyle Lampinen, Stephanie CY Chan, Andrea Banino, and Felix Hill. Towards mental time travel: a hierarchical momory for reinforcement learning agents. arXiv, 2021.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Loc. Network randomization: A simple technique for generalization in deep reinforcement learning. In ICLR, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. arXiv preprint arXi::2006.1500 15055, 2020.
Ofir Marom and Benjamin Rosman. Zero-shot transfer with deictic object-oriented representation in reinforcement learning. In NeurIPS, 2018.
Alex Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, and Danilo J Rezende. Towards interpretable reinforcement learning using attention augmented agents. NewrIPS, 2019.
Junhyuk Oh, Satinder Singh, Honglak Lee, and Kohli. Zero-shot task generalization with multi-task deep reinforcement learning. ICML, abs 1706.66666,, 2017.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. arXiv, 2018.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional Istm network: A machine learning approach for precipitation nowcasting. Advances in neural information processing systems, 28, 2015.
Sungryull Sohn, Junhyuk Oh, and Honglak Lee, Hierarchical reinforcement learning for zero-shot generalization with subtask dependencies. NewrIPS, 2018.
12

---

Published in Transactions on Machine Learning Research (11/2023)
Sungryull Sohn, Hyunjae Woo, Jongwook Choi, and Honglak Lee. Meta reinforcement learning with autonomous inference of subtask dependencies. ICLR, 2021.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In IROS, pp. 23 23-30. IEEE, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv, 2017.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based through unsupervised object discovery and curiosity-driven exploration. ar preprint arXiv:1905.0922 2019.
Amy Zhang, Sainbayar Sukhbaatar, Adam Lerer, Arthur Szlam, and Rob Fergus. Composable planning with attributes. In JCML, pp. 5842 5851. PMLR, 2018.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ICLR, 2021.
Daniel Zoran, Rishabh Kabra, Alexander Lerchner, and Danilo Rezende. Parts: Unsupervised segmentation with slots, attention and independence maximization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10439-1110 10439-10447, 2021.
13


Reviews:
[{"review_id": "AixXsoc3pA", "reviewer": "Reviewer_G2vf", "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* The network architecture is clearly illustrated with graphs and is easy to follow.\n* The proposed method achieves better results on 3 presented tasks than the chosen baselines.\n* The tasks chosen are complex and each task tests different aspects of the method. \n\n### Weaknesses:\n* Some details of the experiments are not clear\n    * Are the network capacities comparable between the proposed method and baselines?\n    * How are the language instructions input to the network?\n*  In section 5, they mention that the baselines apply a weighted sum of observation features across spatial positions while this method attends to a selected set of features shared across all spatial positions. Is this the critical design choice that makes this method better? \n*  [1] also uses a set of modules to represent the state and outperforms RIMs on a set of tasks. How does this method perform compare to [1]?\n\n\n[1] Madan, Kanika, et al. \"Fast and slow learning of recurrent independent mechanisms.\" arXiv preprint arXiv:2105.08710 (2021).\n", "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": "This paper presents a new network architecture that learns state representations for downstream long-horizon tasks. The representations learned are aggregated from multiple modules and each module utilizes attention mechanism to attend to spatial-temporal features. They evaluate this method on three long-horizon tasks that require the agent to reason about object movements and navigation. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Including more details of the experiment setup would help readers understand if the comparison to baseline methods is fair.\n* Comparison to more recent baselines such as [1] would help strengthen the experiment section.\n* Ablation experiments about the dynamic feature attention module would help understand if that design decision is important.\n", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "weEhQ5Z2Xa", "reviewer": "Reviewer_bCvr", "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- The claims are likely interesting to some TMLR readers.\n- The claims are mostly supported by evidence.\n- Enough detailed is included to make the work somewhat reproducible.\n- The writing is clear.\n\n## Weaknesses\n- Some of the claims are missing some evidence: We cannot draw strong conclusions about generalization without understanding the point at which the method _fails_ to generalize. For example, Figure 6 presented overfitting results for the LSTM, which raises the question: how long does it take the reset of the proposed methods to overfit? This is needed to more fully answer the questions about generalization performance.\n- The supplementary material seems to be missing. There are some references to appendices that do not work (e.g. the reference to hyperparameters in S5.2, and references in S5.3.1).\n- Some ambiguities in the method and baselines are present. See the requested changes.", "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": "This paper presents a method for generalization in object-centric RL settings. The method is based on several particular inductive biases about objects, which amount to a particular learnable featurization approach, called Feature Attending Recurrent Modules (FARM). These modules process information from new observations with attention to the prior state representation using queries derived from module-specific hidden states, and the prior action and reward. Information is then shared across across the attention outputs with another attention layer that uses the same attention keys as before. Generalization experiments are presented that compare the proposed approach to its ablated versions and several related prior works. The environments used in the experiments are presented in increasing complexity of test (generalization) conditions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please include the supplementary details.\n- Please run the experiments in Sec 5.3 for 1e9 frames and recompute the figures. Ideally, run them for long enough such that we observe the point at which each method begins to overfit, or that the x-axis (number of frames seen) extends quite far without overfitting (where “quite far” here would be 1e10 frames, if feasible).\n- Please make the x-axis labels consistent. In Fig 4, 5 they say “Frames (billions) … 1e9” and in Fig 6 they say “Frames (billions) … 1e8”.  I assume the latter “(billions)” is incorrect.\n- Please clarify how the language input is used by the model, and how the inputs are set during training.\n- Please add a description of how the baselines were implemented. Include a description of the network architectures and hyperparameters and how they related to the proposed method’s architectures and hyperparmeters. Ideally, the baselines use equivalent architectures to the proposed method’s where applicable, with hyperparameters tuned on a per-baseline basis.\n- Please clarify how the success rate is computed. Is it a threshold on accumulated reward?\n- Please fix typo in first sentence of abstract: “in terms of object” -> “in terms of objects”.\n- Consider rephrasing “objects manifest in a multitude of ways” to plainer language.\n- Change the Fig. 1 caption’s first sentence to summarize the Figure content. It currently summarizes the environments, not the Figure.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "LCGJ2Y02NT", "reviewer": "Reviewer_5aam", "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- Overall stronger performance compared to baselines\n- Ablation study on analysis on the proposed components\n\nWeaknesses\n- No clear intuition of why feature-wise attention can be better than spatial attention\n- Writing has room for improvement, currently it's not self-contained.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "This paper presents Feature-Attending Recurrent Modules (FARM) that learns state representations by introducing multiple recurrent modules that update their internal states with feature-wise attention. The main difference to prior work is that prior work uses spatial attention to learn object-centric representations but FARM utilizes feature-wise attention. Experiments are conducted to compare with prior approaches in three different experimental setups, where ablation studies on (i) introducing multiple moduels and (ii) feature-wise attention.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have no specific concerns on the method itself and experimental results. But the paper is not self-contained so it was very difficult to parse the method and I had to read the other papers to understand what the paper intends to do and how the proposed method works exactly. I'm still confused if my understanding of the method and results is correct. For instance, it's still not clear to me why feature-wise attention should perform better than spatial attention. I think the paper needs to be improved a lot in terms of providing intuition and illustrating the method.\n- It would be nice to have (i) clear illustration of the difference between spatial attention and feature-wise attention and (ii) explanation and motivating example of how feature-wise attention can be suitable for the considered setup.\n- Formal description of previous approaches (AAA, RIM) written in a consistent notation to the proposed method could be helpful to understand the method.\n- It's not clear what does 'attending to features' mean in Intro and Related Work before reading the method section.\n- What does 'Structured' mean in Structued spatiotemporal observation features?\n- In (6), what does [ ] mean? How are the states combined?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The first round of reviews raised concerns largely with lack of clarity in technical exposition, missing supplementary material, and asking for clarifications on relationship to some prior work, and the limits of generalization with the proposed approach.


All reviewers were satisfied with the author responses. After examination of the reviews and responses, and a less comprehensive reading through the paper, I find the proposed approach and findings interesting and valuable, and all key reviewer concerns satisfactorily addressed. 

I recommend accepting the paper after incorporating the changes made during the revision (highlighted in red). 

Decision: Accept (poster)