Q-Learning

# Q-learning on the same grid, with NO transition model given to the agent: it
# learns only from experience. The answer is checked against value iteration,
# which is allowed to see the model, and then SARSA is put beside it.
COLS, ROWS = 4, 3
WALL = (1, 1)
TERMINAL = {(3, 2): 1.0, (3, 1): -1.0}
START = (0, 0)
LIVING, GAMMA = -0.04, 1.0
ACTIONS = ["N", "S", "E", "W"]
MOVE = {"N": (0, 1), "S": (0, -1), "E": (1, 0), "W": (-1, 0)}
SIDEWAYS = {"N": ("W", "E"), "S": ("E", "W"), "E": ("N", "S"), "W": ("S", "N")}
P_INTENDED, P_SLIP = 0.8, 0.1
STATES = [(c, r) for r in range(ROWS) for c in range(COLS) if (c, r) != WALL]
INNER = [s for s in STATES if s not in TERMINAL]

def lcg(seed):
    x = seed
    while True:
        x = (1664525 * x + 1013904223) % (2 ** 32)
        yield x / 2 ** 32

def legal(s):
    c, r = s
    return 0 <= c < COLS and 0 <= r < ROWS and (c, r) != WALL

def move(s, a):
    d = MOVE[a]
    t = (s[0] + d[0], s[1] + d[1])
    return t if legal(t) else s

def transitions(s, a):
    left, right = SIDEWAYS[a]
    out = {}
    for p, act in ((P_INTENDED, a), (P_SLIP, left), (P_SLIP, right)):
        t = move(s, act)
        out[t] = out.get(t, 0.0) + p
    return [(p, t) for t, p in sorted(out.items())]

def draw(grid, fmt="%8s"):
    for r in range(ROWS - 1, -1, -1):
        print(("   " + "".join(fmt % ("#" if (c, r) == WALL else grid.get((c, r), "."))
                               for c in range(COLS))).rstrip())

# ---- the answer the agent must reach, computed WITH the model ---------------
def value_iteration():
    V = {s: 0.0 for s in STATES}
    for t, v in TERMINAL.items():
        V[t] = v
    while True:
        delta, new = 0.0, dict(V)
        for s in INNER:
            best = max(sum(p * (LIVING + GAMMA * V[t]) for p, t in transitions(s, a))
                       for a in ACTIONS)
            delta = max(delta, abs(best - V[s]))
            new[s] = best
        V = new
        if delta < 1e-12:
            return V

V_STAR = value_iteration()
POLICY_STAR = {s: max(ACTIONS,
                      key=lambda a: sum(p * (LIVING + GAMMA * V_STAR[t])
                                        for p, t in transitions(s, a)))
               for s in INNER}

print("THE AGENT IS GIVEN NOTHING. it does not know the slip probabilities, the")
print("step cost, where the pillar is, or where the exits are. it knows the four")
print("actions and it can see which square it is in. everything else it must")
print("learn from what happens.")
print()
print("   Q(s, a) is the value of TAKING action a in state s and behaving well")
print("   afterwards. the update, after every single step:")
print()
print("      Q(s,a) <- Q(s,a) + alpha * [ r + gamma * max over a' of Q(s',a')")
print("                                   - Q(s,a) ]")
print()
print("   the bracket is the TEMPORAL DIFFERENCE ERROR: what the step actually")
print("   suggests the value is, minus what was believed. alpha decides how much")
print("   of that correction to apply.")
print()
print("   NOTE WHERE THE MODEL WOULD HAVE BEEN. the Bellman equation summed over")
print("   every next state weighted by P(s'|s,a). this update uses the ONE next")
print("   state that actually happened. the average is accumulated over many")
print("   visits instead of being computed, which is why no model is needed.")
print()

def step(gen, s, a):
    """The environment. The agent may call it; it may not look inside."""
    r = next(gen)
    left, right = SIDEWAYS[a]
    act = a if r < P_INTENDED else (left if r < P_INTENDED + P_SLIP else right)
    t = move(s, act)
    return t, (TERMINAL[t] if t in TERMINAL else LIVING)

def learn(episodes, alpha, epsilon, seed, sarsa=False, random_behaviour=False):
    gen = lcg(seed)
    Q = {(s, a): 0.0 for s in INNER for a in ACTIONS}
    for ep in range(episodes):
        s = INNER[min(int(next(gen) * len(INNER)), len(INNER) - 1)]

        def choose(state):
            if random_behaviour or next(gen) < epsilon:
                return ACTIONS[min(int(next(gen) * 4), 3)]
            return max(ACTIONS, key=lambda a: Q[(state, a)])

        a = choose(s)
        for _ in range(200):
            t, r = step(gen, s, a)
            if t in TERMINAL:
                Q[(s, a)] += alpha * (r - Q[(s, a)])
                break
            a2 = choose(t)
            nxt = Q[(t, a2)] if sarsa else max(Q[(t, b)] for b in ACTIONS)
            Q[(s, a)] += alpha * (r + GAMMA * nxt - Q[(s, a)])
            s, a = t, a2
    return Q

