---
title: 'KAN-Dreamer: KAN-Based MBRL Architecture'
url: https://www.emergentmind.com/topics/kan-dreamer
type: topic
---

# KAN-Dreamer: KAN-Based MBRL Architecture

KAN-Dreamer denotes a series of model-based reinforcement learning (MBRL) architectures that benchmark Kolmogorov–Arnold Networks (KANs) and their efficient variant FastKAN as functional components in DreamerV3-style world models. By substituting conventional multilayer perceptrons (MLPs) and convolutional sub-modules within DreamerV3 with KAN-based alternatives, KAN-Dreamer systematically evaluates KANs as differentiable function approximators in three principal subsystems: visual perception, latent prediction (reward and continuation heads), and behavior learning (actor-critic). Empirical assessments on DeepMind Control Suite’s walker_walk task illuminate the regimes in which KANs and FastKAN match or underperform relative to established MLP/CNN baselines, characterizing trade-offs in parameter efficiency, interpretability, and computational cost [2512.07437].

## 1. Theoretical Foundations: Kolmogorov–Arnold Networks

KANs are motivated by the Kolmogorov–Arnold representation theorem, which asserts that any continuous multivariate function $f : [0,1]^d \rightarrow \mathbb{R}$ can be constructed via a finite superposition of continuous univariate functions and addition. Operationally, a KAN layer computes
$$
f(x_1,\dots,x_d) = \sum_{q=1}^{m} \varphi_q\left(\sum_{p=1}^{d} \gamma_{pq}(x_p)\right),
$$
where $\gamma_{pq} : \mathbb{R} \to \mathbb{R}$ (inner functions) and $\varphi_q : \mathbb{R} \to \mathbb{R}$ (outer functions) are trainable univariate mappings merged with a linear branch in typical application. KANs generalize universal approximation, offering structural efficiency and potential analytic interpretability through their explicit decomposition [2512.07437].

FastKAN addresses the prohibitive computational overhead of spline-based KANs by employing a bank of $G$ Gaussian radial basis functions (RBFs) per edge:
$$
\varphi(x) = w_b b(x) + w_s \sum_{k=1}^{G} c_k\,\exp\left(-\left(\frac{x - \mu_k}{h}\right)^2\right)
$$
with fixed base activation $b(x)$ (e.g., SiLU), $w_b=w_s=1$, fixed RBF centers $\mu_k$ uniformly spaced on $[-2,2]$, bandwidth $h$, and trainable coefficients $c_k$. This construct enables fully vectorized and simplified implementations in JAX, exclusively learning $c_k$ per edge and eliminating adaptive grid mechanisms [2512.07437].

## 2. KAN-Dreamer Architectural Integration

KAN-Dreamer is instantiated as a drop-in replacement strategy for MLP/CNNs in the following DreamerV3 functional units, preserving the GRU-based RSSM:

- **Visual Perception (Encoder/Decoder):**
  - *Baseline:* CNN encoder / transposed CNN decoder.
  - *KAN-Vis:* Image flattening and fully connected KAN layers.
  - *FKAN-Vis:* Fully connected FastKAN layers.
- **Latent Prediction (Reward & Continue Heads):**
  - *Baseline:* 2-layer MLP on $(h_t, z_t)$.
  - *KAN-Pred:* 2-layer KAN ($G=$8, spline order $k=$3).
  - *FKAN-Pred:* 2-layer FastKAN ($G=$8 RBFs).
- **Behavior Learning (Actor / Critic):**
  - *Baseline:* 2-layer MLP on $s_t=(h_t, z_t)$.
  - *KAN-AC:* 2-layer KAN ($G=$8, $k=$3).
  - *FKAN-AC:* 2-layer FastKAN.

In all regimes, KAN/ FastKAN layers are parameter-matched to baselines (total ≃10.5M), utilizing fixed input clamp ranges ([−5,5] for KAN; [−2,2] for FastKAN) and consistent training hyperparameters (e.g., replay capacity $5 \times 10^6$, batch size 16, sequence length 64, LaProp optimizer, RMSNorm + SiLU activations). Grid sizes and hidden unit counts are subsystem-specific: KAN-Vis uses 8 units per FC, KAN-Pred 20, and KAN-AC 24; FKAN counterparts use 12, 30, and 34 respectively. All FastKAN operations utilize a single `tensordot` contraction and fixed grid management [2512.07437].

## 3. Empirical Evaluation and Ablation Analysis

KAN-Dreamer’s evaluation on walker_walk consists of subsystem ablations, comparing final returns, sample efficiency, throughput (FPS), parameter count, and loss trajectories.

| Subsystem   | Model      | Final Score | FPS (Pol.) | FPS (Trn.) | Params   |
|-------------|------------|-------------|------------|------------|----------|
| Perception  | Baseline   | 977         | 62         | 15.9k      | 10.49M   |
|             | KAN-Vis    | 940         | 72         | 18.5k      | 10.40M   |
|             | FKAN-Vis   | 950         | 89         | 22.8k      | 10.51M   |
| Prediction  | Baseline   | 977         | 62         | 15.9k      | 10.49M   |
|             | KAN-Pred   | 965         | 45         | 11.5k      | 10.45M   |
|             | FKAN-Pred  | 951         | 61         | 15.5k      | 10.51M   |
| Behavior    | Baseline   | 977         | 62         | 15.9k      | 10.49M   |
|             | KAN-AC     | 957         | 31         | 7.8k       | 10.48M   |
|             | FKAN-AC    | 964         | 57         | 14.5k      | 10.47M   |

