Q-Learning on Gymnasium

import numpy as np

# Number of bins per dimension. More bins = more states = slower learning.
# Fewer bins = coarser representation = faster learning but less precision.
N_BINS = 6

# Clip limits. Values beyond these get clamped to the edge bin.
# We use tighter limits than the space's true bounds because the agent
# rarely reaches the extreme ends before failing.
CLIP_LOW  = np.array([-2.4, -3.0, -0.25, -3.0])
CLIP_HIGH = np.array([ 2.4,  3.0,  0.25,  3.0])

def discretize(obs):
    """
    Convert a CartPole continuous observation into a tuple of bin indices.

    This is the same idea as the RL course's discretize() for ultrasonic
    readings, but applied to 4 dimensions at once using NumPy.
    """
    # Clip values to our expected range
    obs_clipped = np.clip(obs, CLIP_LOW, CLIP_HIGH)

    # Scale to [0, N_BINS - 1]
    # (obs - low) / (high - low) gives a value in [0, 1]
    # multiply by N_BINS and floor to get a bin index in [0, N_BINS - 1]
    scaled = (obs_clipped - CLIP_LOW) / (CLIP_HIGH - CLIP_LOW)
    bins = np.floor(scaled * N_BINS).astype(int)

    # Clamp to [0, N_BINS - 1] to handle any edge cases at the exact maximum
    bins = np.clip(bins, 0, N_BINS - 1)

    return tuple(bins)   # e.g. (3, 2, 4, 1)

import gymnasium as gym

env = gym.make("CartPole-v1")
obs, info = env.reset(seed=0)
print("Raw obs:", obs)
print("Discretized:", discretize(obs))
env.close()
# Raw obs: [ 0.04  -0.04   0.01  -0.02]
# Discretized: (3, 2, 3, 2)

import numpy as np
import gymnasium as gym
import random

# ── Discretisation settings ──────────────────────────────────────────────────
N_BINS  = 6
CLIP_LOW  = np.array([-2.4, -3.0, -0.25, -3.0])
CLIP_HIGH = np.array([ 2.4,  3.0,  0.25,  3.0])

def discretize(obs):
    obs_clipped = np.clip(obs, CLIP_LOW, CLIP_HIGH)
    scaled = (obs_clipped - CLIP_LOW) / (CLIP_HIGH - CLIP_LOW)
    bins = np.floor(scaled * N_BINS).astype(int)
    bins = np.clip(bins, 0, N_BINS - 1)
    return tuple(bins)

# ── Q-learning hyperparameters ───────────────────────────────────────────────
ALPHA         = 0.1     # learning rate
GAMMA         = 0.99    # discount factor
EPSILON_START = 1.0     # start fully random
EPSILON_MIN   = 0.01    # never go below 1 % random
DECAY_RATE    = 0.995   # per-episode decay (same as RL course)
NUM_EPISODES  = 500     # enough to see clear improvement

# ── Q-table ──────────────────────────────────────────────────────────────────
q_table = {}

def get_q(state):
    """Return Q-values for state, initialising to zero if unseen."""
    if state not in q_table:
        q_table[state] = [0.0, 0.0]   # one value per action (0=left, 1=right)
    return q_table[state]

def choose_action(state, epsilon):
    if random.random() < epsilon:
        return random.randint(0, 1)    # explore
    return int(np.argmax(get_q(state))) # exploit

def update_q(state, action, reward, next_state, terminated):
    current_q = get_q(state)[action]
    if terminated:
        # No future reward from a terminal state
        td_target = reward
    else:
        td_target = reward + GAMMA * max(get_q(next_state))
    get_q(state)[action] += ALPHA * (td_target - current_q)

# ── Training loop ─────────────────────────────────────────────────────────────
env = gym.make("CartPole-v1")
epsilon = EPSILON_START
recent_rewards = []

for episode in range(1, NUM_EPISODES + 1):
    obs, info = env.reset(seed=episode)
    state = discretize(obs)
    episode_reward = 0
    done = False

    while not done:
        action = choose_action(state, epsilon)
        next_obs, reward, terminated, truncated, info = env.step(action)
        next_state = discretize(next_obs)

        # Only pass terminated to update_q, not truncated.
        # A time-limit ending is not a true terminal state.
        update_q(state, action, reward, next_state, terminated)

        state = next_state
        episode_reward += reward
        done = terminated or truncated

    epsilon = max(EPSILON_MIN, epsilon * DECAY_RATE)
    recent_rewards.append(episode_reward)

    if episode % 50 == 0:
        avg = sum(recent_rewards[-50:]) / 50
        print(f"Episode {episode:4d}  avg(50)={avg:6.1f}  epsilon={epsilon:.3f}  states={len(q_table)}")

env.close()
print(f"\nFinal Q-table: {len(q_table)} states visited")