def evaluate(policy):
    """THE AUTHOR'S measuring instrument, not the agent's: what is a learnt
    policy actually worth, computed with the model the agent never saw."""
    V = {s: 0.0 for s in STATES}
    for t, v in TERMINAL.items():
        V[t] = v
    while True:
        delta, new = 0.0, dict(V)
        for s in INNER:
            x = sum(p * (LIVING + GAMMA * V[t]) for p, t in transitions(s, policy[s]))
            delta = max(delta, abs(x - V[s]))
            new[s] = x
        V = new
        if delta < 1e-12:
            return V

def report(Q):
    V = {s: max(Q[(s, a)] for a in ACTIONS) for s in INNER}
    pol = {s: max(ACTIONS, key=lambda a: Q[(s, a)]) for s in INNER}
    err = max(abs(V[s] - V_STAR[s]) for s in INNER)
    agree = sum(1 for s in INNER if pol[s] == POLICY_STAR[s])
    return V, pol, err, agree, evaluate(pol)[START]

print("LEARNING. alpha = 0.10, epsilon = 0.20, episodes from random squares.")
print("   the last column is what matters: not whether the learnt policy MATCHES")
print("   the optimal one, but what it is WORTH. the optimal policy is worth")
print("   %+.4f from the start square." % V_STAR[START])
print()
print("   episodes | largest value error | actions matching | policy is worth")
Q = None
for k in (10, 100, 1000, 10000, 100000):
    Q = learn(k, 0.10, 0.20, 99991)
    V, pol, err, agree, worth = report(Q)
    print("   %8d | %19.4f | %5d of %d       | %+.4f"
          % (k, err, agree, len(INNER), worth))
print()
V, pol, err, agree, worth = report(Q)
print("   after 100000 episodes, the values Q-learning believes:")
draw({s: "%+.4f" % V[s] for s in INNER}, fmt="%10s")
print("   what value iteration computed WITH the model:")
draw({s: "%+.4f" % V_STAR[s] for s in INNER}, fmt="%10s")
print("   the policy it acts on, then the optimal one:")
g1, g2 = dict(pol), dict(POLICY_STAR)
for t, v in TERMINAL.items():
    g1[t] = g2[t] = "%+.0f" % v
draw(g1)
print("   ---")
draw(g2)
print("   %d of %d actions agree, and the learnt policy is worth %+.4f against"
      % (agree, len(INNER), worth))
print("   %+.4f for the optimal one: a loss of %.4f."
      % (V_STAR[START], V_STAR[START] - worth))
print()
print("   THE VALUES ARE STILL WRONG BY %.4f AND THE POLICY IS RIGHT. with a" % err)
print("   CONSTANT alpha the updates never stop moving: every visit pulls the")
print("   estimate part of the way towards a single sampled outcome, so it keeps")
print("   bouncing around the true value for ever. the bounce is roughly")
print("   proportional to alpha.")
print()
print("   ALPHA AGAINST THE ERROR, at 20000 episodes each:")
print("   alpha | largest value error | actions matching | policy is worth")
for al in (0.50, 0.20, 0.10, 0.05, 0.02):
    Qa = learn(20000, al, 0.20, 99991)
    V2, pol2, err2, agree2, worth2 = report(Qa)
    print("   %5.2f | %19.4f | %5d of %d       | %+.4f"
          % (al, err2, agree2, len(INNER), worth2))
print("   a large alpha learns quickly and settles badly; a small one settles")
print("   well and learns slowly. TO CONVERGE EXACTLY, alpha must SHRINK over")
print("   time: the standard condition is that the alphas sum to infinity, so")
print("   every estimate can still be moved anywhere, while their squares sum to")
print("   something finite, so the bouncing dies away.")
print()

print("OFF-POLICY, which is the property that matters. run it again with the")
print("agent choosing UNIFORMLY AT RANDOM at every step, never once using what")
print("it has learnt:")
Qr = learn(100000, 0.05, 1.0, 99991, random_behaviour=True)
Vr, polr, errr, agreer, worthr = report(Qr)
draw({s: polr[s] for s in INNER})
print("   largest value error %.4f, %d of %d actions matching, and the policy it"
      % (errr, agreer, len(INNER)))
print("   learnt is worth %+.4f against the optimal %+.4f."
      % (worthr, V_STAR[START]))
print("   AN AGENT THAT NEVER ONCE ACTED ON WHAT IT KNEW LEARNT THE OPTIMAL")
print("   POLICY. that is what OFF-POLICY means: the policy learnt")
print("   about is not the policy used to act, because the update takes the")
print("   MAXIMUM over the next state's actions rather than the action the agent")
print("   will actually take. it is also why the method can learn from logged")
print("   data, from a human demonstrator, or from an old policy's records.")
print()

