Add Monte Carlo Control examples (MC-ES, On-policy ε-soft) with GridWorld + tests

Author: srikanthbaride

ch5_monte_carlo/examples/mc_control_es_gridworld.py

@@ -0,0 +1,73 @@
+# ch5_monte_carlo/examples/mc_control_es_gridworld.py
+import numpy as np
+from collections import defaultdict
+from ch2_rl_formulation.gridworld import GridWorld4x4  # your existing env
+
+ACTIONS = [0, 1, 2, 3]  # R, L, D, U (consistent with your env)
+
+def random_start(env: GridWorld4x4):
+    s = env.S[np.random.randint(len(env.S))]
+    a = np.random.choice(ACTIONS)
+    return s, a
+
+def step(env: GridWorld4x4, s, a):
+    # env exposes P[s_idx][a] -> list of (prob, s', r)
+    s_idx = env.s2i[s]
+    trans = env.P[s_idx][a]
+    probs = [p for (p, _, _) in trans]
+    i = np.random.choice(len(trans), p=probs)
+    _, sp_idx, r = trans[i]
+    return env.i2s[sp_idx], r
+
+def generate_episode_es(env, pi, gamma=1.0):
+    s0, a0 = random_start(env)  # Exploring start
+    episode = [(s0, a0, None)]
+    s = s0
+    a = a0
+    done = (s == env.goal)
+    rewards = []
+    while not done:
+        sp, r = step(env, s, a)
+        rewards.append(r)
+        s = sp
+        if s == env.goal:
+            break
+        # follow current policy after the start
+        a = pi[s]
+        episode.append((s, a, None))
+    return episode, rewards  # rewards aligned with transitions
+
+def mc_es_control(env, episodes=5000, gamma=0.9):
+    Q = defaultdict(lambda: 0.0)
+    N = defaultdict(int)
+    # start with arbitrary deterministic policy
+    pi = {s: np.random.choice(ACTIONS) for s in env.S}
+    for _ in range(episodes):
+        ep, rewards = generate_episode_es(env, pi, gamma)
+        G = 0.0
+        visited = set()
+        # process backwards
+        for t in range(len(ep) - 1, -1, -1):
+            s, a, _ = ep[t]
+            G = gamma * G + rewards[t] if t < len(rewards) else G
+            if (s, a) not in visited:
+                N[(s, a)] += 1
+                Q[(s, a)] += (G - Q[(s, a)]) / N[(s, a)]  # incremental mean
+                visited.add((s, a))
+                # greedy improvement
+                best_a = max(ACTIONS, key=lambda act: Q[(s, act)])
+                pi[s] = best_a
+    return Q, pi
+
+if __name__ == "__main__":
+    np.random.seed(0)
+    env = GridWorld4x4(step_reward=0.0, goal=(0, 3))
+    Q, pi = mc_es_control(env, episodes=3000, gamma=0.9)
+    # print a small slice of the learned greedy policy arrows
+    arrows = {0: "→", 1: "←", 2: "↓", 3: "↑"}
+    for i in range(env.n):
+        row = []
+        for j in range(env.n):
+            s = (i, j)
+            row.append(" G " if s == env.goal else f" {arrows[pi[s]]} ")
+        print("".join(row))

ch5_monte_carlo/examples/mc_control_onpolicy_gridworld.py

@@ -0,0 +1,79 @@
+# ch5_monte_carlo/examples/mc_control_onpolicy_gridworld.py
+
+import numpy as np
+from collections import defaultdict
+from ch2_rl_formulation.gridworld import GridWorld4x4  # reuse Chapter 2 env
+
+ACTIONS = [0, 1, 2, 3]  # R, L, D, U
+
+def step(env: GridWorld4x4, s, a):
+    """Sample one step given state and action from env's transition model."""
+    s_idx = env.s2i[s]
+    trans = env.P[s_idx][a]
+    probs = [p for (p, _, _) in trans]
+    i = np.random.choice(len(trans), p=probs)
+    _, sp_idx, r = trans[i]
+    return env.i2s[sp_idx], r
+
+def generate_episode(env, pi, epsilon=0.1, gamma=0.9):
+    """Generate one episode following epsilon-soft policy pi."""
+    s = env.S[np.random.randint(len(env.S))]
+    episode, rewards = [], []
+    done = (s == env.goal)
+    while not done:
+        # choose action epsilon-greedily
+        if np.random.rand() < epsilon:
+            a = np.random.choice(ACTIONS)
+        else:
+            a = max(ACTIONS, key=lambda act: pi[(s, act)])
+        sp, r = step(env, s, a)
+        episode.append((s, a))
+        rewards.append(r)
+        s = sp
+        done = (s == env.goal)
+    return episode, rewards
+
+def mc_control_onpolicy(env, episodes=5000, gamma=0.9, epsilon=0.1):
+    """On-policy MC control with epsilon-soft policies."""
+    Q = defaultdict(float)
+    N = defaultdict(int)
+    pi = {(s, a): 1.0/len(ACTIONS) for s in env.S for a in ACTIONS}  # uniform start
+
+    for _ in range(episodes):
+        ep, rewards = generate_episode(env, pi, epsilon, gamma)
+        G, visited = 0.0, set()
+        # backward return computation
+        for t in range(len(ep) - 1, -1, -1):
+            s, a = ep[t]
+            G = gamma * G + rewards[t]
+            if (s, a) not in visited:
+                N[(s, a)] += 1
+                Q[(s, a)] += (G - Q[(s, a)]) / N[(s, a)]
+                visited.add((s, a))
+                # policy improvement: epsilon-greedy
+                best_a = max(ACTIONS, key=lambda act: Q[(s, act)])
+                for act in ACTIONS:
+                    if act == best_a:
+                        pi[(s, act)] = 1 - epsilon + epsilon/len(ACTIONS)
+                    else:
+                        pi[(s, act)] = epsilon/len(ACTIONS)
+    return Q, pi
+
+if __name__ == "__main__":
+    np.random.seed(1)
+    env = GridWorld4x4(step_reward=0.0, goal=(0, 3))
+    Q, pi = mc_control_onpolicy(env, episodes=3000, gamma=0.9, epsilon=0.1)
+
+    # Print learned greedy policy (arrows)
+    arrows = {0: "→", 1: "←", 2: "↓", 3: "↑"}
+    for i in range(env.n):
+        row = []
+        for j in range(env.n):
+            s = (i, j)
+            if s == env.goal:
+                row.append(" G ")
+            else:
+                # choose most probable action
+                best_a = max(ACTIONS, key=lambda a: pi[(s, a)])
+                row.append(f" {arrows[best_a]} ")
+        print("".join(row))

