DAG-NAS : An explainable neural architecture search framework for reinforcement learning.

Abstract

We present an explainable Neural Architecture Search (NAS) framework for Reinforcement Learning (RL). We model a feed-forward neural network as a Directed Acyclic Graph (DAG) that consists of scalar-level operations and their interconnections. Scalar-level DAGs are trained using a differentiable search method, followed by pruning search results. This approach yields a compact neural architecture that delivers high performance while enhancing explainability by highlighting the critical information needed to solve the problem. We apply our NAS framework to search both actor and critic networks of the Actor-Critic PPO algorithm across various RL tasks. Extensive experiments demonstrate that our architectures achieve comparable performance with significantly fewer parameters while highlighting key features and enhancing explainability.