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

TL;DR

This work benchmarks Kolmogorov-Arnold Networks (KAN) and the efficient FastKAN variant as replacements for DreamerV3 components within a world-model RL framework. It demonstrates that FastKAN can serve as a competitive drop-in for latent predictors (reward and continuation), achieving parity with the MLP baseline in asymptotic performance and training speed on walker_walk, while standard KANs struggle with high-dimensional visual encoding and stable policy optimization. The results highlight the importance of inductive biases and temporal structure for KANs in vision- and control-heavy modules, suggesting targeted future work on spatial priors, recurrent KANs, and component-specific hyperparameters. Overall, FastKAN offers a promising, parameter-efficient alternative for certain world-model heads, forming a basis for further architectural developments in KAN-based latent dynamics models.

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.

Figures (8)

• Figure 1: Architecture of the KAN-Dreamer World Model. (a) Variables: White boxes denote data tensors: $z_t$ (stochastic latent state), $h_t$ (deterministic recurrent state), $r_t$ (reward), $c_t$ (continuation flag), and $a_t$ (action). (b) Components: Colored boxes represent learnable modules. Dashed borders indicate configurable modules (Encoder, Decoder, Predictors) where the backbone can be instantiated as standard architectures (MLP/CNN), KAN, or FastKAN. The solid-bordered Sequence Model (cyan) retains the fixed GRU architecture for stability. Blue arrows indicate the data flow during the forward pass. Note: Proprioceptive inputs $q$ are depicted for completeness but were masked out in our visual-only experiments.
• Figure 2: Actor-Critic architecture
• Figure 3: Training Dynamics across Component Groups. Darker lines indicate smoothed curves (factor 0.8), while lighter traces show raw data.
• Figure 4: Wall-Clock Training Efficiency. Darker lines indicate smoothed curves (factor 0.8), while lighter traces show raw data.
• Figure 5: Image Reconstruction Loss. Comparison of Baseline (orange), KAN-Vis (blue), and FKAN-Vis (red).