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

AI summary

Overview

Research area: Model-based reinforcement learning (MBRL), specifically the architecture of learned world models, combining the DreamerV3 algorithm with Kolmogorov-Arnold Networks (KANs) and their efficient variant FastKAN.

Technical level: Intermediate. Readers need familiarity with reinforcement learning concepts (actor, critic, latent states, reward prediction) and basic neural architecture vocabulary, but the paper is written as an exploratory technical report and does not require deep mathematical background.

One-sentence scope: The paper benchmarks whether KAN and FastKAN layers can replace the MLP and CNN components inside DreamerV3's world model and actor-critic, tested on the DeepMind Control Suite walker_walk task under an iso-parameter budget of approximately 10.5M parameters.

What This Paper Is About

DreamerV3, a leading online model-based RL algorithm, builds its world model and policy networks from conventional Multi-Layer Perceptrons (MLPs) and Convolutional Neural Networks (CNNs). KANs are a newer architecture that places learnable univariate functions on network edges instead of fixed activations on nodes, offering theoretical advantages in parameter efficiency and interpretability, but they have never been tested as function approximators inside a world model. This paper asks whether swapping KAN or FastKAN layers into specific DreamerV3 subsystems is feasible, and what it costs in performance, sample efficiency, and wall-clock training speed.

Key Contributions

1. First KAN-in-world-model integration. The authors implement KAN and FastKAN layers within the DreamerV3 framework, reporting what they describe as the first empirical evidence of their viability as replacements for standard MLPs in complex latent dynamics models.

2. Trade-off analysis between fitting capability and inference latency. They investigate how the learnable univariate functions of KANs affect the speed of imagination-trajectory generation and the resulting sample efficiency of the agent.

3. A systematic three-part ablation. Rather than swapping the whole network, they isolate the effect of KAN/FastKAN on three functional subsystems: Visual Perception (Encoder/Decoder), Latent Prediction (Reward/Continue predictors), and Behavior Learning (Actor/Critic), each compared against a fixed baseline.

4. A tailored JAX implementation. Because DreamerV3 is JAX-based while KAN libraries are predominantly PyTorch-native, they build a fully vectorized JAX version with simplified fixed grid management, single-einsum B-spline aggregation for KAN, and tensor-contraction RBF evaluation for FastKAN, specifically to remain compatible with the world model.

Main Findings

• Baseline reference: Under the iso-parameter setting of approximately 10.5M parameters, standard DreamerV3 achieves a final score of 977, with 62 FPS of environment interaction and 15.9k FPS of training throughput. It crosses the reward threshold of ≥900 at roughly 250k environment steps and about 1.1 hours of wall-clock runtime.

• Best result comes from the Prediction heads. Replacing the MLP-based Reward and Continue predictors with adapted FastKAN produces a final score of 951 with 61 FPS (Pol.) and 15.5k FPS (Trn.), which the authors describe as parity with the 977 / 62 FPS / 15.9k FPS baseline. Standard KAN in the same role scores 965 but drops to 45 FPS (Pol.) and 11.5k FPS (Trn.), the authors' stated trend being Baseline ≈ FastKAN > KAN.

• Visual perception swaps are faster but less accurate. Replacing the CNN encoder/decoder with flattened-input KAN or FastKAN yields final scores of 940 (KAN) and 950 (FastKAN), both below the 977 baseline, while throughput rises above the CNN baseline: FastKAN reaches 89 FPS (Pol.) / 22.8k FPS (Trn.) and KAN reaches 72 FPS (Pol.) / 18.5k FPS (Trn.), giving the ordering FastKAN > KAN > Baseline.

• Behavior learning swaps converge more slowly. KAN-AC scores 957 and FKAN-AC scores 964, but both lag the baseline in convergence speed. FastKAN reaches the ≥900 threshold at approximately 520k steps and standard KAN at nearly 900k steps, compared to roughly 250k steps for the baseline. Efficiency also suffers: KAN-AC operates at 31 FPS (Pol.) and 7.8k FPS (Trn.), while FKAN-AC reaches 57 FPS and 14.5k FPS, against the baseline's 62 FPS and 15.9k FPS.

• Sample efficiency by subsystem. In the Perception group, KAN requires approximately 930k steps and FastKAN approximately 750k steps to exceed a return of 900, both far slower than the baseline, and KAN-based encoders show larger reward fluctuations. In the Prediction group, all models converge around 250k steps, with FastKAN crossing the 900-reward threshold as early as 220k steps and variance comparable to the baseline.

• Wall-clock efficiency. The baseline reaches reward ≥900 in approximately 1.1 hours. FastKAN-Vis needs about 2.5 hours and KAN-Vis about 3.5 hours. In the Prediction group, standard KAN needs approximately 1.5 hours while FastKAN converges in just under 1 hour, slightly faster than the baseline. The wall-clock results for the Behavior group are not available in the provided content, as the text is truncated mid-sentence.

