A Consciousness-Inspired Planning Agent for Model-Based Reinforcement Learning
Abstract
We present an end-to-end, model-based deep reinforcement learning agent which dynamically attends to relevant parts of its state during planning.
The agent uses a bottleneck mechanism over a set-based representation to force the number of entities to which the agent attends at each planning step to be small.
In experiments, we investigate the bottleneck mechanism with several sets of customized environments featuring diﬀerent challenges.
We consistently observe that the design allows the planning agents to generalize their learned task-solving abilities in compatible unseen environments by attending to the relevant objects, leading to better out-of-distribution generalization performance.
Introduction
Whether when planning our paths home from the oﬃce 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 traﬃc.
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 ﬁrst sense (C1) [14], has been theorized to enable humans’ exceptional adaptability and learning eﬃciency [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 diﬀerent 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 signiﬁcant 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
RL techniques [42, 26].
Speciﬁcally, 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 speciﬁc form of OOD generalization, where consistent dynamics are preserved across seemingly diﬀerent environment settings (Sec. 5).
Background & Context
We consider an agent interacting with its environment at discrete timesteps.
At time 𝑡 , the agent receives observation 𝑜 𝑡 and takes action 𝑎 𝑡 , receiving a reward 𝑟 𝑡 + 1 and new observation 𝑜 𝑡 + 1 .
The interaction is episodic.
The agent is also building a latent-space transition model, ℳ , which can be used to sample a next state, ˆ 𝑠 𝑡 + 1 , a reward ˆ 𝑟 𝑡 + 1 and a binary signal ˆ 𝜔 𝑡 + 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 diﬃcult with high-dimensional inputs like images, since they may focus on predictable yet useless aspects of the raw observations [31].
Besides suﬀering 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 eﬀective 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 ﬁxed 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 suﬃciently similar to those observed under improved policies, which is not the case in many environments.
Furthermore, the learned representation may not be eﬀective 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 diﬀerent 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 eﬃciency.
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-
[Figure]
vector vectorized representation [Figure] resentations, where the agents’ observation is transformed into a feature vector of ﬁxed 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.
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-ﬂy 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.
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 diﬀerent 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 𝑄 :
𝒮 → R |
𝒜
| , where 𝒮 is the learned state space by the set -based encoder (hoping to capture the real underlying state space of the MDP
) and
𝒜
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]. [Figure] S A
[Equation] F C [Figure]
m e a n [Figure] Transition Model. The transition model maps from 𝑠 𝑡 , 𝑎 𝑡 to ˆ 𝑠 𝑡 + 1 , ˆ 𝑟 𝑡 and ˆ 𝜔 𝑡 + 1 .
We separate this into: 1) the dynamics model , in charge of simulating how the state would change with the input of 𝑎 𝑡 and 2) the reward-termination estimator which maps 𝑠 𝑡 , 𝑎 𝑡 to ˆ 𝑟 𝑡 and ˆ 𝜔 𝑡 + 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 Hausdorﬀ 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 ˆ 𝑠 𝑡 + 1 (output of the dynamics model) and the training sample 𝑠 𝑡 + 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 Diﬀerence (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 ˆ 𝑠 𝑡 + 1 and 𝑠 𝑡 + 1 , where ˆ 𝑠 𝑡 + 1 is the imagined next (latent) state given 𝑜 𝑡 , 𝑎 𝑡 and 𝑠 𝑡 + 1 is the true next (latent) state encoded from 𝑜 𝑡 + 1 . • Reward Estimation ℒ 𝑟 : the KL-divergence between the imagined reward ˆ 𝑟 𝑡 + 1 predicted by the model and the true reward 𝑟 𝑡 + 1 of the observed transition.
• Termination Estimation ℒ 𝜔 : the binary cross-entropy loss from the imagined termination ˆ 𝜔 𝑡 + 1 to the ground truth 𝜔 𝑡 + 1 , obtained from environment feedback. [Figure] A
F C imaginedfeature [Figure]
The resulting total loss for end-to-end training of this set-based MBRL agent is thus 2 :
[Equation]
Jointly shaping the states avoids the representation collapsing to trivial solutions and makes the representation useful for all signal predictions of interest.
Consciousness-Inspired Bottleneck
In this section, we introduce an inductive bias which facilitates C1-capable planning.
In a nutshell, the planning is expected to focus on the parts of the world that matter for the plan.
Simulations and predictions are all expected to be performed on a (small) bottleneck set, which contains all the important transition-related information.
As illustrated in Figure 4, the model performs 1) selection of the bottleneck set from the full state-set, 2) dynamics simulation on the bottleneck set and 3) integration of predicted bottleneck set to form the predicted next state.
Conditional State Selection We select a bottleneck set 𝑐 𝑡 of 𝑛 objects from the potentially large state set 𝑠 𝑡 of 𝑚 (cid:29) 𝑛 objects. Then we only model the transition for the selected objects in 𝑐 𝑡 .
To make this selection, we use a key-query-value attention mechanism, where the key and the value for each object in 𝑠 𝑡 are obtained from that object, and the query is a function of some learned dedicated set of vectors and of the action considered (see Appendix for details).
Inspired by the work on self-attention for memory access [25], we use a semi-hard top- 𝑘 attention mechanism to facilitate the selection of the bottleneck set.
That is, after the query, the top- 𝑘 attention weights are kept, all others are set to 0 , and then the attention weights are renormalized.
This semi-hard attention technique limits the inﬂuence of the ill-matched objects on the bottleneck set 𝑐 𝑡 while allowing for a gradient to propagate on the assignment of relative weight to diﬀerent objects.
With purely soft attention, weights for irrelevant objects are never 0 and learning to disentangle objects may be more diﬃcult.
Dynamics / Reward-Termination Prediction on Bottleneck Sets.
We use the same architecture as described in Sec. 3, but taking the bottleneck objects as input rather than the full state set.
Details of the architecture are in the Appendix.
Change Integration.
An integration operation, intuitively the inverse operation of selection, is implemented to ‘soft paste-back’ the changes of the bottleneck state onto the state set 𝑠 𝑡 , [Figure] yielding the imagined next state set ˆ 𝑠 𝑡 +
This is also achieved by attention operations, more speciﬁcally querying ˆ 𝑐 𝑡 + 1 with 𝑠 𝑡 , conditioned on the action 𝑎 𝑡 .
Please check the Appendix for more details.
Discussion.
The bottleneck described in this section is a natural complement to the MBRL model with set representations discussed previously.
In particular, planning and training are carried out the same way as discussed in Sec. 3.
We expect the Conscious Planning (CP) agent to demonstrate the following advantages:
• Higher Quality Representation: the interplay between the set representation and the selection / integration forces the representation to be more disentangled and more capable of capturing the locally sparse dynamics.
• More Eﬀective Generalization: only essential objects for the purpose of planning participate in the transition, thus generalization should be improved both in-distribution and OOD, because the transition does not depend on the parts of the state ignored by the bottleneck.
• Lower Computational Complexity: directly employing transformers to simulate the full state dynamics results in a complexity of 𝒪(| 𝑠 𝑡 | 2 𝑑 ) , where 𝑑 is the length of the objects, due to the use of Self-Attention (SA), while the bottleneck lowers it to 𝒪(| 𝑠 𝑡 || 𝑐 𝑡 | 𝑑 ) . (a) In-dist, diﬀ 0 .
[Figure]
Experiments
We present our experimental settings and ablation studies of our CP agent against baselines to investigate the OOD generalization capabilities enabled by the C1-inspired bottleneck mechanism.
To clarify, the OOD generalization we refer to speciﬁcally is the agents’ ability to generalize its learned task skills across seemingly diﬀerent tasks with common underlying dynamics .
Take the set of experiments in this section for example, we want the agent to be able to generalize its navigation skills in unseen environments.
5.1 Environment / Task Description
We use environments based on the MiniGrid-BabyAI framework [11, 10, 21], which can be customized for generating OOD generalization tests with varying diﬃculties.
To make sure we assess the agents as clearly as possible, the customized environments feature clear object deﬁnitions, with well-understood underlying dynamics based on object interactions.
Furthermore, the environments are solvable by Dynamic Programming (DP) and can be easily tuned to generate OOD evaluation tasks.
These characteristics are crucial for the experimental insights we are seeking.
In this section, the experiments are carried out on 8 × 8 gridworlds 3 , as shown in Figure 5.
The agent (red triangle) needs to navigate (by turning left, right or stepping forward) to the goal while dodging the lava cells along the way 4 .
If the agent steps into lava (orange square), the episode terminates immediately with no reward.
If the agent successfully reaches the goal (green square), it receives a reward of + 1 and the episode terminates.
For better generalization, the agent needs to understand how to avoid lava in general (and not at speciﬁc locations, since their placement changes) and to reach the goal as quickly as possible 5 .
The environments provide grid-based observations that are ready to be interpreted as set representations: each cell of the observation array is an object, thus resulting in a set ofobjects in 𝑠 𝑡 for each observation.
For the agent to be able to understand the environment dynamics instead of memorizing speciﬁc task layouts, we generate a new environment for each training or evaluation episode.
In each training episode, the agent starts at a random position on the leftmost or rightmost edge and the goal is placed randomly somewhere along the opposite edge.
In between the two edges, the lava cells are randomly generated according to a diﬃculty parameter which controls the probability of placing a lava cell at each valid position.
The diﬃculty parameter controls partially how seemingly diﬀerent the OOD evaluation tasks are to the in-distribution training tasks, though we know the underlying dynamics of all these tasks are the same.
For training episodes, the diﬃculty is ﬁxed to 0 .
We note that most usual RL benchmarks contain ﬁxed environments, where the agent is expected to acquire a speciﬁc optimal policy.
These environments are ill-suited for our purpose.
For OOD evaluation, the agent is expected to adapt in new tasks with the same underlying dynamics in a 0 -shot fashion, i.e. with the agent’s parameters ﬁxed.
The OOD tasks are crafted to include changes both in the support (orientation) and in the distribution (diﬃculty): the agent is deployed in transposed layouts 6 with varying levels of diﬃculty ( { 0 . 25 , 0.35 , 0 . 45 , 0 . 55 } ).
The diﬀerences of in-distribution (training) and OOD (evaluation) environments are illustrated in Figure 5.
5.2 Agent Setting
We build all the set-based MBRL agents included in the evaluation on a common model-free baseline: a set-based variant of Double-DQN (DDQN) [44] with prioritized replay and distributional outputs.
For more details, please check the Appendix.
We compare the proposed approach, labelled CP in the ﬁgures (for Conscious Planning) against the following methods:
• UP (for Unconscious Planning): the agent proposed in Section 3, lacking the bottleneck. • model-free : the model-free set-based agent is the basis for the set-based model-based agents.
It consists of only the encoder and the value estimator, sharing their architectures with CP and UP. • Dyna : the set-based MBRL agent which includes a model-free agent and an observationlevel transition model, i
e. a transition generator. For the model, we use the CP transition model (with the same hyperparameters as the best performing CP agent) on the original environment features without an encoder.
We also use the same hyperparameters as in the CP model training.
The agent essentially doubles the batch size of the model-free baseline by augmenting training batches with an equal number of generated transitions. • Dyna* : A Dyna baseline that uses the true environment model for transition generation.
This is expected to demonstrate Dyna’s performance limit. • WM-CP : A world model CP variant that diﬀers by following a 2 -stage training procedure [16]. First, the model (together with the encoder) is trained with 10 6 random transitions.
After this, the encoder and the model are ﬁxed and RL begins.
• NOSET :
A UP-counterpart with vectorized representations and no bottleneck mechanism.
Particularly, for CP and UP agents, we also test the following variants: • CP-noplan : A CP agent that trains normally but does not plan in OOD evaluations, i.e. carrying out model-free behavior.
This baseline aims to demonstrate the impact of planning in the training process on the OOD capability of the value estimator. • UP-noplan : UP counterpart of CP-noplan.
Note that the compared methods share architectures as much as possible to ensure fair comparisons.
Details of the compared methods, their design and hyperparameters are provided in the Appendix.
5.3 Performance Evaluation 5.3.1 In-Distribution
In Figure 6, we present the in-distribution evaluation curves for the diﬀerent agents.
For UP, CP and the corresponding model-free baselines, the performance curves show no signiﬁcant diﬀerence, which demonstrates that these agents are eﬀective in learning to solve the in-distribution tasks.
During the “warm-up” period of the WM baseline, the model learns a representation that captures the underlying dynamics.
After the warm-up, the encoder and the model parameters are ﬁxed and only the value estimator learns to predict the state-action values based on the given representation.
The increase in performance is not only delayed due to the warm-up phase (during which rewards are not taken into account) but also harmed, presumably because the value estimator has no ability to shape the representation to better suit its needs.
The Dyna baseline performs badly while the Dyna* baselines perform relatively well.
This is likely due to the delusional transitions generated by the model at the early stages of training, from which the value estimator never recovers.
However, the Dyna* baseline does not achieve satisfactory OOD performance (Figure 7), presumably because its planning only focuses on observed data, and hence only improves the in-distribution performance, due to insuﬃciently strong generalization.
The
NOSET baseline performs very badly even in-distribution, per Figure 6.
In the Appendix, we show that the NOSET baseline seems only able to perform well in a more classical, static RL setting, which may indicate that it relies on memorization.
We provide more results regarding the model accuracy in the Appendix.
5.3.2 OOD Task-Solving Performance
The OOD evaluation focuses on testing the agents’ performance in a set of environments forming a gradient of task diﬃculty.
In Figure 7, we present the performance error bars of [Figure] [Figure] the compared methods under diﬀerent OOD diﬃculty levels.
CP(8), CP with bottleneck size 𝑛 = 8 , shows a clear performance advantage over UP, validating the OOD generalization capability.
The Dyna* baseline, essentially the performance upper bound of Dyna-based planning methods, shows no signiﬁcant performance gain in OOD tests compared to modelfree methods.
WM may have the potential to reach similar performance as CP, yet it needs to warm up the encoder with a large portion of the agent-environment interaction budget, if no free unsupervised phase is provided.
We dive into this matter in the Appendix.
5.3.3 Ablation
We validate design choices with ablation.
Figure 8 visualizes two of these experiments.
For more ablation results, which include validation of the eﬀectiveness of diﬀerent model choices, and further quantitative measurements, e.g. of OOD ability as a function of behavior optimality and model accuracy, please check the Appendix.
5.4 Summary of Experimental Results
With the scope limited to our experiments, the results allow us to draw these conclusions:
• Set-based representations enable at least in-distribution generalization across diﬀerent environment instances in our non-static setting, where the agents are forced to discover dynamics that are preserved across environments; • Model-free methods seem to face more diﬃculties in solving our OOD evaluation tasks which preserved the same environment dynamics to the corresponding in-distribution training settings; • MPC exhibits better performance than Dyna in the tested OOD generalization settings; • Online joint training of the representation with all the relevant signals could bring beneﬁts to RL, as suggested in [22].
Please check Appendix E for more discussions of this matter; • In accordance with our intuition, transition models with bottlenecks tend to learn dynamics better in our tests.
This is likely for they prioritize learning the relevant aspects, while models without bottleneck may have to waste capacity on irrelevance;
agent-env interactions 10 0.
s u cc e ss r a t e
CP(8)noplan(8)
UPnoplan(UP) modelfree ( agent-env interactions 10 0.
s u cc e ss r a t e
CP(16)-best
CP(8)-best
CP(4)-best
CP(16)
CP(8)
CP(4)
(b) Value estimators do not generalize well in our OOD tests : random heuristic signiﬁcantly outperforms best-ﬁrst heuristic OOD.
• From further experiments provided in the Appendix E, we observe that bottleneckequipped agents may also be less aﬀected by larger environmental scales, possibly due to their prioritized learning of interesting entities.
Conclusion & Limitations
We introduced a conscious bottleneck mechanism into MBRL, facilitated by set-based representations, end-to-end learning and tree search MPC.
In the non-static RL settings, the bottleneck allows selecting the relevant objects for planning and hence enables signiﬁcant OOD performance.
One limitation of our work is the experimental focus on only Minigrid environments, due to the need to validate carefully our approach.
For future works, we would also like to extend these ideas to temporally extended models, which could simplify the planning task, and are also better suited as a conceptual model of C1.
Finally, we note that the architectures we use are involved and can require careful tuning for new types of environments.
Acknowledgements
Mingde is grateful for the ﬁnancial support from the Fonds de Recherche du QuébecNature et Technologies (FRQNT).
Yoshua acknowledges the ﬁnancial support from Samsung
Electronics and IBM.
We acknowledge the computational power provided by Compute Canada.
We are also thankful for the helpful discussions with Xiru Zhu (about the design of the environment generation procedure),
David Yu-Tung Hui (about the bag-of-word representations, insights on BabyAI as well as about the writing of the introduction section), Min Lin (about the design of the dynamics model as well as the early stage brainstorming) and Ian Porada (for consistently supporting the student authors).