ch5: geometry-based step; remove dependency on env.P shape

Author: srikanthbaride

ch5_monte_carlo/examples/mc_control_es_gridworld.py

@@ -7,38 +7,41 @@
 ACTIONS = [(0, 1), (0, -1), (1, 0), (-1, 0)]  # R, L, D, U
 
 def _is_terminal(env: GridWorld4x4, s) -> bool:
-    """Robust terminal check even if env.is_terminal is absent."""
     if hasattr(env, "is_terminal"):
         return bool(env.is_terminal(s))
     st = s if isinstance(s, tuple) else env.i2s[int(s)]
     return st == env.goal
 
 def _step(env: GridWorld4x4, s, a):
-    """Use env.step if present; otherwise use P/R (deterministic)."""
+    """Robust step that does NOT depend on env.P; uses geometry."""
     if hasattr(env, "step"):
         return env.step(s, a)
-    s_idx = env.s2i[s] if isinstance(s, tuple) else int(s)
-    probs = env.P[s_idx, a]
-    sp_idx = int(np.argmax(probs))
-    r = float(env.R[s_idx, a, sp_idx])
-    return env.i2s[sp_idx], r
+    st = s if isinstance(s, tuple) else env.i2s[int(s)]
+    i, j = st
+    di, dj = ACTIONS[a]
+    # infer grid size
+    n = getattr(env, "n", int(round(len(env.S) ** 0.5)))
+    ni, nj = i + di, j + dj
+    if not (0 <= ni < n and 0 <= nj < n):
+        ni, nj = i, j  # wall -> stay
+    sp = (ni, nj)
+    # reward: step cost unless entering goal, then 0.0 (matches your ch4 tests)
+    step_reward = float(getattr(env, "step_reward", -1.0))
+    r = 0.0 if sp == getattr(env, "goal", (0, 3)) else step_reward
+    return sp, r
 
 def _greedy_action(q_row: np.ndarray) -> int:
     return int(np.argmax(q_row))
 
 def generate_episode_es(env: GridWorld4x4, Q: np.ndarray, gamma: float, max_steps: int = 10000):
-    """
-    Exploring starts: start from a random non-terminal state and random action,
-    then follow greedy policy thereafter. Returns (states, actions, returns).
-    """
     rng = np.random.default_rng()
     non_terminal = [s for s in env.S if not _is_terminal(env, s)]
     s = non_terminal[rng.integers(len(non_terminal))]
     a = int(rng.integers(len(env.A)))
 
     states = [s]
     actions = [a]
-    rewards = [0.0]  # so rewards[t+1] aligns with action at t
+    rewards = [0.0]
 
     steps = 0
     while not _is_terminal(env, s) and steps < max_steps:
