4.2  Model-Based Methods

import numpy as np
from collections import defaultdict, Counter

np.random.seed(42)

# ══════════════════════════════════════════════════════════
# Dyna-Q for Crop Planting Strategy
# States: soil_quality (0–4), season (0–3) → 20 states
# Actions: crop choice 0=wheat,1=maize,2=rice,3=fallow
# ══════════════════════════════════════════════════════════
N_SOIL    = 5; N_SEASON = 4; N_ACTIONS = 4
N_STATES  = N_SOIL * N_SEASON

def encode(soil, season): return soil * N_SEASON + season
def decode(s): return divmod(s, N_SEASON)

# True environment
YIELD = np.array([  # [soil][action]
    [2.5, 2.0, 1.5, 0.5],
    [3.2, 2.8, 2.0, 0.5],
    [4.0, 3.5, 2.8, 0.5],
    [4.8, 4.2, 3.5, 0.5],
    [5.5, 5.0, 4.2, 0.5]
])
SOIL_EFFECT = [0.1, 0.05, 0.0, -0.1]  # soil impact of crop

def env_step(state, action):
    soil, season = decode(state)
    yield_val    = YIELD[soil, action] * (1 + np.random.normal(0, 0.1))
    new_soil     = int(np.clip(soil + SOIL_EFFECT[action] * 10 + np.random.randint(-1,2), 0, 4))
    new_season   = (season + 1) % N_SEASON
    return encode(new_soil, new_season), yield_val

def run_agent(n_plan=0, n_episodes=200):
    Q = np.zeros((N_STATES, N_ACTIONS))
    # The environment is stochastic, so a model storing the LAST observed
    # transition is a single noisy draw. Store M(s,a) = (r_avg, s'_most_common)
    # instead, matching the definition in the maths box above.
    m_count = defaultdict(int)        # (s,a) -> times visited
    m_rew   = defaultdict(float)      # (s,a) -> running mean reward
    m_next  = defaultdict(Counter)    # (s,a) -> counts over observed s'
    alpha = 0.1; gamma = 0.95; epsilon = 0.3
    ep_returns = []
    
    for ep in range(n_episodes):
        state    = np.random.randint(N_STATES)
        total_r  = 0
        
        for _ in range(16):  # 4 seasons × 4 years
            # ε-greedy
            if np.random.rand() < epsilon: action = np.random.randint(N_ACTIONS)
            else: action = np.argmax(Q[state])
            
            next_state, reward = env_step(state, action)
            
            # Direct Q-update
            Q[state,action] += alpha*(reward + gamma*np.max(Q[next_state]) - Q[state,action])
            # Model update: incremental mean reward + successor histogram
            key = (state, action)
            m_count[key] += 1
            m_rew[key]   += (reward - m_rew[key]) / m_count[key]
            m_next[key][next_state] += 1
            
            # Planning: one visited (s,a) pair is enough to begin
            if n_plan > 0 and m_count:
                keys = list(m_count)
                for i in np.random.choice(len(keys), n_plan, replace=True):
                    ps, pa = keys[i]
                    pr  = m_rew[(ps,pa)]                        # r_avg
                    pns = m_next[(ps,pa)].most_common(1)[0][0]  # s'_most_common
                    Q[ps,pa] += alpha*(pr + gamma*np.max(Q[pns]) - Q[ps,pa])
            
            state = next_state; total_r += reward
        ep_returns.append(total_r)
    return Q, ep_returns

# ── Compare: Q-learning (n=0) vs Dyna-Q (n=20) ───────────────
Q_base, rew_base = run_agent(n_plan=0,  n_episodes=200)
Q_dyna, rew_dyna = run_agent(n_plan=20, n_episodes=200)

print("=== Dyna-Q vs Q-Learning — Crop Planning ===")
print(f"Q-Learning  first 20 ep: {np.mean(rew_base[:20]):.3f}")
print(f"Q-Learning  last  20 ep: {np.mean(rew_base[-20:]):.3f}")
print(f"Dyna-Q (n=20) first 20: {np.mean(rew_dyna[:20]):.3f}")
print(f"Dyna-Q (n=20) last  20: {np.mean(rew_dyna[-20:]):.3f}")

