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

📄 arXiv: 2512.07437v1 📥 PDF

作者: Chenwei Shi, Xueyu Luan

分类: cs.LG, cs.AI, cs.CV, cs.NE, cs.RO

发布日期: 2025-12-08

备注: 23 pages, 8 figures, 3 tables

💡 一句话要点

KAN-Dreamer：探索Kolmogorov-Arnold网络在World Models中作为函数逼近器的性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Kolmogorov-Arnold网络 World Models 模型强化学习 DreamerV3 函数逼近 深度学习 机器人控制

📋 核心要点

  1. DreamerV3虽然性能优异，但其内部模块依赖MLP，存在参数效率和可解释性方面的局限。
  2. KAN-Dreamer通过将DreamerV3中的MLP和卷积层替换为KAN和FastKAN，探索了KAN在World Models中的应用潜力。
  3. 实验表明，使用FastKAN替代奖励和持续预测器后，KAN-Dreamer在样本效率和训练速度上与原始DreamerV3相当。

📝 摘要（中文）

DreamerV3是一种先进的在线模型强化学习(MBRL)算法，以卓越的样本效率著称。与此同时，Kolmogorov-Arnold网络(KANs)作为多层感知机(MLPs)的一种有前景的替代方案崭露头角，提供了卓越的参数效率和可解释性。为了缓解KANs的计算开销，FastKAN等变体利用径向基函数(RBFs)来加速推理。本文研究了将KAN架构集成到DreamerV3框架中。我们引入了KAN-Dreamer，用KAN和FastKAN层替换了DreamerV3的特定MLP和卷积组件。为了确保在基于JAX的World Model中的效率，我们实现了一个定制的、完全矢量化的版本，并简化了网格管理。我们将研究分为三个子系统：视觉感知、潜在预测和行为学习。在DeepMind Control Suite (walker_walk)上的实证评估分析了样本效率、训练时间和渐近性能。实验结果表明，使用我们调整后的FastKAN作为奖励和持续预测器的直接替代品，可以获得与原始基于MLP的架构相当的性能，并在样本效率和训练速度上保持一致。本报告是未来基于KAN的世界模型开发的初步研究。

🔬 方法详解

问题定义：论文旨在探索使用Kolmogorov-Arnold Networks (KANs) 替代 DreamerV3 中传统的多层感知机 (MLP) 的可行性与性能。DreamerV3 虽然在模型强化学习 (MBRL) 领域表现出色，但 MLP 的参数效率和可解释性存在不足，限制了其在复杂环境中的应用潜力。

核心思路：论文的核心思路是将 KANs 作为函数逼近器，嵌入到 DreamerV3 的 World Model 中，以期在保持或提升性能的同时，提高参数效率和模型的可解释性。通过替换 DreamerV3 中的关键 MLP 组件，例如奖励预测器和持续预测器，来评估 KANs 的有效性。

技术框架：KAN-Dreamer 的整体框架基于 DreamerV3，主要修改集中在三个子系统：视觉感知、潜在预测和行为学习。具体而言，论文将 DreamerV3 中的部分 MLP 层替换为 KAN 或 FastKAN 层。为了保证效率，论文实现了一个完全矢量化的 KAN 版本，并简化了网格管理，使其能够与 JAX 框架兼容。

关键创新：该论文的主要创新在于首次将 KANs 引入到 DreamerV3 框架中，并针对 World Model 的特点进行了优化。通过实验验证了 KANs 在保持性能的同时，具备替代 MLP 的潜力。此外，论文还提出了一个高效的 KAN 实现，使其能够更好地适应 World Model 的计算需求。

关键设计：论文的关键设计包括：1) 使用 FastKAN 变体来加速推理，降低计算开销；2) 针对 JAX 框架进行矢量化实现，提高计算效率；3) 简化网格管理，降低内存占用；4) 在 DeepMind Control Suite (walker_walk) 环境中，对比 KAN-Dreamer 与原始 DreamerV3 的性能，评估 KANs 的有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在 DeepMind Control Suite (walker_walk) 环境中，使用 FastKAN 替代 DreamerV3 的奖励和持续预测器后，KAN-Dreamer 在样本效率和训练速度上与原始 DreamerV3 相当。这表明 KANs 具备替代 MLPs 的潜力，并且能够在保持性能的同时，提高参数效率和模型的可解释性。

🎯 应用场景

KAN-Dreamer 的研究成果可应用于机器人控制、游戏AI、自动驾驶等领域。通过使用 KANs 替代 MLPs，有望构建更高效、更易于理解和调试的 World Model，从而提升智能体的学习效率和泛化能力。未来，该研究可以扩展到更复杂的环境和任务中，推动模型强化学习技术的发展。

📄 摘要（原文）

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.