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

arXiv: https://arxiv.org/abs/2106.02097 · HTML: https://mlsanity.com/paper/2106.02097/

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 65 references that are in the map.
- First posted on arXiv 3 Jun 2021.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Inductive Biases for Deep Learning of Higher-Level Cognition](https://mlsanity.com/paper/2011.15091/) (2020)
- [Distributed Prioritized Experience Replay](https://mlsanity.com/paper/1803.00933/) (2018)
- [BabyAI: A Platform to Study the Sample Efficiency of Grounded Language Learning](https://mlsanity.com/paper/1810.08272/) (2018)
- [Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control](https://mlsanity.com/paper/1812.00568/) (2018)
- [How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks](https://mlsanity.com/paper/2009.11848/) (2020)
- [Benchmarking Model-Based Reinforcement Learning](https://mlsanity.com/paper/1907.02057/) (2019)
- [The Consciousness Prior](https://mlsanity.com/paper/1709.08568/) (2017)
- [The Predictron: End-To-End Learning and Planning](https://mlsanity.com/paper/1612.08810/) (2016)
- [Look Before You Leap: Bridging Model-Free and Model-Based Reinforcement Learning for Planned-Ahead Vision-and-Language Navigation](https://mlsanity.com/paper/1803.07729/) (2018)
- [Learning and Planning in Complex Action Spaces](https://mlsanity.com/paper/2104.06303/) (2021)
- [On the role of planning in model-based deep reinforcement learning](https://mlsanity.com/paper/2011.04021/) (2020)
- [Sparse Attentive Backtracking: Temporal CreditAssignment Through Reminding](https://mlsanity.com/paper/1809.03702/) (2018)
- [Conditional Set Generation with Transformers](https://mlsanity.com/paper/2006.16841/) (2020)

## Built on by (most cited first, 27 of 27)

- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [RL Fine-Tuning Heals OOD Forgetting in SFT](https://mlsanity.com/paper/2509.12235/) (2025)
- [A Theory of Consciousness from a Theoretical Computer Science Perspective: Insights from the Conscious Turing Machine](https://mlsanity.com/paper/2107.13704/) (2021)
- [NAVIX: Scaling MiniGrid Environments with JAX](https://mlsanity.com/paper/2407.19396/) (2024)
- [On the link between conscious function and general intelligence in humans and machines](https://mlsanity.com/paper/2204.05133/) (2022)
- [Dynamic Model Predictive Shielding for Provably Safe Reinforcement Learning](https://mlsanity.com/paper/2405.13863/) (2024)
- [Temporally Abstract Partial Models](https://mlsanity.com/paper/2108.03213/) (2021)
- [SPARTAN: A Sparse Transformer World Model Attending to What Matters](https://mlsanity.com/paper/2411.06890/) (2024)
- [Representation Learning on Heterophilic Graph with Directional Neighborhood Attention](https://mlsanity.com/paper/2403.01475/) (2024)
- [Translate Policy to Language: Flow Matching Generated Rewards for LLM Explanations](https://mlsanity.com/paper/2502.12530/) (2025)
- [A Look at Value-Based Decision-Time vs. Background Planning Methods Across Different Settings](https://mlsanity.com/paper/2206.08442/) (2022)
- [Towards Empowerment Gain through Causal Structure Learning in Model-Based RL](https://mlsanity.com/paper/2502.10077/) (2025)
- [Compositional Generalization in Grounded Language Learning via Induced Model Sparsity](https://mlsanity.com/paper/2207.02518/) (2022)
- [Reusable Slotwise Mechanisms](https://mlsanity.com/paper/2302.10503/) (2023)
- [Consciousness-Inspired Spatio-Temporal Abstractions for Better Generalization in Reinforcement Learning](https://mlsanity.com/paper/2310.00229/) (2023)
- [Revealing the Pitfalls and Re-Evaluating the Advancement of Heterophilic Graph Learning](https://mlsanity.com/paper/2409.05755/) (2024)
- [Flexible Diffusion Scopes with Parameterized Laplacian for Heterophilic Graph Learning](https://mlsanity.com/paper/2409.09888/) (2024)
- [Boolean Decision Rules for Reinforcement Learning Policy Summarisation](https://mlsanity.com/paper/2207.08651/) (2022)
- [Minimal Value-Equivalent Partial Models for Scalable and Robust Planning in Lifelong Reinforcement Learning](https://mlsanity.com/paper/2301.10119/) (2023)
- [A Tractable Inference Perspective of Offline RL](https://mlsanity.com/paper/2311.00094/) (2023)
- [Deep Reinforcement Learning via Object-Centric Attention](https://mlsanity.com/paper/2504.03024/) (2025)
- [On Addressing the Limitations of Graph Neural Networks](https://mlsanity.com/paper/2306.12640/) (2023)
- [Rejecting Hallucinated State Targets during Planning](https://mlsanity.com/paper/2410.07096/) (2024)
- [Inductive Biases for Zero-shot Systematic Generalization in Language-informed Reinforcement Learning](https://mlsanity.com/paper/2501.15270/) (2025)
- [Reinforcement Learning with Brain-Inspired Modulation can Improve Adaptation to Environmental Changes](https://mlsanity.com/paper/2205.09729/) (2022)
- [DELTA: Dynamic Embedding Learning with Truncated Conscious Attention for CTR Prediction](https://mlsanity.com/paper/2305.04891/) (2023)
- [Partial Models for Building Adaptive Model-Based Reinforcement Learning Agents](https://mlsanity.com/paper/2405.16899/) (2024)

## Abstract (arXiv)

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 different 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.