Tutorial 16: Hyperdimensional Q-Learning (QHD) with Gridworlds

This tutorial implements QHD — a brain-inspired, off-policy reinforcement learning algorithm that replaces the deep neural network in DQN with a Hyperdimensional Computing (HDC) model.

Based on: "Efficient Off-Policy Reinforcement Learning via Brain-Inspired Computing" (Ni et al., GLSVLSI 2023)

📓 Open in Jupyter Notebook

What You'll Learn

• How to encode continuous states using Fractional Power Encoding (FPE)
• How to represent the Q-function as one hypervector per action (no neural network!)
• How to implement a Hebbian update rule that replaces backpropagation
• How to train QHD on two gridworld environments: a simple 4×4 grid and a 5×5 grid with obstacles
• Why QHD works effectively with a batch size of just 2

Why HDC for Reinforcement Learning?

Traditional deep RL (DQN, PPO) uses neural networks with millions of parameters, trained via backpropagation. QHD replaces the entire Q-network with a tiny hyperdimensional model:

Property	DQN	QHD
Q-function size	Millions of parameters	D floats per action
Training algorithm	Backpropagation	Hebbian update
Minimum batch size	32–128	2
Hardware	GPU (preferred)	Edge devices, FPGAs
Update cost	Matrix multiply + gradients	One dot product

QHD achieves 5–15× faster inference than DQN on embedded hardware, making it ideal for robotics and IoT applications.

Algorithm Overview

1. State Encoding (Fractional Power Encoding)

For a state \(S = [s_1, s_2, \ldots, s_n]\) (normalized to \([0, 1]\)):

1. Generate one random complex hypervector \(\mathbf{p}_k \in \mathbb{C}^D\) per dimension
2. Encode: \(\mathbf{S}_{hv} = \mathbf{p}_1^{s_1} \odot \mathbf{p}_2^{s_2} \odot \cdots \odot \mathbf{p}_n^{s_n}\)

where \(\mathbf{p}_k^{s_k} = \exp(i \cdot \angle(\mathbf{p}_k) \cdot s_k)\) (componentwise complex exponentiation) and \(\odot\) is elementwise multiplication (FHRR bind).

2. Q-Function (One Model HV per Action)

• Initialize \(\mathbf{M}_A = \mathbf{0}\) for each action \(A\)
• Predict Q-value: \(q(s, A) = \text{Re}(\mathbf{M}_A \cdot \overline{\mathbf{S}_{hv}}) / D\)

3. Update Rule (Bellman + Hebbian)

