The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning
Authors: Yujin Tang
David Ha
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: self-organization
attention
reinforcement learning
evolution strategies
zero-shot generalization
meta-learning
permutation invariance
multi-agent reinforcement learning

Abstract: In complex systems, we often observe complex global behavior emerge from a collection of agents interacting with each other in their environment, with each individual agent acting only on locally available information, without knowing the full picture. Such systems have inspired development of artificial intelligence algorithms in areas such as swarm optimization and cellular automata. Motivated by the emergence of collective behavior from complex cellular systems, we build systems that feed each sensory input from the environment into distinct, but identical neural networks, each with no fixed relationship with one another. We show that these sensory networks can be trained to integrate information received locally, and through communication via an attention mechanism, can collectively produce a globally coherent policy. Moreover, the system can still perform its task even if the ordering of its inputs is randomly permuted several times during an episode. These permutation invariant systems also display useful robustness and generalization properties that are broadly applicable. Interactive demo and videos of our results: https://attentionneuron.github.io

The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning
Yujin Tang
David Ha² Google Brain
Google Brain
yujintangggoogle.com
hadavid@google.cco com
Abstract
In complex systems, we often observe complex global behavior emerge from a collection of agents interacting with each other their environment, with each individual agent acting only on locally available information, without knowing the full picture. Such systems have inspired development of artificial intelligence algorithms in areas such as swarm optimization and cellular automata. Motivated by the emergence of collective behavior from complex cellular systems, we build systems that feed each sensory input from the environment into distinct, but identi- cal neural networks, each with no fixed relationship with one another. We show that these sensory networks be trained to integrate information received locally, and through communication via an attention mechanism, can collectively produce a globally coherent policy. Moreover, the system still perform its task even if the ordering of its inputs is randomly permuted several times during an episode. These permutation invariant systems also display useful robustness and generalization properties that are broadly applicable. Interactive demo and videos of our results: https://attene: http:s///attentiomnuuuooo gi github.oo/
1 Introduction
Sensory substitution refers the brain's 's ability to use one sensory modality (e.g., touch) supply environmental information normally gathered by another sense (e.g., vision). Numerous studies have demonstrated that humans can adapt to changes in sensory imputs, even when they are fed into the wrong channels [5, 6, 25, 64]. But difficult adaptations- such as learning to "see" by interpreting visual information emitted from a grid of electrodes placed on one's tongue [6], learning to ride "backwards" bicycle [64]-requir months of training to attain mastery. Can we do better, and create artificial systems that can rapidly adapt to sensory substitutions, without the need to retrained?
Figure 1: Comparison of visual input intended for the game player, and what our system receives. We partition the visual input from CarRacing (Left) and Atari Pong (right) into 2D grid of small patches, and randomly permute their ordering. Each sensory neuron in the system receives a stream of visual input at particular permuted patch location, and through coordination, must complete the task hand, even if the visual ordering randomly permuted again several times during an episode.
Equal Contribution
35th Conference on Neural Information Processing Systems (NeurIPS 2021).

---

Modern deep learning systems are generally unable to adapt to sudden reordering of sensory inputs, unless the model is retrained, or if the user manually corrects the ordering of the inputs for the model. However, techniques from continual meta-learning, such as adaptive weights [3, 36, 65], Hebbian-learning [52, 53, 57], and model-based [2, 20, 37, 38] approaches can help the model adapt to such changes, and remain promising active area of research.
In this work, we investigate agents that are explicitly designed to deal with sudden random reordering of their sensory inputs while performing a task. Motivated by recent developments in self-organizing neural networks [27, 55, 61] related to cellular automata [14, 17, 18, 58, 78], in our experiments, we feed each sensory input (which could be an individual state from a continuous control environment, or a patch of pixels from a visual environment) into an individual neural network module that integrates information from only this particular sensory input channel over time. While receiving information locally, each of these individual sensory neural network modules also continually broadcasts an output message. Inspired by the Set Transformer [48, 74] architecture, an attention mechanism combines these messages to form a global latent code which is then converted into the agent's action space. The attention mechanism can be viewed context, allows for an arbitrary number of sensory inputs that can be processed in any random order.
form of adaptive weights of neural network, and in this
In our experiments, we find that each individual sensory neural network module, despite receiving only localized information, can still collectively produce a globally coherent policy, and that such a system can be trained to perform tasks in several popular reinforcement learning (RL) environments. Furthermore, our system can utilize a varying number of sensory input channels in any randomly permuted order, even when the order is shuffled again several times during an episode.
Permutation invariant systems have several advantages over traditional fixed-input systems. We find that encouraging a system to learn a coherent representation of a permutation invariant observation space leads to policies that are more robust and generalizes better unseen situations. We show that, without additional training, our system continues to function even when we inject additional input channels containing noise or redundant information. In visual environments, system can be trained to perform task even if it is given only small fraction of randomly chosen patches from the screen, and at test time, if given more patches, the system can take advantage of the additional information to perform better. We also demonstrate that our system can generalize to visual environments with novel background images, despite training on a single fixed background. Lastly, to make training more practical, we propose a behavioral cloning scheme to convert policies trained with existing methods into a permutation invariant policy with desirable properties.
show that our
2 Related Work
Self-organization is a process where some form of global order emerges from local interactions between parts of an initially disordered system. It is also property observed in cellular automata (CA) [17, 18, 58], which are mathematical systems consisting of a grid of cells that perform com- putation by having each cell communicate with its immediate neighbors and performing local computation to update its internal state. Such local interactions are useful in modeling complex systems [78] and have been applied to model non-linear dynamics in various fields [14]. Cellular Neural Networks [16] were first introduced in the 1980s to use neural networks in place of the algorithmic cells in CA systems. They were applied to perform image processing operations with parallel computation. Eventually, the concept of orgf-organi neural networks found its way into deep learning in the form of Graph Neural Networks (GNN) [62, 79].
Using modern deep learning tools, recent work demonstrates that neural CA, or self-organized neural networks performing only local computation, can generate (and re-generate) coherent images [55] and voxel scenes [70, 84], and even perform image classification [61]. Self-organizing neural network agents have been proposed in the RL domain [11, 12, 59, 60], with recent work demonstrating that shared local policies at the actuator level [43], through communicating with their immediate neighbors, can learn global coherent policy for continuous control locomotion tasks. While existing Ca-based approaches present a modular, erg-ooraa solution, they are not inherently permutation invariant. In our work, we build on neural CA, and enable each cell to communicate beyond its immediate neighbors via an attention mechanism that enables permutation invariance.
Meta-learning recurrent neural networks (RNN) 22,,40,44 42, 76] have been proposed to approach the problem of learning the learning rules for neural network using the reward or error signal, enabling meta-learners to learn to solve problems presented outside of their original training domains. The

---

