def q_learning(seed, Q, visits, env, num_episodes,
alpha=lambda n: 1 / (1 + 0.1 * n)):
"""Tabular Q-learning; updates Q in place, yields each episode's return."""
rng = np.random.default_rng(seed)
epsilon = linear_schedule(1.0, 0.05, num_episodes // 2)
env.reset(seed=seed)
for episode in range(num_episodes):
s, done, ret = env.reset()[0], False, 0.0
while not done:
a = epsilon_greedy(Q[s], epsilon(episode), rng)
s_next, r, terminated, truncated, _ = env.step(a)
visits[s, a] += 1
delta = r + gamma * (1 - terminated) * Q[s_next].max() - Q[s, a]
Q[s, a] += alpha(visits[s, a]) * delta
s, done, ret = s_next, terminated or truncated, ret + r
yield ret