The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning

1 Introduction . Sensory substitution refers to the brain ’ s ability to use one sensory modality ( e.g. , touch ) to supply environmental information normally gathered by another sense ( e.g. , vision ) . Numerous studies have demonstrated that humans can adapt to changes in sensory inputs , 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 ] , or learning to ride a “ backwards ” bicycle [ 64 ] –require 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 be retrained ? †Equal Contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Modern deep learning systems are generally unable to adapt to a 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 a 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 as a form of adaptive weights of a neural network , and in this context , allows for an arbitrary number of sensory inputs that can be processed in any random order . 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 to 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 , we show that our system can be trained to perform a task even if it is given only a 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 . 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 a property observed in cellular automata ( CA ) [ 17 , 18 , 58 ] , which are mathematical systems consisting of a grid of cells that perform computation by having each cell communicate with its immediate neighbors and performing a 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 self-organizing 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 a global coherent policy for continuous control locomotion tasks . While existing CA-based approaches present a modular , self-organized 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 ) [ 23 , 40 , 42 , 76 ] have been proposed to approach the problem of learning the learning rules for a 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 to 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 metalearning 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 in 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 softmax 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 state-of-the-art methods in natural language processing [ 9 , 21 ] , image recognition [ 22 ] and generation [ 26 ] , audio and video domains [ 31 , 44 , 71 ] . Attention mechanisms have found many uses for RL [ 1 , 13 , 56 , 68 , 72 , 83 ] . 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 as a communication mechanism between independent neural network modules of a GNN , our work focuses on studying the permutation invariant properties of attention-based communication applied to 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 . 3 Method . 3.1 Background . Our goal is to devise an agent that is permutation invariant ( PI ) in the action space to the permutations in the input space . While it is possible to acquire a quasi-PI 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 ) : Rn 7→ Rm such that f ( x [ s ] ) = f ( x ) for any x ∈ Rn , and s 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 ) : Rn 7→ Rn such that h ( x [ s ] ) = h ( x ) [ 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 = σ ( QK > ) V where Q , K ∈ Rn×dq , V ∈ Rn×dv are the Query , Key and Value matrices and σ ( · ) is a non-linear function . In most scenarios , Q , K , V are functions of the input x ∈ Rn ( e.g . linear transformations ) , and permuting x therefore is equivalent to permuting the rows in Q , K , V and based on its definition it 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.1 . 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 π , which accepts the current observation ot and the previous action at−1 as its inputs . We call this new layer AttentionNeuron , and Figure 2 gives an overview of our method . Inside AttentionNeuron , we model the observation ot 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 t , the ith neuron can see only the ith component of the observation ot [ i ] . Combined with the previous action at−1 , each sensory neuron computes messages fk ( ot [ i ] , at−1 ) and fv ( ot [ i ] ) that are broadcast to the rest of the system . We then use attention to aggregate these messages into a global latent code , mt , that is PI with respect to the inputs . The operations inside AttentionNeuron can be described by the following two equations . For clarity , Table 1 summarizes the notations as well as the corresponding setups we used for the experiments . K ( ot , at−1 ) = [ fk ( ot [ 1 ] , at−1 ) · · · fk ( ot [ N ] , at−1 ) ] ∈ RN×dfk , V ( ot ) = [ fv ( ot [ 1 ] ) · · · fv ( ot [ N ] ) ] ∈ RN×dfv ( 1 ) mt = σ ( [ QWq ] [ K ( ot , at−1 ) Wk ] > √ dq ) [ V ( ot ) Wv ] ( 2 ) Equation 1 shows how each of the N sensory neuron independently generates its messages fk and fv , which are functions shared across all sensory neurons . Equation 2 shows the attention mechanism aggregate these messages . Note that although we could have absorbed the projection matrices Wq , Wk , Wv into Q , K , V , we keep them in the equation to show explicitly the formulation . Equation 2 is 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 ot . 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 mt which is PI . As long as we set 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 mt . In other words , our agent can accept inputs of arbitrary length and output a fixed sized mt . Later , we apply this flexibility of input dimensions to RL agents .

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.