@@ -52,7 +55,6 @@ def generate_episode_es(env: GridWorld4x4, Q: np.ndarray, gamma: float, max_step
         actions.append(a)
         steps += 1
 
-    # first-visit returns
     G = 0.0
     returns = np.zeros(len(states), dtype=float)
     for t in range(len(states) - 1, -1, -1):
@@ -61,30 +63,23 @@ def generate_episode_es(env: GridWorld4x4, Q: np.ndarray, gamma: float, max_step
     return states, actions, returns
 
 def mc_es_control(env: GridWorld4x4, episodes: int = 1500, gamma: float | None = None, seed: int | None = None):
-    """
-    On-policy Monte Carlo control with Exploring Starts (ES).
-    Returns:
-        Q:  (S,A) action-value table
-        pi: (S,A) deterministic greedy policy derived from Q
-    """
     if seed is not None:
         np.random.seed(seed)
     if gamma is None:
-        gamma = float(env.gamma)
+        gamma = float(getattr(env, "gamma", 1.0))
 
     S, A = len(env.S), len(env.A)
     Q = np.zeros((S, A), dtype=float)
-    N = np.zeros((S, A), dtype=float)  # first-visit counts
+    N = np.zeros((S, A), dtype=float)
 
     for _ in range(episodes):
         states, actions, returns = generate_episode_es(env, Q, gamma)
         seen = set()
         for t, (s, a) in enumerate(zip(states, actions)):
             s_idx = env.s2i[s]
-            key = (s_idx, a)
-            if key in seen:
+            if (s_idx, a) in seen:
                 continue
-            seen.add(key)
+            seen.add((s_idx, a))
             G = returns[t]
             N[s_idx, a] += 1.0
             alpha = 1.0 / N[s_idx, a]

ch5_monte_carlo/examples/mc_control_onpolicy_gridworld.py

@@ -12,20 +12,26 @@ def _is_terminal(env: GridWorld4x4, s) -> bool:
     return st == env.goal
 
 def _step(env: GridWorld4x4, s, a):
+    """Robust step that does NOT depend on env.P; uses geometry."""
     if hasattr(env, "step"):
         return env.step(s, a)
-    s_idx = env.s2i[s] if isinstance(s, tuple) else int(s)
-    probs = env.P[s_idx, a]
-    sp_idx = int(np.argmax(probs))
-    r = float(env.R[s_idx, a, sp_idx])
-    return env.i2s[sp_idx], r
+    st = s if isinstance(s, tuple) else env.i2s[int(s)]
+    i, j = st
+    di, dj = ACTIONS[a]
+    n = getattr(env, "n", int(round(len(env.S) ** 0.5)))
+    ni, nj = i + di, j + dj
+    if not (0 <= ni < n and 0 <= nj < n):
+        ni, nj = i, j
+    sp = (ni, nj)
+    step_reward = float(getattr(env, "step_reward", -1.0))
+    r = 0.0 if sp == getattr(env, "goal", (0, 3)) else step_reward
+    return sp, r
 
 def _epsilon_greedy(q_row: np.ndarray, epsilon: float, rng: np.random.Generator) -> int:
     return int(rng.integers(len(q_row))) if rng.random() < epsilon else int(np.argmax(q_row))
 
 def generate_episode_onpolicy(env: GridWorld4x4, Q: np.ndarray, epsilon: float,
                               rng: np.random.Generator, max_steps: int = 10000):
-    """Start from a random non-terminal state; follow ε-greedy w.r.t. Q."""
     non_terminal = [s for s in env.S if not _is_terminal(env, s)]
     s = non_terminal[rng.integers(len(non_terminal))]
 
@@ -40,7 +46,6 @@ def generate_episode_onpolicy(env: GridWorld4x4, Q: np.ndarray, epsilon: float,
         states.append(s)
         steps += 1
 
-    # first-visit returns
     gamma = float(getattr(env, "gamma", 1.0))
     G = 0.0
     returns = np.zeros(len(actions), dtype=float)
@@ -52,12 +57,6 @@ def generate_episode_onpolicy(env: GridWorld4x4, Q: np.ndarray, epsilon: float,
 def mc_control_onpolicy(env: GridWorld4x4, episodes: int = 5000,
                         epsilon: float = 0.1, gamma: float | None = None,
                         seed: int | None = None):
-    """
-    On-policy MC control using ε-greedy behavior/target policy (no ES).
-    Returns:
-        Q:  (S,A)
-        pi: (S,A) deterministic greedy policy
-    """
     rng = np.random.default_rng(seed)
     S, A = len(env.S), len(env.A)
     if gamma is None:
@@ -71,10 +70,9 @@ def mc_control_onpolicy(env: GridWorld4x4, episodes: int = 5000,
         seen = set()
         for t, (s, a) in enumerate(zip(states, actions)):
             s_idx = env.s2i[s]
-            key = (s_idx, a)
-            if key in seen:
+            if (s_idx, a) in seen:
                 continue
-            seen.add(key)
+            seen.add((s_idx, a))
             G = returns[t]
             N[s_idx, a] += 1.0
             alpha = 1.0 / N[s_idx, a]