KAN-Dreamer: Benchmarking Kolmogorov-Arnold Networks as Function Approximators in World Models

TL;DR

This paper explores integrating Kolmogorov-Arnold Networks into DreamerV3 for improved parameter efficiency, demonstrating comparable performance and efficiency in reinforcement learning tasks.

Contribution

It introduces KAN-Dreamer, replacing MLP components with KAN and FastKAN layers, and implements a vectorized version for efficient integration within the DreamerV3 framework.

Findings

KAN-based components achieve similar performance to MLPs

FastKAN accelerates inference with maintained accuracy

Sample efficiency and training speed are preserved

Abstract

DreamerV3 is a state-of-the-art online model-based reinforcement learning (MBRL) algorithm known for remarkable sample efficiency. Concurrently, Kolmogorov-Arnold Networks (KANs) have emerged as a promising alternative to Multi-Layer Perceptrons (MLPs), offering superior parameter efficiency and interpretability. To mitigate KANs' computational overhead, variants like FastKAN leverage Radial Basis Functions (RBFs) to accelerate inference. In this work, we investigate integrating KAN architectures into the DreamerV3 framework. We introduce KAN-Dreamer, replacing specific MLP and convolutional components of DreamerV3 with KAN and FastKAN layers. To ensure efficiency within the JAX-based World Model, we implement a tailored, fully vectorized version with simplified grid management. We structure our investigation into three subsystems: Visual Perception, Latent Prediction, and Behavior Learning. Empirical evaluations on the DeepMind Control Suite (walker_walk) analyze sample efficiency, training time, and asymptotic performance. Experimental results demonstrate that utilizing our adapted FastKAN as a drop-in replacement for the Reward and Continue predictors yields performance on par with the original MLP-based architecture, maintaining parity in both sample efficiency and training speed. This report serves as a preliminary study for future developments in KAN-based world models.