The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning
Yujin Tang* (yujintang@google.com) and David Ha* (hadavid@google.com), Google Brain Tokyo
*Equal contribution
Humans adjust to new input fashion after practice
(Sadlin, 2019)
(Sampaio et al., 2001)
(Erismann et al., 1962)
(Bach-y-Rita, et al., 1962)
Can an agent do that too? Even without retraining?
Permutation
in the input space
Action:
“Turn right”
Permutation invariant
Action:
“Turn right”
Sensory neuron as a Transformer: AttentionNeuron layer
It treats each observation as a randomly ordered list of arbitrary length. Together with the action from the previous time step, it creates Keys and Values using the functions fk and fv that are shared among all neurons. The inclusion of previous actions is critical, it allows input identification.
AttentionNeuron is a layer that sits in front of the policy network.
Inspired by Set-transformer, the Queries are input independent. We use the positional encoding from the Transformer paper.
It outputs a permutation invariant message mt using attention.
Permutation invariant CartPoleSwingUp agent is robust to noise
Learning a permutation invariant equivalence using Behavior Cloning
Solving PuzzlePong, even with occlusions unseen in training
CarRacing agent zero-shot transfers to modified environments