print("SARSA, the on-policy relative. one symbol changes in the update:")
print("      Q-learning : ... + gamma * MAX over a' of Q(s',a')")
print("      SARSA      : ... + gamma * Q(s', a'), where a' is the action the")
print("                   agent will ACTUALLY take next")
print("   so SARSA learns the value of the policy it is following, exploration")
print("   and all. the name is the five things its update uses: s, a, r, s', a'.")
print()
print("   both on this grid, alpha 0.10, epsilon 0.20, 100000 episodes:")
print("   method     | largest value error | actions matching | policy is worth")
for name, Qx in (("Q-learning", Q), ("SARSA", learn(100000, 0.10, 0.20, 99991,
                                                    sarsa=True))):
    Vx, polx, errx, agreex, worthx = report(Qx)
    print("   %-10s | %19.4f | %5d of %d       | %+.4f"
          % (name, errx, agreex, len(INNER), worthx))
print("   on a forgiving grid they end in much the same place. the difference")
print("   needs an environment where exploring is DANGEROUS.")
print()

print("SO HERE IS ONE: A CLIFF. a %d by %d corridor, start at the bottom left,"
      % (6, 3))
print("goal at the bottom right, and every square between them along the bottom")
print("is a fall costing -100 that ends the episode. every other step costs -1.")
print("THIS FLOOR IS NOT SLIPPERY: the danger is entirely the agent's own")
print("exploration.")

CW, CH = 6, 3
CLIFF = [(c, 0) for c in range(1, CW - 1)]
CSTART, CGOAL = (0, 0), (CW - 1, 0)
CINNER = [(c, r) for r in range(CH) for c in range(CW) if (c, r) not in CLIFF
          and (c, r) != CGOAL]

def cliff_step(s, a):
    d = MOVE[a]
    t = (s[0] + d[0], s[1] + d[1])
    if not (0 <= t[0] < CW and 0 <= t[1] < CH):
        t = s
    if t in CLIFF:
        return CSTART, -100.0, True
    if t == CGOAL:
        return t, -1.0, True
    return t, -1.0, False

def cliff_run(episodes, alpha, epsilon, seed, sarsa):
    gen = lcg(seed)
    Q = {(s, a): 0.0 for s in CINNER for a in ACTIONS}
    rewards = []
    for ep in range(episodes):
        s = CSTART
        pick = (lambda st: ACTIONS[min(int(next(gen) * 4), 3)]
                if next(gen) < epsilon
                else max(ACTIONS, key=lambda a: Q[(st, a)]))
        a = pick(s)
        total = 0.0
        for _ in range(300):
            t, r, done = cliff_step(s, a)
            total += r
            if done:
                Q[(s, a)] += alpha * (r - Q[(s, a)])
                break
            a2 = pick(t)
            nxt = Q[(t, a2)] if sarsa else max(Q[(t, b)] for b in ACTIONS)
            Q[(s, a)] += alpha * (r + GAMMA * nxt - Q[(s, a)])
            s, a = t, a2
        rewards.append(total)
    return Q, rewards

print("   method     | mean reward per episode over the last 500 | route taken")
for name, flag in (("Q-learning", False), ("SARSA", True)):
    Q, rewards = cliff_run(3000, 0.10, 0.10, 20260930, flag)
    greedy_pol = {s: max(ACTIONS, key=lambda a: Q[(s, a)]) for s in CINNER}
    s, route = CSTART, [CSTART]
    for _ in range(20):
        t, r, done = cliff_step(s, greedy_pol[s])
        route.append(t)
        s = t
        if done:
            break
    rows = sorted({p[1] for p in route})
    print("   %-10s | %41.2f | rows used: %s"
          % (name, sum(rewards[-500:]) / 500, ", ".join(str(r) for r in rows)))
print()
print("   Q-LEARNING learns the OPTIMAL route, which runs along row 1, directly")
print("   above the cliff. SARSA learns a route further from the edge.")
print("   neither is mistaken. Q-learning's update takes the maximum, so it")
print("   learns the value of the best route ASSUMING NO EXPLORATION; while it")
print("   is still exploring, one random step in ten sends it over the edge and")
print("   it collects a worse reward. SARSA's update uses the action it will")
print("   really take, so the cost of its own exploration is built into the")
print("   values, and it prefers a route where a random step is survivable.")
print()
print("   SO: Q-LEARNING LEARNS THE OPTIMAL POLICY AND SARSA LEARNS THE BEST")
print("   POLICY FOR AN AGENT THAT EXPLORES. if the exploration is switched off")
print("   at the end, Q-learning's answer is better. if the agent must keep")
print("   exploring in the real world, and falling off the cliff is a real")
print("   robot falling off a real table, SARSA's is.")