ch5_monte_carlo/tests/__init__.py

@@ -0,0 +1,77 @@
+# ch5_monte_carlo/tests/test_mc_control.py
+import numpy as np
+
+from ch2_rl_formulation.gridworld import GridWorld4x4
+from ch5_monte_carlo.examples.mc_control_es_gridworld import mc_es_control
+from ch5_monte_carlo.examples.mc_control_onpolicy_gridworld import mc_control_onpolicy, ACTIONS
+
+ARROWS = {0: "→", 1: "←", 2: "↓", 3: "↑"}
+
+def rollout_greedy_es(env: GridWorld4x4, pi, max_steps=64):
+    """Roll out deterministic greedy policy 'pi' (state -> action)."""
+    s = env.S[np.random.randint(len(env.S))]
+    steps = 0
+    while s != env.goal and steps < max_steps:
+        a = pi[s]
+        s_idx = env.s2i[s]
+        trans = env.P[s_idx][a]
+        probs = [p for (p, _, _) in trans]
+        i = np.random.choice(len(trans), p=probs)
+        _, sp_idx, _ = trans[i]
+        s = env.i2s[sp_idx]
+        steps += 1
+    return s == env.goal, steps
+
+def greedy_from_soft(pi_soft, s):
+    """Pick argmax_a pi(a|s) from dict keyed by (s,a)."""
+    return max(ACTIONS, key=lambda a: pi_soft[(s, a)])
+
+def rollout_greedy_from_soft(env: GridWorld4x4, pi_soft, max_steps=64):
+    """Roll out using greedy action from ε-soft policy probabilities."""
+    s = env.S[np.random.randint(len(env.S))]
+    steps = 0
+    while s != env.goal and steps < max_steps:
+        a = greedy_from_soft(pi_soft, s)
+        s_idx = env.s2i[s]
+        trans = env.P[s_idx][a]
+        probs = [p for (p, _, _) in trans]
+        i = np.random.choice(len(trans), p=probs)
+        _, sp_idx, _ = trans[i]
+        s = env.i2s[sp_idx]
+        steps += 1
+    return s == env.goal, steps
+
+def success_rate(trial_fn, trials=100):
+    succ = 0
+    total_steps = 0
+    for _ in range(trials):
+        ok, steps = trial_fn()
+        succ += int(ok)
+        total_steps += steps
+    return succ / trials, total_steps / trials
+
+def test_mc_es_gridworld_reaches_goal():
+    np.random.seed(0)
+    env = GridWorld4x4(step_reward=0.0, goal=(0, 3))
+
+    # Train MC-ES (keep episodes modest so CI stays fast)
+    Q, pi = mc_es_control(env, episodes=1500, gamma=0.9)
+
+    sr, avg_steps = success_rate(lambda: rollout_greedy_es(env, pi), trials=100)
+
+    # Expect high success and reasonable path length
+    assert sr >= 0.9, f"MC-ES success rate too low: {sr:.2f}"
+    assert avg_steps <= 25, f"MC-ES average steps too high: {avg_steps:.1f}"
+
+def test_mc_onpolicy_gridworld_reaches_goal():
+    np.random.seed(1)
+    env = GridWorld4x4(step_reward=0.0, goal=(0, 3))
+
+    # Train on-policy MC with ε-soft behavior
+    Q, pi_soft = mc_control_onpolicy(env, episodes=2000, gamma=0.9, epsilon=0.1)
+
+    sr, avg_steps = success_rate(lambda: rollout_greedy_from_soft(env, pi_soft), trials=100)
+
+    # Expect robust success with a slightly looser bound than MC-ES
+    assert sr >= 0.85, f"On-policy MC success rate too low: {sr:.2f}"
+    assert avg_steps <= 28, f"On-policy MC average steps too high: {avg_steps:.1f}"