Skip to yearly menu bar Skip to main content

Workshop: Deep Reinforcement Learning

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

Mingde Zhao · Zhen Liu · Sitao Luan · Shuyuan Zhang · Doina Precup · Yoshua Bengio


We present an end-to-end, model-based deep reinforcement learning agent which dynamically attends to relevant parts of its state, in order to plan and to generalize better out-of-distribution. The agent's architecture uses a set representation and a bottleneck mechanism, forcing the number of entities to which the agent attends at each planning step to be small. In experiments, we investigate the bottleneck mechanism with sets of customized environments featuring different dynamics. We consistently observe that the design allows agents to learn to plan effectively, by attending to the relevant objects, leading to better out-of-distribution generalization.

Chat is not available.