Option Discovery via Differentiable Neural Decomposition

ICLR 2026 Conference Submission20581 Authors

Keywords: options, progammatic representations, program synthesis

TL;DR: This paper introduces a masking approach to extract reusable sub-functions from neural networks, treating neural network decomposition as a differentiable problem.

Abstract: Option discovery via neural network decomposition is a promising way of discovering temporally extended actions in reinforcement learning. The challenge is that the number of sub-functions a network encodes grows exponentially with its size, so finding sub-functions that can be useful in downstream tasks is a difficult combinatorial search problem. In this paper, we turn this combinatorial search problem into a differentiable problem by showing that extracting sub-functions from a network is equivalent to learning masks over the neurons of the network. In addition to extracting sub-functions, we can also learn default input parameters to such sub-functions through masks over the inputs. Neuron masks select what to execute; input masks specify how to call it. We evaluate our masking scheme on grid-world problems with binary and pixel observations, using both feedforward and recurrent policies. Our results show that masking can produce sub-functions with default input parameters that improve sample efficiency on downstream tasks.

Primary Area: reinforcement learning

Submission Number: 20581