goals are to enable agents to continually learn from their environments in a single lifetime episode, and obtain much better data efficiency than conventional learning methods such as stochastic gradient descent (SGD). A meta-learned policy that can adapt the weights of a neural network to its inputs during inference time have been proposed in fast weights [65, 67], associative weights [3], hypernetworks [36], and Hebbian-learning
[52, 53] approaches. Recently works [46, 63] combine ideas of self-organization with meta-learning RNNs, and have demonstrated that modular meta- learning RNN systems not only can learn to perform SGD-like learning rules, but can also discover more general learning rules that transfer to classification tasks on unseen datasets.
In contrast, the system presented here does not use an error or reward signal to meta-learn or fine-tune its policy. But rather, by using the shared modular building blocks from the meta-learning literature, we focus on learning or converting an existing policy to one that is permutation invariant, and we examine the characteristics such policies exhibit a zero-shot setting, without additional training.
Attention can be viewed as an adaptive weight mechanism that alters the weight connections of a neural network layer based on what the inputs are. Linear dot-product attention has first been proposed for meta-learning [66], and versions of linear attention with sof ftaa nonlinearity appeared later (33,51, now made popular with Transformer [74]. The adaptive nature of attention provided the Transformer with a high degree of expressiveness, enabling it to learn inductive biases from large datasets and have been incorporated into ofate-of-thh-arr methods in natural language processing [9, 211, image recognition [22] and generation [26], audio and video domains [31, 44, 71j.
Attention mechanisms have found many uses for RL. [1, 13, 56, 68, 72, 83J. Our work here specifically uses attention to enable communication between arbitrary numbers of modules in an RL. agent. While previous work [32, 45, 54, 75, 81, 85] utilized attention independent neural network modules of a GNN, our work focuses on studying the permutation invariant properties of attention-based communication applied RL agents. Related work [50] used permutation invariant critics to enhance performance of multi-agent RL. Building [34, 82], Set Transformers [48] investigated the use of attention explicitly for permutation invariant problems that deal with set-structured data, which have provided the theoretical foundation for our work.
communication mechanism between
3 Method
3.1 Background
Our goal is devise an agent that is permutation invariant (PI) in the action space the permutations in the input space. While it is possible to acquire a quasi-PI Pl agent by training with randomly shuffled observations and hope the agent's policy network has enough capacity to memorize all the patterns, we aim for a design that achieves true PI even if the agent is trained with fix-ordered observations. Mathematically, we are looking for a non-trivial function f(x) : R" + RT such that (x(ss) f() for any E R", and is any permutation of the indices {1. ... n). A different but closely related concept is permutation equivariance (PE) which can be described by a function h(x)) : R" - R" such that (x((s)) [s). Unlike PI, the dimensions of the input and the output must equal in PE.
Self-attentions can be PE. In its simplest form, self-attention is described as y = (QQKT) where Q K E Raxd, V R R=xd, are the Query, Key and Value matrices and o (-) is non- con-linear function. In most scenarios, Q, K, V are functions of the input x E R (e.g. linear transformations), and permuting æ therefore is equivalent to permuting the rows in Q, K, V and based on its definition is straightforward to verify the PE property. Set Transformer [48] cleverly replaced Q with a set of learnable seed vectors, so it is no longer a function of input x, thus enabling the output to become PI. A simple, intuitive explanation of the PI property of self-attention is available in Appendix A.I.
3.2 Sensory Neurons with Attention
To create PI agents, we propose to add an extra layer in front of the agent's policy network n, which accepts the current observation o, and the previous action a;- as its inputs. We call this new layer AttentionNeuron, and Figure 2 gives an overview of our method. Inside AttentionNeuron, we model the observation or as an arbitrarily ordered, variable-length list of sensory inputs, each of which is passed into its own sensory neuron, a neural network module. Each sensory neuron only has partial access to the agent's observation, at time I, the ith neuron can see only the ith component of the observation or Combined with the previous action dz-1, 1, each sensory neuron computes messages L& (o, [i], ax-1) and [o(o)(i)) that are broadcast to the rest of the system. We then use attention to aggregate these messages into global latent code, my, that is PI with respect to the inputs.
3

---

G1-1
Ot
Sensory Neurons
0x
[2)
Kia 01-1)
[3
Attention
171p
Vine
:
Ge-1
Figure 2: Overview of Method. AttentionNNuron is a standalone layer, which each sensory neuron only has access to a part of the unordered observations Together with the agent's previous action a,-1, each neuron generates messages independently using the shared functions [or] at-1) and [o) The attention mechanism summarizes the messages into global latent code my.
The operations inside AttentionNeuron can be described by the following two equations. For clarity, Table I summarizes the notations well the corresponding setups we used for the experiments.
[u(oi[1),
E R Nxds, V(oi) = [((or(1))
K(o, a-11= =
[s(oi(()) E R Nxdr,
(1)
[s(orl[N), a0
m,= = (Ko [V))))
(2)
ddg
Equation 1 shows how each of the N sensory neuron independently generates its messages fx and fu, which are functions shared across all sensory neurons. Equation 2 shows the attention mecha- aggregate these messages. Note that although we could have absorbed the projection matrices Wg. We, W. into Q, K, V, we keep them in the equation to show explicitly the formulation. Equa- tion 2 almost identical to the simple definition of self-attention mentioned earlier. Following [48], we make our Q matrix a bank of fixed embeddings, rather than depend on the observation
Note that permuting the observations only affects the row orders of K and V, and that applying the same permutation to the rows of both K and V still results in the same m. which is PI. As long as we constant the number of rows in Q, the change in the input size affects only the number of rows in K and V and does not affect the output my In other words, our agent can accept inputs of arbitrary length and output fixed sized m. Later, we apply this flexibility of input dimensions to RL agents.
Table 1: In this notation list, provide the dimensions used in our model for different RL environ- ments, to give the reader sense of the relative magnitudes involved in each part of the system. Description
Notation
CartPole
Ant
CarRacing
Atari Pong 84x84x4
Full observation space Individual sensory input space
Ot
R
R22
R 26x96x4
Dt [
R²
R²
R5x6x4-144
R R6x6x4-144
Number of sensory neurons Dimension of action space Number of embeddings in Q Projection matrix for Projection matrix for K Projection matrix for Post-an aost-attention activation function Global latent code
N
28
(96/6)² 256
(84/6)² 196
[A]
8
(one-hot)
M
16
32
1024
400
W
R8x32
8x
R8X16
Rx32
W
Rxx
8x
R111x1
R 114xa2
W.
114xx16
144xx2
(())
tanh
tanh R32
s0 so ftmax
so softmax
R.
R 16
R. 400x32
3.3 Design Choices
It is worthwhile to have discussion on the design choices made. Since the ordering of the input is arbitrary, each sensory neuron required to interpret and identify their received signal. To achieve this, we want [x(or a1 to have temporal memories. In practice, we find both RNNs and feed-forward neural networks (FNN) with stacked observations work well, with FNNs being more practical for environments with high dimensional observations.

---

In addition to the temporal memory, including previous actions is important for the input identification too. Although the former allows the neurons to infer the imput signals based on the characteristics of the temporal stream, this may not be sufficient. For example, when controlling a legged robot, most of the sensor readings are joint angles and velocities from the legs, which are not only numerically identically bounded but also change in similar patterns. The inclusion of previous actions gives each sensory neuron a chance to infer the casual relationship between the input channel and the applied actions, which helps with the input identification.
Finally, in Equation 2 we could have combined QW E R.dxx as a single learnable parameters matrix, but we separate them for two reasons. First, by factoring into two matrices, we can reduce the number of learnable parameters. Second, we find that instead of making Q learnable, using the positional encoding proposed in Transformer [74] encourages the attention mechanism to generate distinct codes. Here we use the row indices in Q as the positions for encoding.
Experiments
We experiment on several different RL. environments to study various properties of permutation invariant RL agents. Due to the nature of the underlying tasks, will describe the different architectures of the policy networks used and discuss various different training methods. However, the AttentionNeuron layers in all agents similar, so we first describe the common setups. Hyper- parameters and other details for all experiments are summarized in Appendix A.4.
For non- con-vision continuous control tasks, the agent receives an observation vector on E RI at time t. We assign N [O] sensory neurons for the tasks, each of which sees one element from the vector, hence o.[ū] E R , = 1, 1..... (O). We use an LSTM [41] as our fa(oe [i], a;-1) to generate Keys, the input size of which is 1 + [A] (2 for Cart-Pole and 9 for PyBullet Ant). A simple pass-through function (xx = x serves as our (o)(())), and (-) land. For simplicity, we find W, = / works well for the tasks, the learnable components are the LSTM, Wg and WR-
For vision based tasks, we gray-scale and stack k 4 consecutive RGB frames from the environment, and thus our agent observes E R W xk is split into aonoolelapppng patches of size P P=0 6 using a sliding window, so each sensory neuron observes oz(8] E R6x6*k Here, fo (o (i) flattens the data and returns it, hence (or) returns tensor of shape N x dj. = N x (6 x 6 x 4) = N x 144. Due to the high dimensionality for vision tasks, we do not use RNNs for fu, but instead use a simpler method to process each sensory input. Su [o at-1) takes the difference between consecutive frames (o, [i), then flattens the result, appends or-1, and returns the concatenated vector. K(or, 0t-1) thus gives tensor of shape N d = N x [66 x 6 x 3) + |A]] = N x (108 + IAl) (111 for CarRacing and 114 for Atari Pong). We use softmaz as the non-linear activation function (-), and we apply layer normalization [4] to both the input patches and the output latent code.
4.1
Cart-pole swing up
We examine Cart-pole swing up [29, 30, 35, 86] to first illustrate method, and also use it to provide a clear analysis of the attention mechanism. We use CartPoleswingUpiiaadeer [28], a more difficult version of the task where the initial positions and velocities are highly randomized, leading to higher variance of task scenarios. In the environment, the agent observes [z, 2t cos(O), sin(O) 0.. outputs scalar action, and is rewarded at each step for getting x close to 0 and cos(f) close to 1. Inemocatios Orderi (6.42)1122)
Digus
Ixpor 31 coeci)
Exput
Figure 3: Interactive demo of SeiPPoleSwinUuHaadde In our web demo?, the user can shuffle the order of the 5 inputs at any time, and observe how the agent adapts to the new ordering of the inputs.
We use two-layer neural network as our agent. The first layer is an AttentionNeuron layer with N = 5 sensory neurons and outputs m E 'R. 16 A linear layer takes my as input and outputs scalar action. For comparison, we also trained an agent with a two-layer FNN policy with 16 hidden units. We use direct policy search to train agents with CMA-ES [39], an evolution strategies (ES) method.
Interactive demos and videos of our results available at https: attpp://aatenntiommeuro .iniiigiig github. 1o/
5

---

Table 2: Cart-pole Tests. For each experiment, we report the average score and the standard deviation
from 1000 test episodes. Our agent trained only in the environment with 5 sensory inputs.
5 obs
5 obs (shuffled)
10 obs
5 obs 5 noise
FNN (trained with obs) FNN (trained with 10 obs)
593 = 433
38 + 120
N/A
N/A
N/A
N/A
593 + 433
137 + 242
Ours (trained with obs)
472 = 426
471
426
471 425
461 + 410
Our agent can perform the task and balance the cart-pole from an initially random state. Its average score is slightly lower than the baseline (See column of Table 2) because each sensory neuron requires some time steps in each episode to interpret the sensory input signal it receives. However, as trade-off for the performance sacrifice, sensor array is randomly shuffled, which is not the case for an FÑN policy (column 2). Moreover, although our agent is only trained in environment with five inputs, it can accept arbitrary number of inputs in any order without re-training'. We test our agent by duplicating the inputs to give the agent 10 observations (column 3). When we replace the 5 extra signals with white noises with o = 0.1 (column 4), we do not see a significant drop in performance.
agent retain its performance even when the imput
Without Observation Shufflig
With Observation Shuffling
16 messages
t=1000 steps
1=1000 steps
Figure 4: Permutation invariant outputs. The output (16-dimensional global latent code) from the AttentionNeuron layer does not change when we input the sensor array as-is (left) or when we randomly shuffle the array (right). Yellow represents higher values, and blue for lower values.
The AttentionNeuron layer should possess 2 properties to attain these: its output is permutation invariant to its input, and its output carries task-relevant information. Figure is a visual confirmation of the permutation invariant property, whereby we plot the output messages from the layer and their changes over time from two tests. Using the same random seed, we keep the observation as-is in the first test but we shuffle the order in the second. As the figure shows, the output messages are identical in the two roll-outs. We also perform a simple linear regression analysis on the outputs (based on the shuffled inputs) to recover the 5 inputs in their original order. Table 3 shows the Ro values* from this analysis, suggesting that some important indicators (e.g. and 0) are well represented in the output.
Table 4: PyBullet Ant results.
Table 3: Linear regression analysis on the output. For each of the = sensory inputs we have one linear regression model with m, E R the explanatory variables.
Score
FNN (teacher)
2700 28
FNN (shuffled)
232 + 112
Ours (ES, shuffled)
2576 + 75
a
(0)
sin
0
Ours (BC, shuffled)
2034 + 948
R
0.354
0.620
0.626
0.233
0.550
Ours (BC, shuffled, larger)
2579 + 457
4.2 PyBullet Ant
While direct policy search methods such as evolution strategies (ES) can train permutation invariant RL agents, oftentimes we already have access to pre-trained agents or recorded human data performing the task at hand. Behavior cloning (BC) can allow us to convert an existing policy to a version that is permutation invariant with desirable properties associated with it.
In Table 4, we train a standard two-layer FNN policy to perform AntBul SetBuuleeEEnvee v0, a 3D locomotion task in PyBullet [19], and use it as a teacher for BC. For comparison, we also train a two-layer agent with AttentionNeuron for its first layer. Both networks are trained with ES. Similar to CartPole, we expect to see a small performance drop due to some time steps required for the agent to interpret an arbitrarily ordered observation space. We then collect data from the FNN teacher policy to train permutation invariant agents using BC. More details of the BC setup can be found in Appendix A.4.2.
Becaase our agent was not trained with normalization layers, we scaled the output from the AttentionNeuron layer by 0.5 to account for the extra inputs in the last 2 experiments.
R² measures the goodness-of of of model. An Ht² of I implies that the regression perfectly fits the data.

---

The performance of the BC agent is lower than the one trained from scratch with ES, despite having the identical architecture. This suggests that the inductive bias that comes with permutation invariance may not match the original teacher network, so the small model used here may not be expressive enough clone any teacher policy, resulting in larger variance performance. gradient-based BC, compared to RI., is that we can easily train larger networks to fit the behavioral data. We show that increasing the size of the subsequent layers for BC does enhance the performance.
benefit of environments with high dimensional visual observations that may require larger networks.
As we will demonstrate next, BC is a useful technique for training permutation invariant agents in
4.3 Atari Pong
Here, we are interested in solving screen-shuffled versions vision-based RI environments, where each observation frame is divided up into grid of patches, and like a puzzle, the agent must process the patches in a shuffled order to determine a course of action take. A shuffled version of Atari Pong [8] (See Figure 1, right pair) can be especially hard for humans to play when inductive biases from human priors [24] that expect a certain type spatial structure is missing from the observations.
But rather than throwing away the spatial structure entirely from our solution, we find that convolution neural network (CNN) policies work better than fully connected multi-layer perceptron (MLP) policies when trained with behavior cloning for Atari Pong, In this experiment, we reshape the output m, of 2002200332 2D grid of latent codes, and pass this 2D grid into a CNN policy. This way. the role of the AttentionNeuron layer is to take a list of unordered observation patches, and learn to construct 2D grid representation of the inputs be used by a downstream policy that expects some form of spatial structure in the codes. Our permutation invariant policy trained with BC can consistently reach perfect score of 21, even with shuffled screens. The details of the CNN policy and BC training can be found in Appendix A.4.3.
the AttentionNeuron layer from R x
Training Ocolusion Ratio
0.1
21
-21
-21
0.2
-3
-77
-2
-21
-20
-21
0.3
12
-3
-9
-13
22
20
-21
-21
0.4
21
17
5
0
0
-10
16
-18
-21
-21
0.5
21
15
17
15
15
10
-14
-20
-21
13
10
-10
-17
-21
0.8
-15
0.9
-17
-20
20
-19
19
-21
-21
-20
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8 0.9
Test Occlusion Ratio
Figure 5: Mean test scores Atari Pong, and example of a randomly-shufled occluded observation. In the heat map, each value is the average score from 100 test episodes. For comparison, we show the original screen (left), and the agent's observation (right). Discarded patches shown here in black.
Unlike typical CNN policies, our agent can accept subset of the screen, since the agent's 's input is a variable-length list of patches. It would thus be interesting to deliberately randomly discard certain percentage of the patches and see how the agent reacts. The net effect of this experiment for humans similar to being asked to play partially occluded and shuffled version of Atari Pong (see Figure right). During training via BC, we randomly remove a percentage of observation patches. In tests, we fix the randomly selected positions of patches discard during an entire episode.
We present the results in heat map in Figure 5 (left), where the y-axis shows the patches removed during training and the x-axis gives the patch occlusion ratio in tests. The diagram shows clear patterns for interpretation. Looking horizontally along each row, the performance drops because the agent sees less of the screen which increases the difficulty. Interestingly, an agent trained at a high occlusion rate of 80% rarely wins against the Atari opponent, but once it is presented with the full set of patches during tests, it able to achieve fair result by making of the additional information.
To gain insights into understanding the policy, we projected the AttentionNeuron layer's output in test roll-out to 2D space using (-SNE [73]. In Figure 6, highlight several groups and show their
7

---

agent accepts corresponding inputs. The AttentionNeuron layer clearly learned to cluster inputs that share similar features. For example, the sampled inputs in the blue group show the situation when the agent's paddle moved toward the bottom of the screen and stayed there. Similarly, the orange group shows the cases when the ball was not in sight, this happened right before/after game started/ended. We believe these discriminative outputs enabled the downstream policy to accomplish the agent's s task.
Figure 6: 2D embedding of the AttentionNeuron layer's output in test episode. We highlight several representative groups in the plot, and show the sampled inputs from them. For each group, we show 3 corresponding inputs (rows) and unstack each to show the time dimension (columns).
4.4 CarRacing
We find that encouraging an agent to learn a coherent representation deliberately shuffled visual scene leads to agents with useful generalization properties. Such agents still able perform their task even if the visual background of the environment changes, despite being trained only on a single static background. Out-of-domain generalization is active area, and here, combine with AttentionAgent [72], a method that uses selective, ntrd-attention via patch voting mechanism. Attention gents in [72] generalize well several unseen visual environments where task irrelevant elements are modified, but fail to generalize to drastic background changes in zero-shot setting.
method
In this experiment, combine the permutation invariant AttentionNeuron layer with the policy network used in AttentionAgent. As their hard-attentionn-aat policy is non-diffcrentiable, we train the entire system using ES. We reshape the AttentionNeuron layer's outputs to adapt for the policy network. Specifically, we reshape the output message to m, E R32 3232x88 x32x32x11 such that it can be viewed 32-by-32 grid of 16 channels. The end result is a policy with two layers of attention: the first layer outputs a latent code book to represent a shuffled scene, and the second layer performs hard attention to select the top K 10 codes from 2D global latent code book. A detailed description of the selective hard attention policy from [72] and other training details can be found in Appendix A.4.4
We first train the agent in the CarRacing [47] environment, and report the average score from 100 test roll-outs in Table As the first column shows, our agent's performance in the training environment is slightly lower but comparable to the baseline method, expected. But because randomly shuffled inputs, is still able to navigate even when the patches are shuffled. Figure (left pair) gives an illustration, where the right screen is what our agent observes and the left is for human visualization. human will find driving with the shuffled observation to very difficult because we not constantly exposed to such tasks, just like in the "reverse bicycle" example mentioned earlier.
Without additional training or fine-tuning, we test whether the agent can also navigate four modified environments where the green grass background is replaced with various images (See Figure 7). As Table 5 (from column 2) shows, our agent generalizes well most of the test environments with only mild performance drops while the bascline method fails to generalize. We suspect this is because

---

Table 5: CarRacing Test Results.
Training Env
KOF
Mt. Fuji
Ukiyoe
DS
AttentionAgent [72]
901 + 54
-81 +
-57 + 38
- -107 = 50
-66 = 23
NetRand [49]
480 + 144
20 84
356 - 159
533 = 111
-27
= 34
NetRand AttentionAgent
885
-51 +
94
656 = 131
122
Ours
801 + 147
646 + 189
503 152
661 + 140
171 - 146
CTC
6176
Figure 7: Screenshots of test environments. In each pair of images, the left is for human visualization and the right is what our agent sees. From the top left and in the clockwise order, the environments "KOF", "Mt. Fuji", "DS" and "Ukiyoe".
the AttentioonNuron layer has transformed the original RGB space to useful hidden representation (represented by my) that has eliminated task irrelevant information after observing and reasoning about the sequences of (o,, ax-1) during training, enabling the downstream hard attention policy to work with an optimized abstract representation tailored for the policy, instead of raw RGB patches.
We also compare our method to NetRand [49], a simple but effective technique developed to perform similar generalization tasks. In the second row of Table are the results of training NetRand on the base CarRacing task. The CarRacing task proved to be too difficult for NetRand, but despite low performance score of 480 in the training environment, the agent generalizes well to the "Mt. Fuji" and "Ukiyoe" modifications. In order to obtain meaningful comparison, with AttentionAgent so that it can get close to a mean score of 900 on the base task. To do that, use NetRand as input layer to the AttentionAgent policy network, and train the combination end-to-end using ES, which is consistent with our proposed method for this task. The combination attains respectable
combine NetRand
score of 885, and approach also generalizes to a few of the unseen modifications of the CarRacing environment.
in the third
of the above table, this
Our score on the base CarRacing task is lower than NetRand, but this is expected since agent requires some amount of time steps to identify each of the inputs (which could be shuffled), while the NetRand and AttentionAgent agent will simply fail on the shuffled versions of CarRacing. Despite this, our method still compares favorably on the generalization performance.
We visualize the attentions from the AttentionNeuron layer in Figure 8. In CarRacing, the agent has learned to focus its attention (indicated by the highlighted patches) on the road boundaries which are intuitive to human beings and are critical to the task. Notice that the attended positions are consistent before and after the shuffling. More details about this visualization can found Appendix A.4.4.
Discussion and Future Work
In this work, we investigate the properties of RL agents that can treat their observations as an arbitrarily ordered, variable-length list sensory inputs. By processing each input stream independently, and consolidating the processed information using attention, our agents can still perform their tasks even the ordering the observations is randomly permuted several times during an episode, without explicitly training for frequent re-shuffling (See Table 6).

---

Figure 8: Altention visualization. We highlight the observed patches that receive the most attention. Left: Training environment. Right: Test environment with unseen background.
Table 6: Reshuffl observations during roll-out. In each test episode, we reshuffle the observations every t steps. For CartPole, we test for 1000 episodes because of its larger task variance. For the other tasks, report mean and standard deviation from 100 tests. All environments except for Atari Pong have a hard limit of 1000 time steps per episode. In Atari Pong, while the maximum length of an episode does not exist, we observed that an episode usually lasts for around 2500 time steps.
CartPole
PyBullet Ant
Atari Pong
CarRacing
1 = 25
107 = 146
2053 + 225
-2
732 + 161
t 50
163 198
2319 + 188
-2
772 + 163
t = 100
242 = 254
2406 + 178
-10 12
768 + 167
200
318 310
2493 + 105
-2 -2
774 + 182
t = 500
407 = 380
2548 87
18
805 + 158
No reshuffle
472 426
2576 + 75
21 20
801 + 147
Applications By presenting the agent with shuffled, and even incomplete observations, we encourage it to interpret the meaning of each local sensory input and how they relate to the global context. This could be useful in many real world applications. For example, such policies could avoid errors due to cross- cross-wiring wiring or complex, dynamic input-output mappings when being deployed in real robots. A similar setup to the CartPole experiment with extra noisy channels could enable a system that receives thousands of noisy input channels to identify the small subset of channels with relevant information.
experimee
Limitations For visual environments, patch size selection will affect both performance and comput- ing complexity. We find that patches of 6x6 pixels work well for our tasks, as did 4x4 pixels to some extent, but single pixel observations fail to work. Small patch sizes also result in a large attention matrix which may be too costly to compute, unless approximations are used [15, 77, 80].
Another limitation is that the permutation invariant property applies only to the inputs, and not the outputs. While the ordering of the observations can be shuffled, the ordering of the actions cannot. For permutation invariant outputs to work, each action will require feedback from the environment, including reward information, in order to learn the relationship between itself and the environment.
Societal Impact Like most algorithms proposed in computer science and machine learning, method can be applied in ways that will have potentially positive or negative impacts to society. While our small-scale, self-contained experiments study only the properties of agents that are PI to their observations, and we believe our results do not directly cause harm to society, the robustness and flexible properties of the method may be of use for data-collection systems that receive data from a large variable number of sensors. For instance, one could apply permutation invariant sensory systems to process data from millions of sensors for anomaly detection, which may lead to both positive or negative impacts, if used in applications such as large-scale sensor analysis for weather forecasting, deployed in large-scale surveillance systems that could undermine our basic freedoms.
Our work also provides a way to view the Transformer [74] through the lens of self-organizing neural networks. Transformers are known to have potentially negative societal impacts highlighted in studies about possible data-Icakage and privacy vulnerabilities [10], malicious misuse and issues concerning bias and fairness [7]. and energy requirements for training them [69].
Future Work An interesting future direction is also make the action layer have the same properties, and model each "motor neuron" as module connected using attention. With such methods, it may be possible to train an agent with an arbitrary number legs, or control robots with different morphology using single policy that is also provided with reward signal as feedback. We look forward to seeing future works that include signals such as environmental rewards to train PI meta-learning agents that can adapt not only changes in the observed environment, but also to changes to itself.
10

---

Acknowledgements
The authors would like to thank Rishabh Agarwal, Jie Tan, Yingtao Tian, Douglas Eck, Aleksandra Faust and our NeurIPS2021 reviewers for valuable discussion and feedback. The experiments in this work were conducted using virtual machines provided by Google Cloud Platform.
References
[1] K. Ahuja, W. Green, and Y.-P. Li. Learning to optimize molecular geometries using reinforcement learning. Journal of Chemical Theory and Computation, 17(2):818 1--8,,, 2021.
[2] B. Amos, I. D. J. Rodriguez, J. Sacks, B. Boots, and J. Z. Kolter. Differentiable mpe for end-to-end planning and control. arXiv preprint 1rXiv:1810.1340 2018.
[3] J. Ba, G. Hinton, V. Mnih, J. Z. Leibo, and C. Ionescu. Using fast weights attend to the recent past. arXiv preprint
2016.
[4] Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint 1rXiv:607.0666 2016.
[5] P. Bach-y Rita, C. C. Collins, F. A. Saunders, B. White, and L. Scadden. Vision substitution by tactile image projection. Nature, 221(5184):963-996 1969.
[6] P. Bach-y Rita and S. W. Kercel. Sensory substitution and the human human-machine machine interface. Trends
in cognitive sciences, 7(12):541-546, 2003.
[7] E. M. Bender, Gebru, A. McMillan-Maan MMjor, and Shmitchell, On the dangers stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference Fairness, Accountability, and Transparency, pages 610-623 2021.
[8] G. Brockman, V. Cheung, L. Pettersson, Schneider, J. Schulman, Tang, and W. Zaremba. Openai gym. arXiv preprint arXiv: 1606.01540 2016.
[9] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, et Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
[10] N. Carlini, F. Tramer, E. Wallace, M. Jagielski, A. Herbert-Voss, K. Lee, A. Roberts, T. Brown, D. Song, U. Erlingsson, al. Extracting training data from large language models. arXiv
preprint arXiv:2012.07780
2020.
[11] M. Chang, S. Kaushik, S. M. Weinberg, T. Griffiths, and S. Levine. Decentralized reinforcement learning: Global decision-mate making via local economic transactions. In International Conference on Machine Learning, pages 1437-1447. PMLR, 2020.
[12] N. Cheney, R. MacCurdy, J. Clune, and Lipson. Unshackling evolution:
evolving soft robots with multiple materials and powerful generative encoding, ACM SIGEVOlution, 7(1):11-23, 2014.
[13] Choi, B.-J. Lee, and B.-T. Zhang. Multi-focus attention network for efficient deep reinforce- ment learning. arXiv preprint 1712. 2017.
[14] B. Chopard and M. Droz. Cellular automata, volume 1, Springer, 1998.
[15] K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins,
J. Davis, A. Mohiuddin, L.. Kaiser, et al. Rethinking attention with performers. arXiv preprint
arXiv:2009. /4794, 2020.
[16] L. O. Chua and L. Yang. Cellular neural networks: Theory. IEEE Transactions on circuits and systems, 35(10):1257-1222, 1988.
[17] E. F. Codd. Cellular automata. Academic press, 1968.
[18] J. Conway. The game of life. Scientific American, 223(4):4, 1970.
[19] E. Coumans and Y. Bai. Pybullet, python module for physics simulation for games, robotics and machine learning, 2016.
[20] M. Deisenroth and C. E. Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICMI.-I1), pages 465-472. Citeseer, 2011.

---

[21] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, arXiv preprint arXiv:1810.044
2018.
[22] A. Dosovitskiy, L.. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 1 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
[23] Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel. RI2: Fast reinforce- ment learning via slow reinforcement learning, arXiv preprint arXiv:1611.0277 2016.
[24] R. Dubey, P. Agrawal. D. Pathak, T. L. Griffiths, and A. A. Efros. Investigating human priors for playing video games. arXiv preprint
10217, 2018.
[25] D. Eagleman. Livewired: The inside story of the ever-changing brain. Canongate Books, 2020.
[26] P. Esser, R. Rombach, and B. Ommer. Taming transformers for high-resolution image synthesis. arXiv preprint arXiv:2012.09944 2020.
[27] V. Fortuin, M. Hüser, F. Locatello, H. Strathmann, and G. Rätsch. Som-vae: Interpretable discrete representation learning on time series. arXiv preprint orXiv:/806.02199 2018.
[28] D. Freeman, D. Ha, and L. Metz. Learning to predict without looking ahead: World models without forward prediction. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. attps://lest htap:://llarmmngtopreedicc .ig.ggggttooreiiiiiiiiibub.. 1o.
[29] A. Gaier and D. Ha. Weight agnostic neural networks. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. https:/ http::///weighttnnsseee github.io io.
[30] Y. Gal, R. McAllister and C. E. Rasmussen. Improving PILCO with Bayesian neural network dynamics models. In Data- Efficient Machine Learning workshop, ICML, Apr. 2016.
[31] R. Girdhar, J. Carreira, C. Doersch, and A. Zisserman. Video action transformer network. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 244 253, 2019.
[32] A. Goyal, A. Lamb, J. Hoffmann, S. Sodhani, S. Levine, Y. Bengio, and B. Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
[33] A. Graves, G. Wayne, and I. Danihelka. Neural turing machines. arXiv preprint arXiv: 2014.
[34] N. Guttenberg. N. Virgo, O. Witkowski, H. Aoki, and R. Kanai. coomuuatiio-eqqivariin neural networks applied dynamics prediction. arXiv preprint drXiv:1612.04530. 2016.
[35] D. Ha. Evolving stable strategies. htp:://blog.otoo ott.b..oll..ooo.eel..l 2017.
[36] D. Ha, A. Dai, and Q V. Le. Hypernetworks. arXiv preprint 1609.9/609.09100 2016.
[37] D. Ha and J. Schmidhuber. Recurrent world models facilitate policy evolution. Advances in Neural Information Processing Systems 31, pages 2451-2463. Curran Associates, Inc., 2018.
https: htt::///wlllmmdeloooo github. io.
[38] D. Hafner, T. Lillicrap, Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. arXiv preprint /81 1..00551, 2018.
[39] N. Hansen. The ema evolution strategy: comparing review. Towards a new evolutionary computation, pages 75-102, 2006.
[40] M. Haruno, D. M. Wolpert, and M. Kawato. Mosaic model for sensorimotor learning and control. Neural computation, 13(10)2201-2220, 2001.
[41] S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation, 18): 1780, 1997.
[42] S. Hochreiter, A. S. Younger, and P. R. Conwell. Learning to learn using gradient descent. In International Conference on Artificial Neural Networks, pages 87-94. Springer, 2001.
[43] W. Huang, Mordatch, and D. Pathak. One policy to control them all: Shared modular policies for agent-agnostic control. In Proceedings of the 37th International Conference Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 1 119 of Proceedings of Machine Learning Research, pages 4455 4464. PMLR, 2020.
12

---

[44] A. Jaegle, F. Gimeno, A. Brock, A. Zisserman, O. Vinyals, and J. Carreira. Perceiver: General perception with iterative attention. arXiv preprint arXiv:2103.03226. 2021.
[45] C. Joshi. Transformers are graph neural networks. The Gradient, 2020.
[46] I.. Kirsch and Schmidhuber. Meta learning backpropagation and improving it. arXiv preprint
arXiv:2012.14905, 2020.
[47] O. Klimov. Carracing-vO, v0, 2016.
[48] J. Lee, Lee, J. Kim, A. Kosiorek, S. Choi, and Y. W. Teh. Set transformer: A framework for neural networks. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3744-3753 PMLR, 09-15 Jun 2019.
attention-ban permutation-invariee
[49] K. Lee, K. Lee, J. Shin, and H. Lee. Network randomization: A simple technique for general- ization in deep reinforcement learning. arXiv preprint arXiv:1910.053300 00.0999,, 2019.
[50] I.-J. Liu, R. A. Yeh, and A. G. Schwing. Pic: permutation invariant critic for multi-agent deep
reinforcement learning. In Conference on Robot Learning, pages 590-602. PMLR, 2020.
[51] M.-T. Luong, H. Pham, and C. D. Manning. Effective approaches to attention- based neural machine translation. arXiv preprint arXiv: 040::0000804,55 2015.
[52] T. Miconi, A. Rawal, J. Clune, and K. O. Stanley. Backpropamine: training self-modifying neural networks with differentiable neuromodulated plasticity. arXiv preprint arXiv:2002.1058 2020.
[53] T. Miconi, K. Stanley, and J. Clune. Differentiable plasticity: training plastic neural networks with backpropagation. In International Conference on Machine Learning, pages 3559-3568. PMLR, 2018
[54] F. Monti, D. Boscaini, J. Masci, E. Rodola, J. Svoboda, and M. M. Bronstein. Geometric deep learning on graphs and manifolds using mixture model cnns. In Proceedings the IEEE conference on computer vision and pattern recognition, pages 5115-5124 2017.
the
[55] A. Mordvintsev, E. Randazzo, E. Niklasson, and M. Levin. Growing neural cellular automata.
Distill, 2020. sttps://distee wwtiiillpppp////// pub/2020/ (00020/rowing-co
[56] A. Mout, D. Zoran, M. Chrzanowski, D. Wierstra, and D. J. Rezende. Towards interpretable reinforcement learning using attention augmented agents. arXiv preprint arXiv: 190.002000000000 2019.
[57] E. Najarro and S. Risi. Meta-learning through hebbian plasticity in random networks. arXiv preprint rrXiv:2007.02686 2020.
[58] J. Neumann, A. W. Burks, et al. Theory of self-reproducing automata, volume 1102024. University of Illinois press Urbana, 1966.
[59] Ohsawa, K. Akuzawa, Matsushima, G. Bezerra, Iwasawa, H. Kajino, S. Takenaka, and
Y. Matsuo. Neuron as agent, 2018.
[60] J. Ott. Giving up control: Neurons reinforcement learning agents. arXiv preprint 2xX1v:2003.11642, 2020.
[61] E. Randazzo, A. Mordvintsev, E. Niklasson, M. Levin, and S. Greydanus. Self-classifying mnist digits. Distill, 2020.
[62] B. Sanchez-Lengeling neural networks. Distill, 2021. https://distill.pub/200 1ttpa///222//gnnntrro.....
Reif, A. Pearce, and A. Wiltschko. A gentle introduction to graph
[63] M. Sandler, M. Vladymyrov, A. Zhmoginov, N. Miller, A. Jackson, T. Madams, et al. Meta- learning bidirectional update rules. arXiv preprint 2XXx:2100400466 2021.
[64] D. Sandlin. The backwards brain bicycle: Un-doing understanding, 2019.
[65] Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
[66] J. Schmidhuber. Reducing the ratio between learning complexity and number of time varying variables in fully recurrent nets. In International Conference on Artificial Neural Networks, pages 460-463. Springer, 1993.
13

---

[67] J. Schmidhuber. A 'self-referential' weight matrix. In International Conference on Artificial Neural Networks, pages 446 450. Springer, 1993.
[68] I. Sorokin, A. Seleznev, M. Pavlov, A. Fedorov, and A. Ignateva. Deep attention recurrent q-network. arXiv preprint arXiv: /512.0169, 2015.
[69] E. Strubell, A. Ganesh, and A. McCallum. Energy and policy considerations for deep learning in nlp. arXiv preprint ar.ivv:996.0000 02243, 2019.
[70] S. Sudhakaran, D. Grbic, S. Li, A. Katona, E. Najarro, C. Glanois, and S. Risi. Growing 3d arte- facts and functional machines with neural cellular automata. arXiv preprini arXiv:2 103.08837 2021.
[71] C. Sun, F. Baradel, K. Murphy, and C. Schmid. Learning video representations using contrastive bidirectional transformer. arXiv preprint (et:::9900055444 2019.
[72] Y. Tang. D. Nguyen, and D. Ha. Neuroevolution of interrnnteeeee agents. In Proceedings of the Genetic and Evolutionary Computation Conference, 2020. ttap::///attentonagem. github. io.
[73] I.. Van der Maaten and G. Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(1 1), 2008.
[74] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, u. Kaiser, and I. Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus. S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[75] P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Lio, and Y. Bengio. Graph attention networks. arXiv preprint arXiv: 10900, 2017.
[76] J. X. Wang, 7. Kurth-Nelson, D. Tirumala, H. Soyer, J. 7. Leibo, R. Munos, C. Blundell, D. Ku- maran, and M. Botvinick. Learning to reinforcement learn. arXiv preprint 1rXiv:1611.05770 2016.
[77] S. Wang, B. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04777 2020.
[78] S. Wolfram. Cellular automata models of complexity. Nature, 3199555541998 1984.
[79] 7. Wu, S. Pan, F. Chen, G. Long, C. Zhang, and S. Y. Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 2020.
[80] Y. Xiong, Z. Zeng, R. Chakraborty, M. Tan, G. Fung. Y. Li, and V. Singh. Nyströmformer: nyström-base algorithm for approximating self-attention. arXiv preprint arXiv:2102.03900, 2021.
[81] S. Yun, M. Jeong, R. Kim, J. Kang, and H. Kim. Graph transformer networks. arXiv preprint /rXiv: /XXv:/9/1.06455. 1.0655, 2019.
[82] M. Zaheer, Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. arXiv preprint arXiv:1703.06111 2017.
[83] V. Zambaldi, D) Raposo, A. Santoro, V. Bapst, Y. Li, I. Babuschkin, K. Tuyls, D. Reichert, T. Lillicrap, E. Lockhart, M. Shanahan, V. Langston, R. Pascanu, M. Botvinick, O. Vinyals, and P. Battaglia. Deep reinforcement learning with relational inductive biases. In International Conference on Learning Representations, 2019.
[84] D. Zhang, C. Choi, Kim, and Y. M. Kim. Learning to generate 3d shapes with generative cellular automata. arXiv preprint arXiv:2103.04130 2021.
[85] J. Zhang. X. Shi, J. Xie, H. Ma, King, and D.-Y. Yeung. Gaan: Gated attention networks for learning large and spatiotemporal graphs. arXiv preprint arXiv:1803.072. 0.72944 2018.
[86] X. Zuo. Pytorch implementation of improving pilco with bayesian neural network dynamics models, 2018. https:// //gi thub. com wwtmm//uuuiignog//////e......
14


Reviews:
[{"review_id": "rX6mmPiA1G3", "reviewer": "Reviewer_Bz2K", "summary": "This paper presents a sensory neuron architecture as a layer in a neural network trained for reinforcement learning and claims that this sensory neuron makes a network generally permutation-invariant. There is an extensive related-work section covering examples of sensory modality adaptation in neuroscience, cellular automata, meta-learning, and attention, setting up the presented approach which uses collective intelligence of only distributed modules with only local awareness, self-organization through attention, and adaptation of an existing policy to be permutation invariant. The paper then presents a mathematical formulation for the AttentionNeuron architecture, including its temporal memory, to be used as the first layer in an RL agent. The paper next tests its approach on four different tasks, some with multiple environments and concludes with some discussion of the results which appear to be robust to permutation of the input vector unlike the baseline feedforward network. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Originality:* I think this paper’s originality comes from the fact that it is a novel combination and then application of well-known techniques and concepts. The paper notes that Transformer’s attention mechanism is known to be permutation invariant and applies it here to reinforcement learning, which is the primary novelty I see in this work. The paper suggests an angle of collective intelligence due to the AttentionNeuron architecture, where each of N sensory neurons sees one of each of N observation components. It seems that the transformation from input to sensory neuron layer is a generalized permutation matrix, and the paper’s point is that this can be any generalized permutation matrix as long as the magnitude of each column remains the same (I think; I may not have all the details totally right). The connection to collective intelligence seems limited to a high-level abstraction, but I do think it’s an elegant and unique approach to permutation invariance whether or not it would be considered collective. \n\n*Quality:* I think the methods and results of this paper are quite well-conceived and -executed. Though the method seems complex at first and to some degree needlessly complex, the discussion of design choices such as temporal memory and reduced network size is helpful. The solution is actually quite thoughtful, and the issues are more of presentation than content (discussed below). The experiments are impressive - the approach is used successfully in four different testbeds, some with multiple environments. Though there are some caveats with each one, particularly Atari Pong, the core message of permutation invariance seems to persist through the results. In my opinion, the paper is really strengthened by forgoing complicated cuts of data and various statistics to just show performance on a large range of testbeds. The only improvement may be ablations or analysis to rule out confounds.\n\n*Clarity:* This paper is somewhat difficult to comprehend. Section-specific critique: \n\nThe related work is quite ambitious. I think all but the attention section could be collapsed into a section about permutation invariance and *maybe* collective intelligence, though as I said in *Originality*, that connection seems vague. Many of these are only related loosely, and the authors say as much at the end of each paragraph. This paper presents an interesting layer architecture that achieves a clear and concrete goal, and related work about that goal and similar architectures is enough in my opinion. This related work section is confusing and to some degree feels overpromising despite a paper that doesn’t need to overpromise. This makes the paper difficult to understand. \n\nIn the methods section, the notation is difficult to follow, partially because it is mostly in paragraph form. Figure 2 is quite helpful, though the 1:1 nature of the observation components to sensory neurons could be visually clearer, but the rest of the text on page 4 is dense with words and notation combined. What makes it more frustrating is that it is notation, not derivation or theory, so it shouldn’t be difficult to follow; a diagram or creative visual aids for the central definitions might help. \n\nFinally, the results are very hefty. That’s one of the biggest value-adds of this paper, but it’s also overwhelming because it’s presented as one long laundry list. I realize that each testbed is very different so it may be hard to have a claims-driven structure to this section, but it would be nice to see unifying stats or even unifying qualitative analysis earlier than Section 5. \n\n*Significance:* I think the ambitious and extensive related work section causes this paper to feel like it’s overstating its significance; I don’t think neuroscience is really a related field at all, nor am I convinced that this is an example of collective intelligence, nor do I know why it would make a difference. It’s also hard to tell if permutation invariance itself is significant - I think the significance of the main contribution comes from the fact that it’s 1) an elegant architectural approach (which builds on Transformer), unlike existing permutation-invariance works that require specific input formats and 2) application to reinforcement learning. If the language focused on what this paper is really doing, as opposed to big ideas that sound kind of like what it is doing, it would come across better. Nonetheless, I do think that this result is significant. The applications listed in section 5 are a good taste of what this method can do, and I suspect there are many more, and more immediate, applications that could be explored from this work.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations receive some discussion, though I think the paper would benefit from more beyond resource usage. I also don’t think the lack of invariance to action permutation is as significant as a lack of invariance to observation permutation would be, though this could be a limit of my own knowledge. Ablation experiments may have been helpful to yield this section. The societal impact section is considerable and far-reaching, and I think it briefly hits on many important points. I want to clarify that I think the way this paper is written calls for exactly this type of societal impact section, so the authors have done a good job. That said, I personally prefer depth and engagement with societal implications rather than a broad laundry list, and if the paper had a narrower focus, Section 6 could focus on just the implications of permutation invariance (i.e. some of its main applications) and/or RL agents. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "qRcecezV-yA", "reviewer": "Reviewer_TUZP", "summary": "This work unites previous research on Transformers and concepts from cellular automata to offer a novel model for learning global policies from local arbitrarily ordered variable-length sensory inputs. The authors evaluate their model using four standard RL tasks and show that their model, compared to the state-of-art, is robust to perturbations and occlusions of the inputs and also is significantly more effective for out-of-distribution generalization", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work provides an interesting and highly valuable method for building RL models robust to various types of noise and capable of out-of-distribution generalization, which is a highly important and widely pursued goal in RL. To attain this goal, the authors propose a novel framework based on Set Transformers (a Transformer-based method for permutation-invariant neural-network analysis of datasets) and also based on concepts of cellular automata. The authors rigorously examine the properties of their new framework using a diverse set of standard RL tasks. As a result, the authors show impressive capabilities of their model to perform robustly in the presence of large-scale occlusions, noise, and out-of-distribution versions of the tasks. With that said, the text could use numerous clarifications and a comprehensive methods section, which would make the paper accessible for the broader NeurIPS community. Overall, the novel model, sound methodology, thorough testing, and highly significant results presented in this paper make it a valuable contribution for NeurIPS.\n\nSpecific comments:\n\nAlthough the paper seems to have all necessary details to be understood, these details appear to be scattered throughout the text (and throughout the referenced literature) which may require reading the paper 2-3 times to cover all the important details. To this end, I suggest making use of unlimited supplementary space to write comprehensive methods, which – alone – would be sufficient to reproduce the model and the experiments.\n\nThe paper may also benefit from adding the details of the framework to the main text. Notably, the Set Transformers used here are much less known than the original Transformer framework, so it may be unfair to expect the reader to know its details or to be motivated to read a separate paper on it. As a good example, the Set Transformers paper itself starts with a brief yet sufficient introduction of the Transformer framework and explains the design choices for the Set Transformer framework immediately upon their appearance. One of such design choices – not explained in the current paper – is that SensoryNeuron, just like Set Transformers, attends the input data (Keys, Values) to a (smaller) set of separately learned/handcrafted Query embeddings – unlike the usual Transformers attending the elements of input data to each other. It is important to explain that learnable/predefined Query embeddings do not depend on the input and decrease the algorithm’s complexity from quadratic to linear.\n\nThe mathematical notation here can also be simplified, as it appears to have multiple variables which are not used. For example, the matrices Wq, Wk, and Wv were originally introduced in Transformer to be used in multi-head attention, where the different versions of these matrices were applied to the same Queries, Keys, and Values. It appears that the framework presented in this submission does not use multi-head attention, so these matrices can be sacrificed for the clarity of how the main model is explained. The matrix Wk, in particular, can be described together with the function Fk as both of them are implemented using neural networks – stacked with each other and thus forming a single neural network. The same goes for Fv and Wv – both being implemented as the identity matrices and thus unnecessary. Decomposition of Q and Wq is justified as it decreases the overall number of parameters and because Q is not learnable – but this distinction does not need to be immediately introduced. Focusing rather on the roles of Q, K, and V would make the paper more accessible for the broad audience.\n\nFinally, I found the neuroscience-oriented introduction to this paper questionable. While the paper presents interesting, significant, and valuable results regarding the robustness of the algorithm to various types of noise, taking advantage of additional data, and out-of-distribution generalization – these results are summarized only at the end of the paper. At the same time, a reader who starts reading the text from the beginning may falsely think that the purpose of the paper is to learn from/replicate a biological phenomenon. Moreover, the introduction – in light of the details of the proposed model – may leave a reader with the impression that the brain information processing is permutation invariant, like in zero-shot learning, whereas it’s not. The neuroscience papers cited in the introduction, such as reverse biking or learning to “see” using an electrode matrix, clearly mention the significant time tolls of acquiring such skills. Although the presented work may have future neuroscience implications, the line of the argument currently presented here does not make a solid connection to neuroscience and, as such, needs to be revised/shrunk. Instead, the paper can start from the useful properties of the model, as outlined in the conclusions. The same goes for the abstract. The reclaimed space can be used for a detailed description of the SensoryNeuron, which is a central part of this work.\n\nI would like to highlight here that all my suggestions above do not devalue the significance and the validity of the results. They are intended to highlight that a thorough explanation of the design choices, a broader overview of the Transformer, and starting from the robustness/generalization properties will make the paper a more accessible read for the broad NeurIPS audience. Even in the current format, this work is publishable in NeurIPS.\n\nMinor comments:\n\n-Why is the pass-through function Fv needed at all? Could you please briefly comment on the intuition why Fv, normally being a learnable parameter in Transformers, does not require learning here?\n\n-Could you please also compute the statistical significance for the RL results? Given the huge standard deviation in the CartPole task, the SensoryNeuron’s apparent slight deficiency may not be statistically significant, which will strengthen your claim.\n\n-Can you please explain why you spatially arrange the messages Mt in the AtariPong task? Naively, it would seem that, as you permute your sensory neurons, the 2D arrangement of their messages will be also random and the convolution on top of these messages will learn the same as the feedforward connectivity?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "9"}, "extra_text": {}}, {"review_id": "j82ZbnJy39", "reviewer": "Reviewer_rWCR", "summary": "This work investigates the problem of shuffled or dropped out sensory inputs for reinforcement learning agents. The authors propose AttentionNeuron, a permutation-invariant attention layer operating on neurons which independently process the input currently present at its location. Each neuron utilizes its past context and the past action to interpret the seen input, and the attention mechanism combines the outputs of each neuron to generate a globally-aware embedding to be used for a standard policy architecture. The authors show that (1) AttentionNeuron is robust to shuffled and noisy observations on CartPole, (2) performing behavior cloning with an AttentionNeuron architecture can yield a policy robust to shuffling on Ant, (3) AttentionNeuron is robust to patch-level shuffling and dropouts on Pong, (4) AttentionNeuron improves robustness to drastic background changes on CarRacing, and (5) AttentionNeuron is robust to shuffling within a single episode.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros\n1. Studies novel and biologically inspired permutation invariance question in the context of RL\n2. Experiments evaluate on several diverse environments (CartPole, Ant, Pong, CarRacing)\n3. Writing is generally clear\n\nCons\n1. It is not clear to me that permutation invariance is a practically significant or useful property\n2. Minimal ablations/comparisons to other methods in the experiments section\n\nAt a high level, I think this work is promising: it investigated a little-studied problem, permutation invariance to sensory observations in RL, motivated by observations in neuroscience regarding malleability in human brains. While it is not clear to me why this is a useful property from an engineering perpsective (i.e. as to why this can't be mostly tackled by the sensor interface), I think the direction studying biologically plausible RL/models is exciting. Furthermore, I think additional practical benefits gained as a result of permutation invariance (such as better generalization) are particularly exciting/important.\n\nMy main concerns are about the lack of comparisons/ablations to other methods, as the authors propose a method with several components but only compare to a feedforward network. I don't necessarily believe that my choices below would outperform the proposed method, but I think they would strengthen the paper and better contextualize the presented numbers. To propose a non-exhaustive but simple set of choices:\n1. Transformer/vision transformer: a transformer operating directly on the sensory inputs/image patches without positional embeddings is permutation invariant in the same sense as AttentionNeuron, but is a simpler approach (i.e. AttentionNeuron without the LSTM input cells). I think this ablation would make clear whether or not the recurrent layers are necessary, and if it is important that the bottom-level cells \"orient\" themselves,or if simply the upper-level attention layers are capable of making sense of the disordered information. I also think a comparison with positional embeddings would be useful, as vision transformers currently represent a state-of-the-art in image-based processing.\n2. LSTM after image processing: similarly to my above comment, this would clarify if this recurrent reorganization should be done at a low-level or a high-level. I think this is also biologically interesting, and practically it would be less computationally expensive since there wouldn't need to be an LSTM at every input location\n3. Other generalization methods: as I mentioned above, I think a particularly exciting experiment is Section 4.4/Table 4 because it promises that generalization could be improved via permutation invariance. However, the experiment only compares to a minimal baseline. I think comparing to other representation learning methods, such as bisimulation [1] or contrastive learning [2] (or others), would make this experiment more exciting.\n\nOverall, I think the direction toward permutation-invariant representations for sensory observations is underexplored and novel in the RL setting. Also, I think the authors investigate a few different interesting properties of the method, though I think the baselines are weak. As a result of the aforementioned weaknesses, I think this work currently does not meet the bar for publication. I am open to changing my score if the above issues are addressed, particularly if the comparisons to the transformer and high-level LSTM are included.\n\n=====\n\nReferences\n\n[1] Zhang, Amy et al 2020. \"Learning Invariant Representations for Reinforcement Learning without Reconstruction.\"\n\n[2] Srinivas, Aravind et al 2020. \"CURL: Contrastive Unsupervised Representations for Reinforcement Learning.\"\n\n=====\n\nUpdates\n\nI am satisfied with the authors' responses and additional baselines, therefore I am increasing my score from 4 to 6.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors adequately addressed the limitations and potential negative societal impact", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "cuElRNecla", "reviewer": "Reviewer_dfwx", "summary": "This paper extends the idea of set transformer to motor control, so that an agent can perform well even when its sensory input channels are shuffled and partly presented.\nThe performance is verified in cart-pole, ant, pong, and car racing tasks by using evolutionary optimization and behavioral closing of pre-trained policies.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\nThe present result is an extension of the set transformer (ICML 2019). Although the title includes \"reinforcement learning\", actual learning takes evolutionary optimization and supervised learning (behavioral cloning).\n\nQuality:\nI wish the role of the previous action input to the attention network is better discussed. My guess is that the network compares the changes in the observations and action taken to infers which channel carries the critical information for control. It would be interesting to analyze which channels actually gathered attention (e.g., the patches including the car and road boundary?) and how that is affected if the past action is not presented.\n\nClarity:\nThe math is clearly presented.\nIn equation (1), most RL folks imagine Q as action values and V as state value, so their meanings as Query and Value may be better explained earlier.\n\nSignificance:\nAlthough the architecture is reported to be inspired from sensory neurons, it has only remote link to neuroscience.\nAlthough not tested in this paper, permutation-invariant feature would be helpful in the context of multi-agent learning, where an agent need to work with a variable number of agents coming and going in random order.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "These are well considered and documented.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1h"}, "extra_text": {}}]

Author rebuttal:

Meta-review: I thank the authors for their submission and active participation in the discussions. Reviewer dfwx has pledged that with the revisions by the authors, they would raise their score. Since reviewer dfwx has not done that yet and did not actively participate in the discussion with the other reviewers, I assume a higher rating than the one currently recorded. Taking this into account, it seems reviewers unanimously agree that this paper is worthy of publication. In particular, reviewers appreciated the diverse set of environments used for evaluation [rWCR,TUZP], that is well executed [Bz2K] and that it presents an original synergy between existing techniques [Bz2K,TUZP]. During rebuttal and discussions, reviewer rWCR's concerns regarding baselines were addressed and reviewers TUZP and Bz2K agreed that the paper is improved based on the author rebuttal. Thus, I recommend acceptance and encourage the authors to further improve the clarity of their paper based on the reviewer feedback.
Decision: Accept (Spotlight)