Monte Carlo with SimplePURT#
Environment#
Same environment from the Dynamic Programming demo
# Simple PURT Environment
import gymnasium as gym
from gymnasium import spaces
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from enum import Enum
class Action(Enum):
Right = 0
Up = 1
Left = 2
Down = 3
class SimplePURT(gym.Env):
# PYLON_COORDINATES: Bottom left corner is (0,0), Top right corner is (1,1)
def __init__(self, GRID_N=10, MAX_STEPS=100, REWARDS={"neutral": 0, "midpoint": 1, "endpoint": 1, "pylon": -1}, NOISE=0, INIT_STATE=np.asarray([0, 0]), PYLON_COORDINATES=np.array([[0.2, 0.2],[0.2,0.8],[0.8,0.8],[0.8,0.2]])):
super(SimplePURT, self).__init__()
self.GRID_N = GRID_N
self.MAX_STEPS = MAX_STEPS
self.REWARDS = REWARDS
self.NOISE = NOISE # [0,1] -- 0: No Noise, 1: 100% Noise
self.INIT_STATE = INIT_STATE.copy()
self.PYLON_GRID_COORDINATES = np.round(PYLON_COORDINATES * (GRID_N - 1)).astype(int) # Convert pylon coordinates to grid coordinates
self.PYLON_GRID_INDICES = self.PYLON_GRID_COORDINATES[:, 0] + self.PYLON_GRID_COORDINATES[:, 1] * GRID_N # Convert grid coordinates to grid indices
self.PYLON_MIDPOINT_COORDINATES = np.zeros_like(PYLON_COORDINATES)
for i in range(PYLON_COORDINATES.shape[0]):
self.PYLON_MIDPOINT_COORDINATES[i] = 0.5 * (PYLON_COORDINATES[i] + PYLON_COORDINATES[(i + 1) % PYLON_COORDINATES.shape[0]]) # Midpoint between pylon i and pylon (i+1)
self.PYLON_MIDPOINT_GRID_COORDINATES = np.round(self.PYLON_MIDPOINT_COORDINATES * (GRID_N - 1)).astype(int) # Convert midpoint coordinates to grid coordinates
self.PYLON_MIDPOINT_INDICES = self.PYLON_MIDPOINT_GRID_COORDINATES[:, 0] + self.PYLON_MIDPOINT_GRID_COORDINATES[:, 1] * GRID_N # Convert grid coordinates to grid indices
self.observation_space = spaces.MultiDiscrete([
GRID_N * GRID_N, # Grid index number
len(self.PYLON_MIDPOINT_INDICES) + 1 # Number of midpoints visited
])
self.action_space = spaces.Discrete(4) # Action space
self.state = INIT_STATE.copy() # Initial state
self.n_steps = 0 # Keep track of the number of steps taken
self.state_history = [INIT_STATE.copy()] # Keep track of the state history
self.action_history = [] # Keep track of the action history
self.reward_history = [] # Keep track of the reward history
def reset(self):
self.state = self.INIT_STATE.copy()
self.n_steps = 0
self.state_history = [self.state.copy()]
self.action_history = []
self.reward_history = []
return self.get_observation(), self.get_info()
def step(self, action, ignore_truncation=False):
new_state = self.state.copy()
reward = self.REWARDS["neutral"]
done = False
truncated = False
# 2D movement
new_state[0] = self.get_new_position(action)
# Reward
if new_state[0] != self.state[0]:
if new_state[0] in self.PYLON_GRID_INDICES:
reward += self.REWARDS["pylon"] # Penalty for hitting a pylon
done = True
elif new_state[1] < len(self.PYLON_MIDPOINT_INDICES):
if new_state[0] == self.PYLON_MIDPOINT_INDICES[new_state[1]]:
reward += self.REWARDS["midpoint"] # Reward for visiting midpoints in order
new_state[1] += 1 # Move to the next midpoint
elif new_state[0] == self.INIT_STATE[0]:
reward += self.REWARDS["endpoint"] # Reward for returning to initial state
done = True
else:
done = True # If the agent tries to move out of bounds, end the episode
self.state = new_state
self.state_history.append(new_state.copy())
self.action_history.append(action)
self.reward_history.append(reward)
if not ignore_truncation:
self.n_steps += 1
truncated = self.n_steps >= self.MAX_STEPS # Truncate the episode after MAX_STEPS
return self.get_observation(), reward, done, truncated, self.get_info()
def get_observation(self):
return self.state
def get_info(self):
return {"state_history": self.state_history, "action_history": self.action_history, "reward_history": self.reward_history}
def get_possible_next_states(self, action):
# 1. Determine probabilities for noise tendencies
# Noise = 0 -> prob_intended = 1.0, prob_dev = 0.0
p_dev = self.NOISE / 3.0
p_intended = 1.0 - (2.0 * p_dev)
# Map tendencies (-1: left shift, 0: intended, 1: right shift) to probabilities
tendencies = [
(0, p_intended),
(-1, p_dev),
(1, p_dev)
]
possible_transitions = {} # Key: (next_position, next_midpoint_count), Value: [prob, reward]
for tendency, prob in tendencies:
if prob <= 0:
continue # Skip zero-probability transitions
# Determine effective action based on tendency logic from get_new_position
transition = Action(action)
if tendency != 0:
match transition:
case Action.Up:
transition = Action.Right if tendency == 1 else Action.Left
case Action.Down:
transition = Action.Left if tendency == 1 else Action.Right
case Action.Left:
transition = Action.Up if tendency == 1 else Action.Down
case Action.Right:
transition = Action.Down if tendency == 1 else Action.Up
# Compute next grid position index
curr_pos = self.state[0]
match transition:
case Action.Up:
new_pos = (curr_pos + self.GRID_N) if curr_pos < self.GRID_N * (self.GRID_N - 1) else curr_pos
case Action.Down:
new_pos = (curr_pos - self.GRID_N) if curr_pos >= self.GRID_N else curr_pos
case Action.Left:
new_pos = (curr_pos - 1) if curr_pos % self.GRID_N != 0 else curr_pos
case Action.Right:
new_pos = (curr_pos + 1) if (curr_pos + 1) % self.GRID_N != 0 else curr_pos
# Calculate reward and next state logic matching step()
new_state = self.state.copy()
reward = self.REWARDS["neutral"]
done = False
if new_pos != curr_pos:
new_state[0] = new_pos
if new_pos in self.PYLON_GRID_INDICES:
reward += self.REWARDS["pylon"]
done = True
elif new_state[1] < len(self.PYLON_MIDPOINT_INDICES):
if new_pos == self.PYLON_MIDPOINT_INDICES[new_state[1]]:
reward += self.REWARDS["midpoint"]
new_state[1] += 1
elif new_pos == self.INIT_STATE[0]:
reward += self.REWARDS["endpoint"]
else:
# Out-of-bounds boundary bump
done = True
pass
state_key = (int(new_state[0]), int(new_state[1]))
# Aggregate probabilities if different noise tendencies lead to identical outcome states
if state_key in possible_transitions:
possible_transitions[state_key][0] += prob
else:
possible_transitions[state_key] = [prob, reward, done]
# Convert results into required format [[next_state_array, prob, reward], ...]
results = []
for (pos, midpoints_visited), (p, r, d) in possible_transitions.items():
next_state_arr = np.array([pos, midpoints_visited], dtype=int)
results.append([next_state_arr, p, r, d])
return results
def get_new_position(self, action):
# 1. Simplify noise sampling (equivalent to your normal ppf logic)
p_dev = self.NOISE / 3.0
tendency = np.random.choice([-1, 0, 1], p=[p_dev, 1 - 2 * p_dev, p_dev])
# 2. Resolve direction shift if noise occurs
transition = Action(action)
if tendency != 0:
match transition:
case Action.Up:
transition = Action.Right if tendency == 1 else Action.Left
case Action.Down:
transition = Action.Left if tendency == 1 else Action.Right
case Action.Left:
transition = Action.Up if tendency == 1 else Action.Down
case Action.Right:
transition = Action.Down if tendency == 1 else Action.Up
# 3. Calculate grid position with boundary limits
curr_pos = self.state[0]
match transition:
case Action.Up:
return (curr_pos + self.GRID_N) if curr_pos < self.GRID_N * (self.GRID_N - 1) else curr_pos
case Action.Down:
return (curr_pos - self.GRID_N) if curr_pos >= self.GRID_N else curr_pos
case Action.Left:
return (curr_pos - 1) if curr_pos % self.GRID_N != 0 else curr_pos
case Action.Right:
return (curr_pos + 1) if (curr_pos + 1) % self.GRID_N != 0 else curr_pos
def render(self, ax=None, title=None, legend_outside=True):
# 1. Initialize the plot (or draw into a caller-provided axes for side-by-side comparisons)
standalone = ax is None
if standalone:
fig, ax = plt.subplots(figsize=(7, 7))
ax.set_xlim(-0.5, self.GRID_N - 0.5)
ax.set_ylim(-0.5, self.GRID_N - 0.5)
# Configure grid lines to match your matrix index boundary points
ax.set_xticks(np.arange(self.GRID_N))
ax.set_yticks(np.arange(self.GRID_N))
ax.grid(True, which='both', color='gainsboro', linestyle='-', linewidth=0.5)
# 2. Extract and Plot Pylons (|)
px = self.PYLON_GRID_COORDINATES[:, 0]
py = self.PYLON_GRID_COORDINATES[:, 1]
ax.scatter(px, py, color='firebrick', marker='|', s=250, linewidths=3, label='Pylon')
# 3. Extract and Plot Midpoints (o)
mx = self.PYLON_MIDPOINT_GRID_COORDINATES[:, 0]
my = self.PYLON_MIDPOINT_GRID_COORDINATES[:, 1]
ax.scatter(mx, my, color='darkorange', marker='o', s=100, facecolors='none', edgecolors='darkorange', linewidths=2, label='Midpoint')
# 4. Process and Plot Agent Path Line (With Smooth Continuous Gradient)
if self.state_history:
path_x = np.array([state[0] % self.GRID_N for state in self.state_history])
path_y = np.array([state[0] // self.GRID_N for state in self.state_history])
# --- Continuous Gradient Line Implementation via Interpolation ---
if len(path_x) > 1:
# Interpolate 100 fine points between every step to blend colors flawlessly
num_interp_points = (len(path_x) - 1) * 100
t_original = np.arange(len(path_x))
t_fine = np.linspace(0, len(path_x) - 1, num_interp_points)
fine_x = np.interp(t_fine, t_original, path_x)
fine_y = np.interp(t_fine, t_original, path_y)
# Reshape the upscaled points into micro-segments
points = np.array([fine_x, fine_y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# --- CONTROLLING THE COLOR SHIFT RATE ---
# Default was np.linspace(0, 1, len(segments))
# Adjust 'max_shift' to control how far along the colormap the path gets.
# 0.4 means it will only shift through 40% of the colormap over the entire path.
max_shift = 1
path_progression = np.linspace(0, max_shift, len(segments))
# ----------------------------------------
lc = LineCollection(segments, cmap='plasma', linewidths=3, alpha=0.9, zorder=2)
lc.set_array(path_progression)
# CRITICAL: Fix the colormap scaling limits so 1.0 remains the absolute ceiling
lc.set_clim(0, 1)
ax.add_collection(lc)
# Add a dummy handle so 'Agent Path' still shows up nicely in the legend
from matplotlib.lines import Line2D
legend_dummy = Line2D([0], [0], color=plt.get_cmap('plasma')(0.5), lw=3, label='Agent Path')
# We will handle adding this dummy to the legend safely below
# -----------------------------------------------------------------
# Overlay directional action markers at each step along the line
action_markers = {Action.Up.value: '^', Action.Down.value: 'v', Action.Left.value: '<', Action.Right.value: '>'}
for i, action in enumerate(self.action_history):
if action in action_markers:
ax.scatter(path_x[i], path_y[i], color='dodgerblue', marker=action_markers[action], s=80, zorder=3)
# Distinctly mark Start and Finish states
ax.scatter(path_x[0], path_y[0], color='forestgreen', marker='s', s=120, label='Start', zorder=4)
ax.scatter(path_x[-1], path_y[-1], color='gold', marker='*', s=250, edgecolor='black', label='Finish', zorder=4)
# 5. Labels and Legends
ax.set_title(title or "SimplePURT Grid Visualization", fontsize=14, fontweight='bold', pad=15)
ax.set_aspect('equal')
# Adjust legend rendering to cleanly include our gradient dummy line
handles, labels = ax.get_legend_handles_labels()
if self.state_history and len(path_x) > 1:
handles.append(legend_dummy)
labels.append('Agent Path')
if legend_outside:
ax.legend(handles, labels, loc='upper left', bbox_to_anchor=(1.02, 1), borderaxespad=0)
else:
# Compact in-axes legend so multiple subplots can sit side by side without overlapping
ax.legend(handles, labels, loc='upper right', fontsize=8, framealpha=0.9)
if standalone:
plt.tight_layout()
plt.show()
---------------------------------------------------------------------------
ModuleNotFoundError Traceback (most recent call last)
Cell In[1], line 3
1 # Simple PURT Environment
----> 3 import gymnasium as gym
4 from gymnasium import spaces
5 import numpy as np
ModuleNotFoundError: No module named 'gymnasium'
Function to run a full episode
def sample_action(policy, state):
action_probs = policy[*state] # Action probabilities (must sum to 1)
return np.random.choice(np.size(action_probs,0), p=action_probs) # Action sampling
def run_episode(env, policy, render_result=False):
while True:
action = sample_action(policy, env.state)
observation, reward, done, truncated, info = env.step(action)
if done or truncated:
break
if render_result:
env.render()
return env.get_info()
Value iteration to use as a baseline
# Value iteration function
def value_iteration(env, discount_factor, theta):
V = np.zeros(env.observation_space.nvec, dtype=np.float64) # Initialize state-value function
n_iter = 0
# 1. Find V* by sweeping the Bellman Optimality equation until it converges
while True:
n_iter += 1
delta = 0
for s in np.ndindex(*env.observation_space.nvec):
s = np.asarray(s)
old_v = V[*s]
Q = np.zeros(env.action_space.n, dtype=np.float64)
for a in range(env.action_space.n):
env.state = s # Set the environment state to the current state
for next_state, prob, reward, done in env.get_possible_next_states(a):
Q[a] += prob * (reward + discount_factor * V[*next_state] * (1 - done))
V[*s] = np.max(Q) # max instead of following a fixed policy
delta = max(delta, np.abs(old_v - V[*s]))
if delta < theta:
break
# 2. Extract the greedy policy from V* (with one final pass, no loop)
policy = np.random.randint(env.action_space.n, size=env.observation_space.nvec, dtype=np.int8)
for s in np.ndindex(*env.observation_space.nvec):
s = np.asarray(s)
Q = np.zeros(env.action_space.n, dtype=np.float64)
for a in range(env.action_space.n):
env.state = s
for next_state, prob, reward, done in env.get_possible_next_states(a):
Q[a] += prob * (reward + discount_factor * V[*next_state] * (1 - done))
policy[*s] = np.argmax(Q)
policy = np.eye(env.action_space.n)[policy]
return policy, V, n_iter
Monte Carlo Prediction#
First-Visit Monte Carlo Prediction#
from collections import defaultdict
def fv_mc_pred(env, policy, discount_factor, n_iter):
# Initialize V arbitrarily
V = np.random.random(env.observation_space.nvec)
# Use a dictionary mapping state to list of returns
returns = defaultdict(list)
for i in range(n_iter):
# Run episode
env.reset()
episode_info = run_episode(env, policy)
state_history = episode_info["state_history"]
reward_history = episode_info["reward_history"]
G = 0.0
n_states = len(state_history)
# Convert state history to tuples for indexing
states_as_tuples = [tuple(s) for s in state_history]
# Iterate backward through the episode
for t in reversed(range(n_states-1)):
G = discount_factor * G + reward_history[t]
state_t = states_as_tuples[t]
# Check if state appeared prior to time t
if state_t not in states_as_tuples[:t]:
returns[state_t].append(G)
V[*state_t] = np.mean(returns[state_t])
return V
from tqdm import tqdm
env1 = SimplePURT(NOISE=0.25)
discount_factor1 = 0.99
VI_policy, VI_V, n_iter = value_iteration(env1,discount_factor1,0.01)
nrange1 = range(1,101,1)
fv_mc_V1 = []
for n in tqdm(nrange1):
fv_mc_V1.append(fv_mc_pred(env1, VI_policy, discount_factor1, n)[1,4])
plt.plot(nrange1, fv_mc_V1)
plt.title("Monte Carlo Value of One-Before-Endpoint State for a Value Iteration Policy")
plt.xlabel("Iterations")
plt.ylabel("Value")
plt.grid(True)
100%|██████████| 100/100 [00:14<00:00, 6.68it/s]
Off-Policy Monte Carlo Prediction with Epsilon-Greedy Behavior Policy#
def off_mc_pred(env, target_policy, discount_factor, epsilon, n_iter):
# Initialize Q arbitrarily and C as zero
Q = np.random.random([*env.observation_space.nvec, env.action_space.n])
C = np.zeros([*env.observation_space.nvec, env.action_space.n], dtype=np.float64)
for i in range(n_iter):
# Epsilon-greedy behavior policy
behavior_policy = np.ones_like(target_policy) * epsilon / env.action_space.n + (1 - epsilon) * np.eye(env.action_space.n, dtype=int)[np.argmax(target_policy, axis=2)]
# Run episode
env.reset()
episode_info = run_episode(env, behavior_policy)
state_history = episode_info["state_history"]
action_history = episode_info["action_history"]
reward_history = episode_info["reward_history"]
G = 0.0
W = 1.0
n_states = len(state_history)
# Convert state history to tuples for indexing
states_as_tuples = [tuple(s) for s in state_history]
# Iterate backward through the episode
for t in reversed(range(n_states-1)):
G = discount_factor * G + reward_history[t]
state_t = states_as_tuples[t]
action_t = action_history[t]
C[*state_t, action_t] += W
Q[*state_t, action_t] += W / C[*state_t, action_t] * (G - Q[*state_t, action_t]) if C[*state_t, action_t] != 0 else 0
W *= target_policy[*state_t, action_t] / behavior_policy[*state_t, action_t] if behavior_policy[*state_t, action_t] != 0 else 1
return Q
off_mc_pred_Q1 = []
eps_range1 = np.arange(0,1,0.1)
for n in tqdm(nrange1):
off_mc_pred_Q1.append(off_mc_pred(env1, VI_policy, discount_factor1, 0.3, n)[1,4,2])
plt.plot(nrange1, off_mc_pred_Q1)
plt.title("Off-Policy MC Action-Value of Endgame for a Value Iteration Policy (epsilon=0.3)")
plt.xlabel("Iterations")
plt.ylabel("Action-Value")
plt.grid(True)
100%|██████████| 100/100 [00:11<00:00, 8.65it/s]
Monte Carlo Control#
Monte Carlo Exploring Starts#
def es_mc_control(env, discount_factor, n_iter):
# Initialize policy and Q arbitrarily
policy = np.ones([*env.observation_space.nvec, env.action_space.n]) / env.action_space.n
Q = np.random.random([*env.observation_space.nvec, env.action_space.n])
# Use a dictionary mapping state to list of returns
returns = defaultdict(list)
for i in range(n_iter):
# Exploring start episode
state_0 = env.observation_space.sample()
action_0 = env.action_space.sample()
env.reset()
env.state = state_0
env.state_history = [state_0]
env.step(action_0)
episode_info = run_episode(env, policy)
state_history = episode_info["state_history"]
action_history = episode_info["action_history"]
reward_history = episode_info["reward_history"]
G = 0.0
n_states = len(state_history)
# Convert state history to tuples for indexing
states_as_tuples = [tuple(s) for s in state_history]
visited_sa_pairs = [(states_as_tuples[t], action_history[t]) for t in range(n_states-1)]
# Iterate backward through the episode
for t in reversed(range(n_states-1)):
G = discount_factor * G + reward_history[t]
state_t = states_as_tuples[t]
action_t = action_history[t]
sa_key = state_t + (action_t,)
# Check if state + action appeared prior to time t
if (state_t, action_t) not in visited_sa_pairs[:t]:
returns[sa_key].append(G)
Q[*state_t, action_t] = np.mean(returns[sa_key])
# Greedy policy
policy[*state_t] = np.eye(env.action_space.n)[np.argmax(Q[*state_t])]
return policy, Q
es_mc_policy, es_mc_Q = es_mc_control(env1, discount_factor1, 10000)
env1.reset()
while len(env1.state_history) < 20:
env1.reset()
run_episode(env1, es_mc_policy)
env1.render(title="Exploring Starts MC Control after 10000 Iterations")
On-Policy First-Visit Monte Carlo Control with Epsilon-Greedy#
def on_fv_mc_control(env, discount_factor, epsilon, n_iter):
# Initialize policy arbitrarily as epsilon-soft and Q arbitrarily
policy = np.ones([*env.observation_space.nvec, env.action_space.n]) / env.action_space.n
Q = np.random.random([*env.observation_space.nvec, env.action_space.n])
# Use a dictionary mapping state to list of returns
returns = defaultdict(list)
for i in range(n_iter):
# Run episode
env.reset()
episode_info = run_episode(env, policy)
state_history = episode_info["state_history"]
action_history = episode_info["action_history"]
reward_history = episode_info["reward_history"]
G = 0.0
n_states = len(state_history)
# Convert state history to tuples for indexing
states_as_tuples = [tuple(s) for s in state_history]
visited_sa_pairs = [(states_as_tuples[t], action_history[t]) for t in range(n_states-1)]
# Iterate backward through the episode
for t in reversed(range(n_states-1)):
G = discount_factor * G + reward_history[t]
state_t = states_as_tuples[t]
action_t = action_history[t]
sa_key = state_t + (action_t,)
# Check if state + action appeared prior to time t
if (state_t, action_t) not in visited_sa_pairs[:t]:
returns[sa_key].append(G)
Q[*state_t, action_t] = np.mean(returns[sa_key])
# Epsilon-greedy policy
policy[*state_t] = np.ones(env.action_space.n) * epsilon / env.action_space.n + (1 - epsilon) * np.eye(env.action_space.n, dtype=int)[np.argmax(Q[*state_t])]
return policy, Q
on_fv_mc_policy, on_fv_mc_Q = on_fv_mc_control(env1, discount_factor1, 0.3, 10000)
env1.reset()
while len(env1.state_history) < 20:
env1.reset()
run_episode(env1, on_fv_mc_policy)
env1.render(title="On-Policy First-Visit MC Control after 10000 Iterations (epsilon=0.3)")
Off-Policy Monte Carlo Control with Epsilon-Greedy Behavior Policy#
def off_mc_control(env, discount_factor, epsilon, n_iter):
# Initialize Q arbitrarily and C as zero
Q = np.random.random([*env.observation_space.nvec, env.action_space.n])
C = np.zeros([*env.observation_space.nvec, env.action_space.n], dtype=np.float64)
target_policy = np.eye(env.action_space.n, dtype=int)[np.argmax(Q, axis=2)]
for i in range(n_iter):
# Epsilon-greedy behavior policy
behavior_policy = np.ones_like(target_policy) * epsilon / env.action_space.n + (1 - epsilon) * np.eye(env.action_space.n, dtype=int)[np.argmax(target_policy, axis=2)]
# Run episode
env.reset()
episode_info = run_episode(env, behavior_policy)
state_history = episode_info["state_history"]
action_history = episode_info["action_history"]
reward_history = episode_info["reward_history"]
G = 0.0
W = 1.0
n_states = len(state_history)
# Convert state history to tuples for indexing
states_as_tuples = [tuple(s) for s in state_history]
# Iterate backward through the episode
for t in reversed(range(n_states-1)):
G = discount_factor * G + reward_history[t]
state_t = states_as_tuples[t]
action_t = action_history[t]
C[*state_t, action_t] += W
Q[*state_t, action_t] += W / C[*state_t, action_t] * (G - Q[*state_t, action_t]) if C[*state_t, action_t] != 0 else 0
best_action = np.argmax(Q[*state_t])
target_policy[*state_t] = np.eye(env.action_space.n)[best_action]
# Check if the action is different
if action_t != best_action:
break
W *= 1 / behavior_policy[*state_t, action_t] if behavior_policy[*state_t, action_t] != 0 else 1
return target_policy, Q
off_mc_policy, off_mc_Q = off_mc_control(env1, discount_factor1, 0.3, 10000)
env1.reset()
while len(env1.state_history) < 20:
env1.reset()
run_episode(env1, off_mc_policy)
env1.render(title="Off-Policy MC Control after 10000 Iterations (epsilon=0.3)")
