notebook

import sys
# Dependencies are provisioned by this course edition's pinned runtime profile.

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import random

np.random.seed(2026)
random.seed(2026)

env = gym.make("CartPole-v1", render_mode="rgb_array")
print(env.action_space)
print(env.observation_space)
print(env.action_space.sample())

env.reset()

for _ in range(100):
    env.render()
    _, _, terminated, truncated, _ = env.step(env.action_space.sample())
    if terminated or truncated:
        break

env.reset()

done = False
while not done:
   env.render()
   obs, rew, terminated, truncated, info = env.step(env.action_space.sample())
   done = terminated or truncated
   print(f"{obs} -> {rew}")

print(env.observation_space.low)
print(env.observation_space.high)

def discretize(x):
    return tuple((x/np.array([0.25, 0.25, 0.01, 0.1])).astype(np.int64))

def create_bins(i,num):
    return np.arange(num+1)*(i[1]-i[0])/num+i[0]

print("Sample bins for interval (-5,5) with 10 bins\n",create_bins((-5,5),10))

ints = [(-5,5),(-2,2),(-0.5,0.5),(-2,2)] # intervals of values for each parameter
nbins = [20,20,10,10] # number of bins for each parameter
bins = [create_bins(ints[i],nbins[i]) for i in range(4)]

def discretize_bins(x):
    return tuple(np.digitize(x[i],bins[i]) for i in range(4))

env.reset()

done = False
while not done:
   #env.render()
   obs, rew, terminated, truncated, info = env.step(env.action_space.sample())
   done = terminated or truncated
   #print(discretize_bins(obs))
   print(discretize(obs))

Q = {}
actions = (0,1)

def qvalues(state):
    return [Q.get((state,a),0) for a in actions]

# hyperparameters
alpha = 0.3
gamma = 0.9
epsilon = 0.90

def probs(v,eps=1e-4):
    v = v-v.min()+eps
    v = v/v.sum()
    return v

Qmax = 0
cum_rewards = []
rewards = []
for epoch in range(200):
    obs, _ = env.reset(seed=2026 + epoch)
    done = False
    cum_reward=0
    # == do the simulation ==
    while not done:
        s = discretize(obs)
        if random.random()<epsilon:
            # exploitation - chose the action according to Q-Table probabilities
            v = probs(np.array(qvalues(s)))
            a = random.choices(actions,weights=v)[0]
        else:
            # exploration - randomly chose the action
            a = np.random.randint(env.action_space.n)

        obs, rew, terminated, truncated, info = env.step(a)

        done = terminated or truncated
        cum_reward+=rew
        ns = discretize(obs)
        bootstrap = 0.0 if terminated else gamma * max(qvalues(ns))
        Q[(s,a)] = (1 - alpha) * Q.get((s,a), 0) + alpha * (rew + bootstrap)
    cum_rewards.append(cum_reward)
    rewards.append(cum_reward)
    # == Periodically print results and calculate average reward ==
    if epoch % 50==0:
        print(f"{epoch}: {np.average(cum_rewards)}, alpha={alpha}, epsilon={epsilon}")
        if np.average(cum_rewards) > Qmax:
            Qmax = np.average(cum_rewards)
            Qbest = Q
        cum_rewards=[]

plt.plot(rewards)

def running_average(x,window):
    return np.convolve(x,np.ones(window)/window,mode='valid')

plt.plot(running_average(rewards,100))

obs, _ = env.reset()
done = False
while not done:
   s = discretize(obs)
   env.render()
   v = probs(np.array(qvalues(s)))
   a = random.choices(actions,weights=v)[0]
   obs, _, terminated, truncated, _ = env.step(a)
   done = terminated or truncated

from pathlib import Path
from PIL import Image

output_dir = Path("outputs")
output_dir.mkdir(parents=True, exist_ok=True)

obs, _ = env.reset()
done = False
i = 0
ims = []
while not done:
   s = discretize(obs)
   img = env.render()
   if img is None:
      raise RuntimeError("CartPole rgb_array rendering returned no frame")
   ims.append(Image.fromarray(img))
   v = probs(np.array([Qbest.get((s,a),0) for a in actions]))
   a = random.choices(actions,weights=v)[0]
   obs, _, terminated, truncated, _ = env.step(a)
   done = terminated or truncated
   i += 1
env.close()

if not ims:
   raise RuntimeError("CartPole produced no frames")
output_path = output_dir / "cartpole-balance.gif"
ims[0].save(output_path,save_all=True,append_images=ims[1::2],loop=0,duration=5)
print(f"Saved {i} frames to {output_path}")