• Fixed-grid design choices. The authors deliberately use fixed uniform grids instead of the original KAN's adaptive grid updates, configuring KAN with grid size G=8 and spline order k=3 over a clamped input range of [-5.0, 5.0], and FastKAN with G=8 Gaussian centers over [-2.0, 2.0]. They note this may limit approximation precision but stabilizes gradient flow during online RL.

Methodology in Plain English

The authors start from the open-source DreamerV3 codebase, which learns a latent "world model" of an environment and then trains the agent's policy by imagining trajectories inside that learned model. Rather than rebuilding the whole system, they keep the architecture's data flow identical and make individual feed-forward modules swappable, leaving the recurrent GRU sequence model untouched for stability because KAN behavior in recurrent settings is under-explored.

They divide the model into three testable groups. In the Perception group, the CNN encoder and decoder are replaced by KAN layers that receive flattened image vectors and, on the decoding side, output a flat vector that is reshaped back to image dimensions, inspired by the Kolmogorov-Arnold Auto-Encoder. In the Prediction group, the low-dimensional Reward and Continue predictors are swapped for KAN or FastKAN as direct drop-in replacements. In the Behavior group, the Actor and Critic MLPs are swapped the same way.

Because the original DreamerV3 is written in JAX while KAN libraries are mostly PyTorch, the authors write their own vectorized JAX versions: KAN packs spline coefficients into a 3D tensor and aggregates with a single einsum, while FastKAN keeps RBF outputs as an [N, D, G] tensor and applies weights by tensor contraction. They also simplify initialization, using native Dreamer linear layers for the base branch, and cutting auxiliary features like masking and pruning to prioritize runtime performance.

Every variant is adjusted to roughly 10.5M parameters (±0.1M) so comparisons are iso-parameter, and all other hyperparameters follow the official DreamerV3 configuration. Evaluation uses the walker_walk task from the DeepMind Control Suite with pure image observations and no proprioceptive state vectors. Runs were performed on a single NVIDIA A800 Tensor Core GPU with an AMD EPYC 9654 CPU (7 cores per run), on Anolis OS 8.6 with Python 3.11 and JAX 0.4.33.

Why This Matters

Impact on research: The paper fills a specific gap: prior KAN work in reinforcement learning covers only model-free methods such as PPO, SAC, A2C, and DQN/Rainbow, and prior KAN generative work was largely limited to low-dimensional datasets like MNIST. This is the first benchmark of KANs inside a world model, covering visual reconstruction, latent-space reward prediction, and imagination-based policy learning. Its most useful takeaway is negative-leaning but practical: KAN variants are not automatic wins, but FastKAN is a viable drop-in for low-dimensional regression heads, while visual and policy components still pay a meaningful price in convergence speed.

Real-world applications:

• Robotics and embodied AI, where environment interaction is expensive and world models are used to reduce real-world sample needs.
• Autonomous driving and other control systems that rely on learned latent dynamics for planning.
• Scientific computing and simulation, where KANs have already shown promise as interpretable function approximators and parameter-efficient surrogates.
• Resource-constrained training pipelines, where the higher throughput observed for FastKAN visual layers (89 FPS Pol. / 22.8k FPS Trn.) could matter even at some accuracy cost.

Industry relevance: Teams already running DreamerV3 at scale now have concrete data on which components can be swapped without hurting results and which should be avoided. The throughput numbers are particularly relevant, since FastKAN's Prediction-head variant matched the baseline almost exactly (61 vs 62 FPS Pol.), while KAN-AC more than halved interaction speed (31 FPS), a difference that translates directly into compute cost.

Future Directions

• Extending beyond a single task. The entire study is confined to walker_walk from the DeepMind Control Suite; whether these conclusions hold across the roughly 150 heterogeneous tasks DreamerV3 was designed for remains untested.

• Making the visual and behavior swaps competitive. The KAN/FastKAN encoder, decoder, actor, and critic all converge more slowly than the baseline, so closing that gap, perhaps through adaptive grids, hybrid KAN-MLP designs, or alternative bases such as sinusoids, is the obvious next step.

• Reconsidering the recurrent component. The authors deliberately kept the GRU sequence model fixed because no established KAN architecture for recurrent settings exists and efficiency there is under-explored, leaving the recurrence itself as an open research question.

• Scaling evaluation. Since the paper presents itself as a preliminary study, larger parameter budgets, longer training horizons, and more seeds would be needed to confirm whether the observed variance in KAN-based visual encoders is a fundamental property or a small-scale artifact.

Target Audience

This paper suits reinforcement learning researchers and engineers who work with DreamerV3 or other latent world models and want to know whether newer architecture families are worth adopting. It is also useful for practitioners evaluating KANs outside the usual toy-regression and AI-for-science settings, and for anyone looking for a concrete, iso-parameter comparison of throughput versus sample efficiency across neural backbones. Because it is framed as an exploratory technical report with a single benchmark task, readers seeking definitive architectural recommendations should treat it as a starting point rather than a final verdict.

Authors’ 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.