Reduce, Reuse, Recycle: Categories for Compositional Reinforcement Learning
Georgios Bakirtzis, Michail Savvas, Ruihan Zhao, Sandeep Chinchali, Ufuk Topcu
TL;DR
This work introduces a category-theoretic framework for compositional reinforcement learning, modeling MDPs as objects and morphisms with subprocesses and pushouts to enable modular task composition. The zig-zag diagram formalism provides a denotational account of sequential task completion, with conditions for when subtask policies remain optimal for the overall task. Empirical results on robosuite manipulation tasks demonstrate faster convergence and higher final performance compared to baselines, and show the potential for reusing and recycling learned subtask policies. While promising, the authors acknowledge the need for standardized compositional generalization benchmarks and pursue future work on scalability and interpretability of the categorical RL approach.
Abstract
In reinforcement learning, conducting task composition by forming cohesive, executable sequences from multiple tasks remains challenging. However, the ability to (de)compose tasks is a linchpin in developing robotic systems capable of learning complex behaviors. Yet, compositional reinforcement learning is beset with difficulties, including the high dimensionality of the problem space, scarcity of rewards, and absence of system robustness after task composition. To surmount these challenges, we view task composition through the prism of category theory -- a mathematical discipline exploring structures and their compositional relationships. The categorical properties of Markov decision processes untangle complex tasks into manageable sub-tasks, allowing for strategical reduction of dimensionality, facilitating more tractable reward structures, and bolstering system robustness. Experimental results support the categorical theory of reinforcement learning by enabling skill reduction, reuse, and recycling when learning complex robotic arm tasks.