Brain-Inspired Planning for Better Generalization in Reinforcement Learning

TL;DR

This paper introduces brain-inspired planning mechanisms to improve the generalization of reinforcement learning agents by incorporating reasoning behaviors, spatial abstraction, task decomposition, and safety measures against hallucinated targets.

Contribution

It presents novel planning strategies inspired by human cognition, including spatial abstraction, the Skipper framework for task decomposition, and a feasibility evaluator to enhance RL generalization and safety.

Findings

Significant improvement in systematic generalization outside training tasks

Robustness against distributional shifts in complex environments

Enhanced safety by rejecting hallucinated, infeasible targets

Abstract

Existing Reinforcement Learning (RL) systems encounter significant challenges when applied to real-world scenarios, primarily due to poor generalization across environments that differ from their training conditions. This thesis explores the direction of enhancing agents' zero-shot systematic generalization abilities by granting RL agents reasoning behaviors that are found to help systematic generalization in the human brain. Inspired by human conscious planning behaviors, we first introduced a top-down attention mechanism, which allows a decision-time planning agent to dynamically focus its reasoning on the most relevant aspects of the environmental state given its instantaneous intentions, a process we call "spatial abstraction". This approach significantly improves systematic generalization outside the training tasks. Subsequently, building on spatial abstraction, we developed the Skipper framework to automatically decompose complex tasks into simpler, more manageable sub-tasks. Skipper provides robustness against distributional shifts and efficacy in long-term, compositional planning by focusing on pertinent spatial and temporal elements of the environment. Finally, we identified a common failure mode and safety risk in planning agents that rely on generative models to generate state targets during planning. It is revealed that most agents blindly trust the targets they hallucinate, resulting in delusional planning behaviors. Inspired by how the human brain rejects delusional intentions, we propose learning a feasibility evaluator to enable rejecting hallucinated infeasible targets, which led to significant performance improvements in various kinds of planning agents. Finally, we suggest directions for future research, aimed at achieving general task abstraction and fully enabling abstract planning.