print("\nOptimal crop by soil quality (Dyna-Q):")
crops = ['Wheat','Maize','Rice','Fallow']
for soil in range(N_SOIL):
    best_actions = [np.argmax(Q_dyna[encode(soil, s)]) for s in range(N_SEASON)]
    print(f"  Soil {soil}: {[crops[a] for a in best_actions]}")

import numpy as np

np.random.seed(42)

# ══════════════════════════════════════════════════════════
# Value Iteration for Medical Treatment MDP
# States: disease stage (0–3) × response (0=good,1=poor) → 8 states
# Actions: 0=watch, 1=targeted, 2=chemo
# ══════════════════════════════════════════════════════════
N_STAGE    = 4   # stages 0–3
N_RESPONSE = 2   # 0=good, 1=poor
N_STATES   = N_STAGE * N_RESPONSE
N_ACTIONS  = 3

def encode(stage, resp): return stage * N_RESPONSE + resp
def decode(s): return divmod(s, N_RESPONSE)

gamma = 0.95
TERMINAL = encode(3, 1)  # advanced + poor response = terminal state

# ── Transition and Reward matrices ───────────────────────────
# P[s,a,s'] = transition probability
P = np.zeros((N_STATES, N_ACTIONS, N_STATES))
R = np.zeros((N_STATES, N_ACTIONS))

def fill_transitions():
    """Build P and R so every row of P sums to 1 BY CONSTRUCTION.

    Renormalising a row afterwards would silently absorb any modelling
    mistake, so we assert instead — in a worked MDP the probabilities
    should be right before they are used, not made right afterwards.
    """
    for s in range(N_STATES):
        stage, resp = decode(s)
        if s == TERMINAL:
            P[s, :, s] = 1.0; continue
        for a in range(N_ACTIONS):
            # Reward = QALYs gained; p_prog = disease advances a stage
            # p_improve / p_worsen = response flips
            if a == 0:    # watchful wait
                R[s,a]    = 0.8 if resp == 0 else 0.3
                p_prog    = 0.10 + 0.20*resp
                p_improve = 0.10
                p_worsen  = 0.15
            elif a == 1:  # targeted therapy
                R[s,a]    = 1.2 if resp == 0 else 0.7
                p_prog    = 0.05 + 0.10*resp
                p_improve = 0.30
                p_worsen  = 0.05
            else:         # chemo
                R[s,a]    = 0.9 if resp == 0 else 1.0
                p_prog    = 0.04 + 0.05*resp
                p_improve = 0.30
                p_worsen  = 0.08

            # Only one response transition is reachable from a given state:
            # a good responder can worsen, a poor responder can improve.
            p_flip = p_worsen if resp == 0 else p_improve
            p_stay = 1.0 - p_prog - p_flip
            assert p_stay >= 0, f"negative stay probability at s={s}, a={a}"

            next_stage = min(stage + 1, N_STAGE - 1)
            P[s,a, encode(next_stage, resp)] += p_prog   # stage advances
            P[s,a, encode(stage, 1 - resp)]  += p_flip   # response flips
            P[s,a, encode(stage, resp)]      += p_stay   # no change

    assert np.allclose(P.sum(axis=2), 1.0), "every transition row must sum to 1"

fill_transitions()

# ── Value Iteration ───────────────────────────────────────
V      = np.zeros(N_STATES)
theta  = 1e-8
iters  = 0
while True:
    delta = 0
    for s in range(N_STATES):
        if s == TERMINAL: continue
        v_old = V[s]
        V[s]  = max(sum(P[s,a,sp] * (R[s,a] + gamma*V[sp])
                        for sp in range(N_STATES))
                    for a in range(N_ACTIONS))
        delta = max(delta, abs(V[s] - v_old))
    iters += 1
    if delta < theta: break

# ── Extract Optimal Policy ────────────────────────────────
actions = ['WatchWait', 'Targeted', 'Chemo']
policy  = np.argmax([[sum(P[s,a,sp]*(R[s,a]+gamma*V[sp])
                          for sp in range(N_STATES))
                       for a in range(N_ACTIONS)]
                      for s in range(N_STATES)], axis=1)

print(f"=== Value Iteration — Treatment MDP ===")
print(f"Converged in {iters} iterations")
print(f"\nOptimal Policy:")
for s in range(N_STATES-1):
    stage, resp = decode(s)
    resp_str = 'Good' if resp == 0 else 'Poor'
    print(f"  Stage {stage}, {resp_str:4s}: {actions[policy[s]]:12s}  V={V[s]:.4f}")