A Consciousness-Inspired Planning Agent for Model-Based Reinforcement Learning

1 Introduction . Whether when planning our paths home from the office or from a hotel to an airport in an unfamiliar city , we typically focus on a small subset of relevant variables , e.g . the change in position or the presence of traffic . An interesting hypothesis of how this path planning skill generalizes across scenarios is that it is due to computation associated with the conscious processing of information [ 2 , 3 , 14 ] . Conscious attention focuses on a few necessary environment elements , with the help of an internal abstract representation of the world [ 43 , 14 ] . This pattern , also known as consciousness in the first sense ( C1 ) [ 14 ] , has been theorized to enable humans ’ exceptional adaptability and learning efficiency [ 2 , 3 , 14 , 43 , 7 , 15 ] . A central characterization of conscious processing is that it involves a bottleneck , which forces one to handle dependencies between very few environmental characteristics at a time [ 14 , 7 , 15 ] . Though focusing on a subset of the available information may seem limiting , it facilitates Out-Of-Distribution ( OOD ) and systematic generalization to other situations where the ignored variables are different and yet still irrelevant [ 7 , 15 ] . In this paper , we encode some of these ideas into reinforcement learning agents . Reinforcement learning ( RL ) is an approach for learning behaviors from agent-environment interactions [ 41 ] . However , most of the big successes of RL have been obtained by deep , model-free agents [ 30 , 37 , 38 ] . While Model-Based RL ( MBRL ) has generated significant research due to the potentials of using an extra model [ 31 ] , its empirical performance has typically lagged behind , with some recent notable exceptions [ 36 , 24 , 17 ] . Our proposal is to take inspiration from human consciousness to build an architecture which learns a useful state space and in which attention can be focused on a small set of variables at any time , where the aspect of “ partial planning ” 1 is enabled by modern deep 1Partial planning is interpreted in different ways . For example , concurrent work [ 26 ] focuses on modelling “ affordable ” temporally extended actions , s.t . an “ intent ” could be achievedmore efficiently . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . RL techniques [ 42 , 26 ] . Specifically , we propose an end-to-end latent-space MBRL agent which does not require reconstructing the observations , as in most existing works , and uses Model Predictive Control ( MPC ) framework for decision-time planning [ 34 , 35 ] . From an observation , the agent encodes a set of objects as a state , with a selective attention bottleneck mechanism to plan over selected subsets of the state ( Sec . 4 ) . Our experiments show that the inductive biases improve a specific form of OOD generalization , where consistent dynamics are preserved across seemingly different environment settings ( Sec . 5 ) . 2 Background & Context . We consider an agent interacting with its environment at discrete timesteps . At time C , the agent receives observation > C and takes action 0C , receiving a reward AC+1 and new observation > C+1 . The interaction is episodic . The agent is also building a latent-space transition model , ℳ , which can be used to sample a next state , B̂C+1 , a reward ÂC+1 and a binary signal $ ̂C+1 which indicates if the model predicts termination after the transition . We will now compare and contrast our approach with some existing methods from the MBRL literature , explaining the rationale for our design choices . Observation Level Planning and Reconstruction vs Latent Space Planning Many MBRL methods plan in the observation space or rely on reconstruction-based losses to obtain state representations [ 24 , 36 , 17 , 48 ] . Appropriate as these methods may be for some robotic tasks with few sensory inputs , e.g . continuous control with joint states , they are arguably difficult with high-dimensional inputs like images , since they may focus on predictable yet useless aspects of the raw observations [ 31 ] . Besides suffering from the need to reconstruct noise or irrelevant parts of the signal , it is not clear if representations built by a reconstruction loss ( e.g . ! 2 in the observation space ) are effective for an MBRL agent to plan or predict the desired signals [ 39 , 17 , 18 ] , e.g . values ( in the RL sense ) , rewards , etc .. In this work , we use an approach similar to those in [ 39 , 36 , 17 ] , building a latent space representation that is jointly shaped by all the relevant RL signals ( to serve value estimation and planning ) without using reconstruction . Staged Training vs End-to-End Training Some MBRL agents based on a world model [ 16 , 24 , 31 ] use two explicit stages of training : ( 1 ) an inner representation of the world is trained using exploration ( usually with random trajectories ) ; ( 2 ) the representation is fixed and used for planning and MBRL . Despite the advantages of being more stable and easier to train , this procedure relies on having an environment where the initial exploration provides transitions that are sufficiently similar to those observed under improved policies , which is not the case in many environments . Furthermore , the learned representation may not be effective for value estimation , if these transitions do not contain reward information that can be used to update the input-torepresentation encoder . End-to-end MBRL agents , e.g . [ 39 , 36 ] , are able to learn the representation online , simultaneously with the value function , hence adapting better to non-stationarity in the transition distribution and rewards . Type of planning MBRL agents can use the model in different ways . Dyna [ 40 ] learns a model to generate “ imaginary ” transitions , which contribute to the training of the value estimator [ 40 ] , in addition to the real observations , thus boosting sample efficiency . However , if the model is inaccurate , the transitions it generates may be “ delusional '' , which may alter the value estimator and negatively impact performance . Moreover , Dyna is typically used to generate extra transitions from the states visited in a trajectory , and updates the model based on the observed transitions as well . This means Dyna is focused on the data distribution encountered by the agent and may have trouble generalizing OOD . In contrast , simulationbased model-predictive control ( MPC ) and its variants [ 34 , 35 , 18 ] only update the value estimator based on real data , using the model simply to perform lookahead at decision-time . Hence , model inaccuracies impact less , with more favorable OOD generalization capabilities . Hence , MPC is adopted in our approach . Vectorized vs Set Representations for RL Most Deep Reinforcement Learning ( DRL ) work focus on learning vectorized state rep- resentations , where the agents ’ observation is transformed into a feature vector of fixed dimensionality [ 30 , 19 ] . Instead , set-based encoders , a.k.a . object-oriented architectures , are designed to extract a set of unordered vectors from which to predict the desired signals via permutation-invariant computations [ 50 ] , as illustrated in Fig . 1 . Recent works in RL have shown the promise of set-based representations in capturing environmental states , in terms of generalization , as well as their similarities to human perception [ 13 , 47 , 32 , 46 , 29 ] . Additionally in this work , we utilize the compositionality of set representations to enable the discovery of sparse interactions among objects , i.e . underlying dynamics , as well as to facilitate the bottleneck mechanism , analogous to C1 selection . The set-based representation coupled with the bottleneck provides an inductive bias consistent with selecting only the relevant aspects of a situation on-the-fly through an attention mechanism . The small size of the working memory bottleneck also enforces sparsity of the dependencies [ 7 , 15 ] captured by the learned dynamics model : each transition can only relate a few objects together , no more than the size of the bottleneck . 3 MBRL with Set Representations . We present an end-to-end baseline MBRL agent that uses a set-based representation and carries out latent space planning , but without a consciousness-inspired small bottleneck . This agent serves as a baseline to investigate the OOD generalization capabilities brought by the bottleneck , which is to be introduced later in Sec . 4 . The mapping from observations to values is a combination of an encoder and a value estimator . The encoder maps an observation vector to a set of objects , which constitutes the latent state . The value estimator is a permutation-invariant set-to-vector architecture that maps the latent state to a value estimate . Note that the same state set is used for all the agents ’ predictions , including future states , rewards etc. , as we will discuss later . Encoder . For image-based observations , we use the features at each position of the CNN output feature map to characterize the feature of an object , similar to [ 9 ] , as shown in Figure 1 . To recover positional information lost during the process , we concatenate each object feature vector with a positional embedding to form a complete object embedding . Such approach is different from the common practice of mixing positional information by addition [ 45 ] . This is for the compatibility with our dynamics model training procedure , discussed below . ( State-Action ) Value Estimator takes the form & : S → R|A| , where S is the learned state space by the set-based encoder ( hoping to capture the real underlying state space of the MDP ) and A is a discrete action set . We use an improved architecture upon DeepSets [ 50 ] , depicted in Figure 2 . The architecture performs reasoning on a set of encoded objects , resembling pervasive usage in natural language processing , where the objects are typically word tokens [ 33 ] . Transition Model . The transition model maps from BC , 0C to B̂C+1 , ÂC and $ ̂C+1 . We separate this into : 1 ) the dynamics model , in charge of simulating how the state would change with the input of 0C and 2 ) the reward-termination estimator which maps BC , 0C to ÂC and $ ̂C+1 . While designing reward-termination estimator is straightforward ( a two-headed augmented architecture similar to the value estimator ) , the dynamics model requires regression on unordered sets of objects ( set-to-set ) . A common approach is to use matching methods , e.g . Chamfer matching or Hausdorff distance , However , they are computationally demanding and subject to local optima [ 5 , 8 , 28 ] . Targeting this , our feature-position separated set encoding not only makes the permutation-invariant computations position-aware , but also allows simple end-to-end training over the dynamics . By forcing the positional tails to be immutable during the computational pass , we can use them to solve the matching trivially : objects “ labeled ” with the same positional tail in the prediction B̂C+1 ( output of the dynamics model ) and the training sample BC+1 ( state obtained from the next observation ) are aligned , forming pairs of objects with changes only in the feature , as shown in Figure 3 . Tree Search MPC . The agent employs a tree-search based behavior policy ( with & -greedy exploration ) . During planning , each tree search call maintains a priority queue of branches to simulate with the model . When a designated budget ( e.g . number of steps of simulation ) is spent , the agent greedily picks the immediate action that leads to the most promising path . We present the pseudocode of the Q-value based prioritized tree-search MPC in Appendix . Equivalence could be drawn from this planning approach to Monte-Carlo Tree Search ( MCTS ) [ 37 , 38 ] . While this method is far more simplistic and require fewer simulations for each planning call ( see example in Appendix ) . Training . The proposed agent is trained from sampled transitions with the following losses : • Temporal Difference ( TD ) ℒTD : regresses the current value estimate to the update target , e.g . calculated according to DQN or Double DQN ( DDQN ) [ 30 , 44 ] . In experiments , a distributional output is used for both value and reward estimation , making this loss a KL-divergence [ 6 ] . • Dynamics Consistency ℒdyn : A ! 2 penalty established between the aligned B̂C+1 and BC+1 , where B̂C+1 is the imagined next ( latent ) state given > C , 0C and BC+1 is the true next ( latent ) state encoded from > C+1 . • Reward Estimation ℒA : the KL-divergence between the imagined reward ÂC+1 predicted by the model and the true reward AC+1 of the observed transition . • Termination Estimation ℒ $ : the binary cross-entropy loss from the imagined termination $ ̂C+1 to the ground truth $ C+1 , obtained from environment feedback . The resulting total loss for end-to-end training of this set-based MBRL agent is thus2 : ℒ = ℒTD + ℒdyn + ℒA + ℒ $ Jointly shaping the states avoids the representation collapsing to trivial solutions and makes the representation useful for all signal predictions of interest .

Summary: The authors introduce an architecture for focusing agents in model-based reinforcement learning on just the relevant subset of the environment representation for making decisions (reward, termination, and successor state prediction). The approach leans heavily on recent advances in set-based representations and the approach is tested in the MiniGrid-BabyAI framework. Results show promising behaviour, particularly in the ability to generalize across different settings.