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

2025-12-08

AI Summary

This work addresses the parameter redundancy and poor interpretability of MLPs in the DreamerV3 world model. We propose the first integration of Kolmogorov–Arnold Networks (KANs) and their efficient variant, FastKAN, into an online model-based reinforcement learning framework. Leveraging a JAX-native fully vectorized implementation and a lightweight grid management strategy, we systematically evaluate KAN-based approximators across three core subsystems: visual encoding, latent dynamics modeling, and reward/continuation prediction. Experimental results on the *walker_walk* task show that replacing only the reward and continuation prediction modules with FastKAN achieves performance, sample efficiency, and training speed comparable to the original MLP baseline—while substantially improving parameter efficiency and function-level interpretability. This work establishes a new paradigm for designing interpretable and parameter-efficient world models.

Technology Category

Machine Learning: Deep Neural Architectures and Foundation ModelsComputer Vision: Generative Adversarial Networks (GANs) for VisionHumans and AI: Human-Aware Planning and Behavior Prediction

Application Category

Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphsSearch and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for rankingSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Integrating KAN architectures into DreamerV3 for model-based reinforcement learning.
• Replacing MLP and convolutional components with KAN and FastKAN layers.
• Evaluating performance in visual perception, latent prediction, and behavior learning.

Innovation

• Replaced MLP and convolutional layers with KAN and FastKAN
• Implemented a fully vectorized version for efficiency in JAX
• Used FastKAN as a drop-in replacement for Reward and Continue predictors