Key observations include:

- **Sample Efficiency:** Baseline achieves reward 900 at ~250k steps. FKAN-Pred requires ~220k, KAN-Pred ~250k, FKAN-AC ~520k, FKAN-Vis ~750k, KAN-AC and KAN-Vis slower at ~900k and ~930k steps respectively.
- **Wall-Clock Performance:** FKAN-Pred matches baseline at ~1.0 h; KAN-Pred at ~1.5 h; FKAN-Vis at ~2.5 h; KAN-Vis at ~3.5 h; KAN-AC exceeds 8 h.
- **Loss Dynamics:** KAN/ FastKAN-Vis plateau at reconstruction loss ~100 (indicating poor visual fidelity vs. baseline), whereas reward/continue predictors and policy/value heads converge comparably in loss but with higher variance and slower value approximation for actor-critic.

## 4. Functional Insights and Interpretative Context

The empirical parity of FastKAN and MLP architectures in the low-dimensional latent prediction regime (Reward and Continue heads) results from the alignment of task properties (deterministic, low-dimensional target functions) with the structural decomposition advocated by the Kolmogorov–Arnold theorem. FastKAN’s entirely vectorized, grid-based inference substantially reduces spline overhead, yielding throughput comparable to MLPs.

Conversely, KANs and FastKAN underperform in vision and policy/value estimation due to architectural incompatibilities: 

- **Visual Perception Limitation:** Fully connected KANs lack convolutional spatial biases, impeding learning of pixel locality and resulting in poor reconstructions and diminished sample efficiency.
- **Policy/Value Approximation Instability:** The recursive, non-stationary critic landscape imposes learning dynamics for which current KAN/ FastKAN parameterizations are insufficient, manifesting as high-variance and slow convergence.
- **Overhead:** KAN layers exhibit ~2× computational overhead relative to MLPs. FastKAN addresses this with minimal loss of task-relevant expressivity in latent heads, but not in high-dimensional or structurally non-aligned tasks [2512.07437].

## 5. Implementation and Vectorization Strategies

The FastKAN implementation is designed for JAX, employing a fully vectorized approach for scalability. Grid centers $\mu_k$ are initialized uniformly over $[-2,2]$ for all edges, with no adaptive updates, and all operations are compiled to a single `tensordot` contraction for batch efficiency. Only the RBF weights $c_k$ are learnable per (output, input, gridpoint). This implementation ensures that FastKAN, in the tested low-dimensional subsystems, can saturate MLP-equivalent training throughput.

A representative pseudocode fragment:

```python
import jax.numpy as jnp

def FastKANLayer(x, W, b, c, mu, h):
    # x: [batch, d_in]
    # W, b: linear branch params
    # c: [d_out, d_in, G], mu: [G], h: bandwidth
    y_base = x @ W + b  # [batch, d_out]
    x_exp = x[..., None] - mu[None, None, :]  # [batch, d_in, G]
    phi = jnp.exp(-((x_exp / h) ** 2))        # [batch, d_in, G]
    y_rbf = jnp.tensordot(phi, c, axes=([1, 2], [1, 2]))  # [batch, d_out]
    return y_base + y_rbf
```

## 6. Limitations and Prospective Advances

KAN-Dreamer’s primary limitations arise from representational mismatch and computational constraints:

- **Spatiotemporal Bias Absence:** Off-the-shelf KANs do not retain convolutional priors necessary for high-fidelity image modeling, limiting their utility in visual systems.
- **Critic/Policy Head Instability:** Value head instability under KAN/FastKAN parameterizations suggests a need for enhanced compositional or recurrent architectures.
- **Scalability:** Despite FastKAN improvements, computational efficiency and expressivity remain suboptimal outside low-dimensional, structurally aligned tasks.

Future research avenues include:

1. **Tailored Hyperparameters and Architectures:** Subsystem-specific co-optimization for grid size, RBF bandwidth, and input normalization.
2. **Symbolic Interpretability:** Applying symbolic regression to analytic forms of $\varphi$ and $\gamma$—facilitated by regression reward heads—to uncover explicit function mappings.
3. **Recurrent and KAN-RSSM Models:** Integration of stateful KAN-RNNs or KAN-parametrized transition models for enhanced long-horizon modeling.
4. **Spatial Priors in KANs:** Developing convolutional or locality-aware KANs for improved visual perception [2512.07437].

A plausible implication is that KAN and FastKAN, when appropriately integrated into world model subsystems matching their strengths, can serve as interpretable, competitive alternatives to MLPs in sample-efficient MBRL, but that significant architectural innovation is required to extend these advantages to high-dimensional and recursive domains.

Source: https://www.emergentmind.com/topics/kan-dreamer