Implementing Inductive bias for different navigation tasks through diverse RNN attrractors

1 INTRODUCTION . Spatial navigation is an important task that requires a correct internal representation of the world , and thus its mechanistic underpinnings have attracted the attention of scientists for a long time ( O ’ Keefe & Nadel , 1978 ) . A standard tool for navigation is a euclidean map , and this naturally leads to the hypothesis that our internal model is such a map . Artificial navigation also relies on SLAM ( Simultaneous localization and mapping ) which is based on maps ( Kanitscheider & Fiete , 2017a ) . On the other hand , both from an ecological view and from a pure machine learning perspective , navigation is firstly about reward acquisition , while exploiting the statistical regularities of the environment . Different tasks and environments lead to different statistical regularities . Thus it is unclear which internal representations are optimal for reward acquisition . We take a functional approach to this question by training recurrent neural networks for navigation tasks with various types of statistical regularities . Because we are interested in internal representations , we opt for a two-phase learning scheme instead of end-to-end learning . Inspired by the biological phenomena of evolution and development , we first pre-train the networks to emphasize several aspects of their internal representation . Following pre-training , we use Q-learning to modify the network ’ s readout weights for specific tasks while maintaining its internal connectivity . We evaluate the performance for different networks on a battery of simple navigation tasks with different statistical regularities and show that the internal representations of the networks manifest in differential performance according to the nature of tasks . The link between task performance and network structure is understood by probing networks ’ dynamics , exposing a low-dimensional manifold of slow dynamics in phase space , which is clustered into three major categories : continuous attractor , discrete attractors , and unstructured chaotic dynamics . The different network attractors encode different priors , or inductive bias , for specific tasks which corresponds to metric or topology invariances in the tasks . By combining networks with different inductive biases we could build a modular system with improved multiple-task learning . Overall we offer a paradigm which shows how dynamics of recurrent networks implement different priors for environments . Pre-training , which is agnostic to specific tasks , could lead to dramatic difference in the network ’ s dynamical landscape and affect reinforcement learning of different navigation tasks . 2 RELATED WORK . Several recent papers used a functional approach for navigation ( Cueva & Wei , 2018 ; Kanitscheider & Fiete , 2017b ; Banino et al. , 2018 ) . These works , however , consider the position as the desired output , by assuming that it is the relevant representation for navigation . These works successfully show that the recurrent network agent could solve the neural SLAM problem and that this could result in units of the network exhibiting similar response profiles to those found in neurophysiological experiments ( place and grid cells ) . In our case , the desired behavior was to obtain the reward , and not to report the current position . Another recent approach did define reward acquisition as the goal , by applying deep RL directly to navigation problems in an end to end manner ( Mirowski et al. , 2016 ) . The navigation tasks relied on rich visual cues , that allowed evaluation in a state of the art setting . This richness , however , can hinder the greater mechanistic insights that can be obtained from the systematic analysis of toy problems – and accordingly , the focus of these works is on performance . Our work is also related to recent works in neuroscience that highlight the richness of neural representations for navigation , beyond Euclidian spatial maps ( Hardcastle et al. , 2017 ; Wirth et al. , 2017 ) . Our pre-training is similar to unsupervised , followed by supervised training ( Erhan et al. , 2010 ) . In the past few years , end-to-end learning is a more dominant approach ( Graves et al. , 2014 ; Mnih et al. , 2013 ) . We highlight the ability of a pre-training framework to manipulate network dynamics and the resulting internal representations and study their effect as inductive bias . 3 RESULTS . 3.1 TASK DEFINITION . Navigation can be described as taking advantage of spatial regularities of the environment to achieve goals . This view naturally leads to considering a cognitive map as an internal model of the environment , but leaves open the question of precisely which type of map is to be expected . To answer this question , we systematically study both a space of networks – emphasizing different internal models – and a space of tasks – emphasizing different spatial regularities . To allow a systematic approach , we design a toy navigation problem , inspired by the Morris water maze ( Morris , 1981 ) . An agent is placed in a random position in a discretized square arena ( size 15 ) , and has to locate the reward location ( yellow square , Fig 1A ) , while only receiving input ( empty/wall/reward ) from the 8 neighboring positions . The reward is placed in one of two possible locations in the room according to an external context signal , and the agent can move in one of the four cardinal directions . At every trial , the agent is placed in a random position in the arena , and the network ’ s internal state is randomly initialized as well . The platform location is constant across trials for each context ( see Methods ) . The agent is controlled by a RNN that receives the proximal sensory input , as well as a feedback of its own chosen action ( Fig . 1B ) . The network ’ s output is a value for each of the 4 possible actions , the maximum of which is chosen to update the agent ’ s position . We use a vanilla RNN ( see Appendix for LSTM units ) described by : ht+1 = ( 1− 1 τ ) ht + 1 τ tanh ( Wht +Wif ( zt ) +WaAt +WcCt ) ( 1 ) Q ( ht ) = Woht + bo ( 2 ) where ht is the activity of neurons in the networks ( 512 neurons as default ) , W is connectivity matrix , τ is a timescale of update . The sensory input f ( zt ) is fed through connections matrixWs , and action feedback is fed throughWa . The context signal Ct is fed through matrix Wc . The network outputs a Q function , which is computed by a linear transformation of its hidden state . Beyond the basic setting ( Fig . 1A ) , we design several variants of the task to emphasize different statistical regularities ( Fig . 1C ) . In all cases , the agent begins from a random position and has to reach the context-dependent reward location in the shortest time using only proximal input . The ” Hole ” variant introduces a random placement of obstacles ( different numbers and positions ) in each trial . The ” Bar ” variant introduces a horizontal or vertical bar in random positions in each trial . The various ” Scale ” tasks stretch the arena in the horizontal or vertical direction while maintaining the relative position of the rewards . The ” Implicit context ” task is similar to the basic setting , but the external context input is eliminated , and instead , the color of the walls indicates the reward position . For all these tasks , the agent needs to find a strategy that tackles the uncertain elements to achieve the goals . Despite the simple setting of the game , the tasks are not trivial due to identical visual inputs in most of the locations and various uncertain elements adding to the task difficulty . 3.2 TRAINING FRAMEWORK . We aim to understand the interaction between internal representation and the statistical regularities of the various tasks . In principle , this could be accomplished by end-to-end reinforcement learning of many tasks , using various hyper-parameters to allow different solutions to the same task . We opted for a different approach - both due to computation efficiency ( see Appendix III ) and due to biological motivations . A biological agent acquires navigation ability during evolution and development , which shapes its elementary cognitive ability such as spatial or object memory . This shaping provides a scaffold upon which the animal could adapt and learn quickly to perform diverse tasks during life . Similarly , we divide learning into two phases , a pre-training phase that is task agnostic and a Q learning phase that is task-specific ( Fig . 2A ) . During pre-training we modify the network ’ s internal and input connectivity , while Q learning only modifies the output . Pre-training is implemented in an environment similar to the basic task , with an arena size chosen randomly between 10 to 20 . The agent ’ s actions are externally determined as a correlated random walk , instead of being internally generated by the agent . Inspired by neurophysiological findings , we emphasize two different aspects of internal representation - landmark memory ( Identity of the last encountered wall ) and position encoding ( O ’ Keefe & Nadel , 1978 ) . We thus pre-train the internal connectivity to generate an ensemble of networks with various hyperparameters that control the relative importance of these two aspects , as well as which parts of the connectivityW , Wa , Wiare modified . We term networks emphasizing the two aspects respectively MemNet and PosNet , and call the naive random network RandNet ( Fig . 2A ) . This is done by stochastic gradient descent on the following objective function : S = −α n∑ i=1 P̂ ( zt ) logP ( zt ) − β n∑ i=1 ÎtlogP ( It ) − γ n∑ i=1 ÂtlogP ( At ) ( 3 ) with z = ( x , y ) for position , I for landmark memory ( identity of the last wall encountered ) , A for action . The term on action serves as a regularizer . The three probability distributions are estimated from hidden states of the RNN , given by : P ( I|ht ) = exp ( Wmht + bm ) ∑ m ( exp ( Wmht + bm ) ) ( 4 ) P ( A|ht−1 , ht ) = exp ( Wa [ ht−1 , ht ] + ba ) ∑ a exp ( Wa [ ht−1 , ht ] + ba ) ( 5 ) P ( z|ht ) = exp ( ( z − ( Wpht + bp ) ) 2/σ2 ) ∑ z exp ( ( z − ( Wpht + bp ) ) 2/σ2 ) ( 6 ) where Wm , Wp , Wa are readout matrices from hidden states and [ ht−1 , ht ] denotes the concatenation of last and current hidden states . Tables 1,2,3 in the Appendix show the hyperparameter choices for all networks . The ratio between α and β controls the tradeoff between position and memory . The exact values of the hyperparameters were found through trial and error . Having obtained this ensemble of networks , we use a Q-learning algorithm with TD-lambda update for the network ’ s outputs , which are Q values . We utilize the fact that only the readout matrix Wo is trained to use a recursive least square method which allows a fast update of weights for different tasks ( Sussillo & Abbott , 2009 ) . This choice leads to a much better convergence speed when compared to stochastic gradient descent . The update rule used is : Wo ( n+ 1 ) = Wo ( n ) − e ( n ) P ( n ) H ( n ) T ( 7 ) P ( n+ 1 ) = ( C ( n+ 1 ) + αI ) −1 ( 8 ) C ( n+ 1 ) = λC ( n ) +H ( n ) TH ( n ) ( 9 ) e ( n ) = WoH ( n ) − Y ( n ) ( 10 ) whereH is a matrix of hidden states over 120 time steps , αI is a regularizer and λ controls forgetting rate of past data . We then analyze the test performance of all networks on all tasks ( Figure 2B and Table 3 in appendix ) . Figure 2B , C show that there are correlations between different tasks and between different networks . We quantify this correlation structure by performing principal component analysis of the performance matrix . We find that the first two PCs in task space explain 79 % of the variance . The first component corresponds to the difficulty ( average performance ) of each task , while the coefficients of the second component are informative regarding the nature of the tasks ( Fig . 2B , right ) : Bar ( -0.49 ) , Hole ( -0.25 ) , Basic ( -0.21 ) , Implicit context ( -0.12 ) , ScaleX ( 0.04 ) , ScaleY ( 0.31 ) , Scale ( 0.74 ) . We speculate these numbers characterize the importance of two different invariances inherent in the tasks . Negative coefficients correspond to metric invariance . For example , when overcoming dynamic obstacles , the position remains invariant . This type of task was fundamental to establish metric cognitive maps in neuroscience ( O ’ Keefe & Nadel , 1978 ) . Positive coefficients correspond to topological invariance , defined as the relation between landmarks unaffected by the metric information . Observing the behavior of networks for the extreme tasks of this axis indeed confirms the speculation . Fig . 3A shows that the successful agent overcomes the scaling task by finding a set of actions that captures the relations between landmarks and reward , thus generalizing to larger size arenas . Fig3B shows that the successful agent in the bar task uses a very different strategy . An agent that captures the metric invariance could adjust trajectories and reach the reward each time when the obstacle is changed . This ability is often related to the ability to use shortcuts ( O ’ Keefe & Nadel , 1978 ) . The other tasks intepolate between the two extremes , due to the presence of both elements in the tasks . For instance , the implicit context task requires the agent to combine landmark memory ( color of the wall ) with position to locate the reward . We thus define metric and topological scores by using a weighted average of task performance using negative and positive coefficients respectively . Fig . 3C shows the various networks measured by the two scores . We see that random networks ( blue ) can achieve reasonable performance with some hyperparameter choices , but they are balanced with respect to the metric topological score . On the other hand , PostNet networks are pushed to the metric side and MemNet networks to the topological side . This result indicates that the inductive bias achieved via task agnostic pre-training is manifested in the performance of networks on various navigation tasks .

Summary: This paper studies the internal representations of recurrent neural networks trained on navigation tasks. By varying the weight of different terms in an objective used for supervised pre-training, RNNs are created that either use path integration or landmark memory for navigation. The paper shows that the pretraining method leads to differential performance when the readout layer of these networks networks is trained using Q-learning on different variants of a navigation task. The main result of the paper is obtained by finding the slow points of the dynamics of the trained RNNs. The paper finds that the RNNs pre-trained to use path integration contain 2D continuous attractors, allowing position memory. On the other hand, the RNNs pre-trained for landmark memory contain discrete attractors corresponding to the different landmarks.