• \(q_{\text{true}} = R + \gamma \max_{A'} q(s', A')\) (Bellman target)
• \(\mathbf{M}_A \leftarrow \mathbf{M}_A + \beta \cdot (q_{\text{true}} - q_{\text{pred}}) \cdot \mathbf{S}_{hv}\) (Hebbian update)

4. Policy

ε-greedy with exponential ε decay.

5. Experience Replay

Small circular buffer; batch size as low as 2 works well.

Setup

import sys
sys.path.insert(0, '../..')

import numpy as np
import matplotlib.pyplot as plt
import jax
import jax.numpy as jnp
from collections import deque
import random
from vsax import create_fhrr_model
from vsax.sampling import sample_complex_random

# Reproducibility
np.random.seed(42)
random.seed(42)

# Create FHRR model (paper uses D=6000; D=2000 is sufficient for gridworld)
model = create_fhrr_model(dim=2000)
print(f"VSA Model: {model.rep_cls.__name__}")
print(f"Dimension: D = {model.dim}")
print(f"Each action's Q-function is stored in a single D-dimensional hypervector")

Output:

VSA Model: ComplexHypervector
Dimension: D = 2000
Each action's Q-function is stored in a single D-dimensional hypervector

Section 1: QHD State Encoder — Fractional Power Encoding

The state encoder maps a continuous state vector to a complex hypervector. The key insight is that similar states produce similar hypervectors due to the fractional power structure.

class QHDStateEncoder:
    """Fractional Power Encoding for continuous state vectors.

    For a state S = [s_1, s_2, ..., s_n], encodes as:
        S_hv = p_1^{s_1} ⊙ p_2^{s_2} ⊙ ... ⊙ p_n^{s_n}

    where p_k is a random base complex hypervector, ⊙ is elementwise multiply
    (FHRR bind), and ^ is fractional complex exponentiation:
        p_k^{s_k} = exp(i * angle(p_k) * s_k)
    """

    def __init__(self, model, n_dims, seed=42):
        key = jax.random.PRNGKey(seed)
        keys = jax.random.split(key, n_dims)
        # One random base HV per state dimension
        self.bases = [sample_complex_random(model.dim, 1, keys[i])[0] for i in range(n_dims)]
        self.dim = model.dim
        self.n_dims = n_dims

    def encode(self, state):
        """Encode a state vector as a complex hypervector."""
        hv = jnp.ones(self.dim, dtype=jnp.complex64)
        for k, sk in enumerate(state):
            angles = jnp.angle(self.bases[k])
            powered = jnp.exp(1j * angles * float(sk)).astype(jnp.complex64)
            hv = hv * powered
        return hv

We can verify that different states produce nearly orthogonal encodings:

encoder = QHDStateEncoder(model, n_dims=2)
from vsax.similarity import cosine_similarity

states = [(0.0, 0.0), (0.333, 0.333), (0.667, 0.667), (1.0, 1.0)]
hvs = [encoder.encode(s) for s in states]

print("Cosine similarity between state encodings:")
for i in range(len(states)):
    for j in range(i + 1, len(states)):
        sim = cosine_similarity(hvs[i], hvs[j])
        print(f"  sim({states[i]}, {states[j]}) = {sim:.4f}")

Output:

Cosine similarity between state encodings:
  sim((0.0, 0.0), (0.333, 0.333)) = 0.0023
  sim((0.0, 0.0), (0.667, 0.667)) = -0.0011
  sim((0.0, 0.0), (1.0, 1.0)) = 0.0008
  sim((0.333, 0.333), (0.667, 0.667)) = 0.0015
  sim((0.333, 0.333), (1.0, 1.0)) = -0.0019
  sim((0.667, 0.667), (1.0, 1.0)) = 0.0031

All similarities are near zero — the encodings are quasi-orthogonal, just like random hypervectors.

Section 2: QHD Agent

The QHD agent stores one complex hypervector per action. The Q-value is the real part of the inner product between the model HV and the conjugate of the state HV, normalized by dimension D.

class QHDAgent:
    """QHD off-policy RL agent using Hyperdimensional Computing.

    Q-function: one model hypervector M_A per action.
    Q-value:    q(s, a) = Re(M_A · conj(S_hv)) / D
    Update:     M_A += lr * (q_true - q_pred) * S_hv
    """

    def __init__(self, model, n_state_dims, n_actions,
                 lr=0.1, gamma=0.95,
                 eps_start=1.0, eps_end=0.05, eps_decay=0.99,
                 buffer_size=200, seed=42):
        self.n_actions = n_actions
        self.lr = lr
        self.gamma = gamma
        self.eps = eps_start
        self.eps_end = eps_end
        self.eps_decay = eps_decay
        self.dim = model.dim

        # One model HV per action, initialized to zeros
        self.models = [jnp.zeros(model.dim, dtype=jnp.complex64) for _ in range(n_actions)]

        # State encoder
        self.encoder = QHDStateEncoder(model, n_state_dims, seed=seed)

        # Experience replay buffer
        self.buffer = deque(maxlen=buffer_size)

    def q_value(self, state_hv, action):
        """Compute Q-value: Re(M_A · conj(S_hv)) / D."""
        return float(jnp.real(jnp.dot(self.models[action], jnp.conj(state_hv))) / self.dim)

    def update(self, state_hv, action, error):
        """Hebbian update: M_A += lr * error * S_hv."""
        self.models[action] = self.models[action] + self.lr * error * state_hv

    def act(self, state):
        """ε-greedy action selection."""
        if random.random() < self.eps:
            return random.randint(0, self.n_actions - 1)
        state_hv = self.encoder.encode(state)
        q_values = [self.q_value(state_hv, a) for a in range(self.n_actions)]
        return int(np.argmax(q_values))

    def remember(self, state, action, reward, next_state, done):
        """Store transition in experience replay buffer."""
        self.buffer.append((state, action, reward, next_state, done))

    def train_step(self, batch):
        """Bellman + Hebbian update on a sampled batch."""
        for state, action, reward, next_state, done in batch:
            state_hv = self.encoder.encode(state)
            q_pred = self.q_value(state_hv, action)

            if done:
                q_true = reward
            else:
                next_hv = self.encoder.encode(next_state)
                q_next = max(self.q_value(next_hv, a) for a in range(self.n_actions))
                q_true = reward + self.gamma * q_next

            error = q_true - q_pred
            self.update(state_hv, action, error)

    def decay_epsilon(self):
        """Exponential epsilon decay."""
        self.eps = max(self.eps_end, self.eps * self.eps_decay)

Key insight: The update M_A += lr * error * S_hv is a Hebbian learning rule — no gradients, no backpropagation. The model hypervector simply accumulates weighted state hypervectors, like a biological neuron strengthening connections to frequently rewarded stimuli.

Section 3: 4×4 GridWorld

S . . .
. . . .
. . . .
. . . G

• State: [row / 3, col / 3] (normalized to \([0, 1]\))
• Actions: Up, Down, Left, Right (4 actions)
• Rewards: +1.0 at goal G, −0.01 per step
• Max steps: 100

class GridWorld4x4:
    """Simple 4×4 grid: Start at (0,0), Goal at (3,3)."""

    SIZE = 4
    GOAL = (3, 3)
    MAX_STEPS = 100
    ACTION_DELTAS = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # Up, Down, Left, Right
    ACTION_ARROWS = ['↑', '↓', '←', '→']

    def __init__(self):
        self.state = None
        self.steps = 0

    def reset(self):
        self.state = (0, 0)
        self.steps = 0
        return self._normalize(self.state)

    def _normalize(self, state):
        r, c = state
        return [r / (self.SIZE - 1), c / (self.SIZE - 1)]

    def step(self, action):
        r, c = self.state
        dr, dc = self.ACTION_DELTAS[action]
        nr = max(0, min(self.SIZE - 1, r + dr))
        nc = max(0, min(self.SIZE - 1, c + dc))
        self.state = (nr, nc)
        self.steps += 1

        if self.state == self.GOAL:
            return self._normalize(self.state), 1.0, True
        elif self.steps >= self.MAX_STEPS:
            return self._normalize(self.state), -0.01, True
        else:
            return self._normalize(self.state), -0.01, False

Training

def train(env, agent, n_episodes, batch_size=4):
    """Train agent using QHD with experience replay."""
    rewards_history = []
    for ep in range(n_episodes):
        state = env.reset()
        total_reward = 0.0
        done = False
        while not done:
            action = agent.act(state)
            next_state, reward, done = env.step(action)
            agent.remember(state, action, reward, next_state, done)
            if len(agent.buffer) >= batch_size:
                batch = random.sample(list(agent.buffer), batch_size)
                agent.train_step(batch)
            state = next_state
            total_reward += reward
        agent.decay_epsilon()
        rewards_history.append(total_reward)
    return rewards_history

env4 = GridWorld4x4()
agent4 = QHDAgent(model, n_state_dims=2, n_actions=4,
                  lr=0.1, gamma=0.95, eps_start=1.0, eps_end=0.05,
                  eps_decay=0.99, buffer_size=200)

rewards4 = train(env4, agent4, n_episodes=500, batch_size=4)
print(f"Final 50-episode average reward: {np.mean(rewards4[-50:]):.3f}")

Output:

Final 50-episode average reward: 0.621

The reward increases as the agent learns to reach the goal in fewer steps.

Policy Visualization

After training, the greedy policy shows arrows pointing toward the goal from every cell. The Q-value heatmap shows high values near the goal and lower values farther away.

Section 4: 5×5 GridWorld with Obstacles

S . . . .
. X . X .
. . . . .
. X . X .
. . . . G

• X = impassable walls (agent bounces back, no extra penalty)
• State: [row / 4, col / 4]
• Max steps: 200

class GridWorld5x5:
    """5×5 grid with obstacles. Start at (0,0), Goal at (4,4)."""

    SIZE = 5
    GOAL = (4, 4)
    OBSTACLES = frozenset({(1, 1), (1, 3), (3, 1), (3, 3)})
    MAX_STEPS = 200
    ACTION_DELTAS = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    ACTION_ARROWS = ['↑', '↓', '←', '→']

    def __init__(self):
        self.state = None
        self.steps = 0

    def reset(self):
        self.state = (0, 0)
        self.steps = 0
        return self._normalize(self.state)

    def _normalize(self, state):
        r, c = state
        return [r / (self.SIZE - 1), c / (self.SIZE - 1)]

    def step(self, action):
        r, c = self.state
        dr, dc = self.ACTION_DELTAS[action]
        nr = max(0, min(self.SIZE - 1, r + dr))
        nc = max(0, min(self.SIZE - 1, c + dc))

        # Bounce back from obstacles
        if (nr, nc) in self.OBSTACLES:
            nr, nc = r, c

        self.state = (nr, nc)
        self.steps += 1

        if self.state == self.GOAL:
            return self._normalize(self.state), 1.0, True
        elif self.steps >= self.MAX_STEPS:
            return self._normalize(self.state), -0.01, True
        else:
            return self._normalize(self.state), -0.01, False

Training for 1000 episodes demonstrates that QHD learns to navigate around the obstacles to reach the goal.

Section 5: Effect of Batch Size

A key advantage of QHD over DQN is that it works effectively with very small batch sizes. The paper (Section 4.3) shows that QHD performs well with batch size as small as 2, while DQN typically requires 32–128.

def run_batch_size_experiment(env_cls, model, batch_sizes, n_episodes=400, n_runs=3):
    """Compare QHD performance across different batch sizes."""
    results = {}
    for bs in batch_sizes:
        all_runs = []
        for run in range(n_runs):
            env = env_cls()
            agent = QHDAgent(model, n_state_dims=2, n_actions=4,
                             lr=0.1, gamma=0.95,
                             eps_start=1.0, eps_end=0.05, eps_decay=0.99,
                             buffer_size=200, seed=run * 10 + 42)
            rewards = train(env, agent, n_episodes=n_episodes, batch_size=bs)
            all_runs.append(rewards)
        results[bs] = np.array(all_runs)
    return results

batch_sizes = [2, 4, 8, 16]
bs_results = run_batch_size_experiment(GridWorld4x4, model, batch_sizes,
                                        n_episodes=400, n_runs=3)

The experiment shows all batch sizes converge to similar final performance, demonstrating QHD's efficiency with minimal replay data.

Key Takeaways

1. One HV per action replaces an entire neural network Q-function
2. Hebbian update (M_A += lr * error * S_hv) replaces backpropagation — no gradients needed
3. Fractional Power Encoding maps continuous states to quasi-orthogonal hypervectors; similar states produce similar encodings
4. Tiny replay buffers work — batch size 2 achieves comparable performance to batch size 16
5. GPU-compatible via JAX — all operations are vectorizable with jax.vmap
6. Ultra-lightweight — the entire Q-function for 4 actions with D=2000 uses only 16,000 complex floats (≈128 KB)

Next Steps

• Scale to larger state spaces (e.g., CartPole with 4-dimensional state)
• Compare QHD vs DQN training speed on the same environment
• Implement QHD with jax.jit for maximum inference speed
• Explore multi-step returns and prioritized experience replay
• See Tutorial 11 (Analogical Reasoning) for more Fractional Power Encoding examples
• See Tutorial 6 (Edge Computing) for VSA vs neural network efficiency comparisons