Interpretable Reinforcement Learning for Load Balancing using Kolmogorov-Arnold Networks

TL;DR

This work addresses the interpretability gap in reinforcement learning for network load balancing by introducing Kolmogorov-Arnold Networks (KAN) as the actor in a PPO framework, paired with an MLP critic. It demonstrates that a 1-layer KAN can achieve competitive performance while enabling extraction of symbolic, human-interpretable policies. The authors present two policy-extraction pathways: direct symbolic extraction from the KAN activations (PPO-KAN) and distillation of symbolic equations from RL trajectories (PPO-DS). Through packet-level NS-3 simulations, the approach yields improvements in throughput utility and loss reduction with interpretable policies, highlighting its potential for deployment in real networks and for extending to other TE problems.

Abstract

Reinforcement learning (RL) has been increasingly applied to network control problems, such as load balancing. However, existing RL approaches often suffer from lack of interpretability and difficulty in extracting controller equations. In this paper, we propose the use of Kolmogorov-Arnold Networks (KAN) for interpretable RL in network control. We employ a PPO agent with a 1-layer actor KAN model and an MLP Critic network to learn load balancing policies that maximise throughput utility, minimize loss as well as delay. Our approach allows us to extract controller equations from the learned neural networks, providing insights into the decision-making process. We evaluate our approach using different reward functions demonstrating its effectiveness in improving network performance while providing interpretable policies.

Figures (6)

• Figure 1: Load-balancing topology
• Figure 2: KAN Actor and MLP Critic Architecture
• Figure 3: KAN-PPO trained activations using loss reward
• Figure 4: Two methods for extracting symbolic policies
• Figure 5: CCDF of network metrics for PPO and KAN-PPO