WindyGridworld.py

# -*- coding: utf-8 -*-
"""Untitled6.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/16HOegynM_qxasZhVWHxBUBo4CmyTHFI-
"""

import numpy as np
import matplotlib.pyplot as plt

class WindyGridworld:
    def __init__(self):
        self.num_rows = 7
        self.num_cols = 10
        self.start_state = (3, 0)
        self.goal_state = (3, 7)
        self.current_state = self.start_state
        self.wind = [0, 0, 0, 1, 1, 1, 2, 2, 1, 0]
        
    def reset(self):
        self.current_state = self.start_state
        return self.current_state
    
    def step(self, action):
        # Update the state based on the chosen action
        if action == 0:  # Up
            next_state = (self.current_state[0] - 1, self.current_state[1])
        elif action == 1:  # Down
            next_state = (self.current_state[0] + 1, self.current_state[1])
        elif action == 2:  # Left
            next_state = (self.current_state[0], self.current_state[1] - 1)
        elif action == 3:  # Right
            next_state = (self.current_state[0], self.current_state[1] + 1)
        
        # Apply the wind effect
        next_state = (next_state[0] - self.wind[self.current_state[1]], next_state[1])
        
        # Ensure the next state is within the gridworld boundaries
        next_state = (max(0, min(next_state[0], self.num_rows - 1)), max(0, min(next_state[1], self.num_cols - 1)))
        
        # Calculate the reward based on the next state
        if next_state == self.goal_state:
            reward = 0
        else:
            reward = -1
        
        # Update the current state
        self.current_state = next_state
        
        # Determine if the episode is done
        done = self.current_state == self.goal_state
        
        return self.current_state, reward, done, {}

def epsilon_greedy_policy(Q, state, epsilon):
    if np.random.rand() < epsilon:
        return np.random.randint(Q.shape[1])
    else:
        return np.argmax(Q[state])

def n_step_sarsa(env, n, alpha, gamma, epsilon, num_episodes):
    num_states = env.num_rows * env.num_cols
    num_actions = 4
    Q = np.zeros((num_states, num_actions))
    counter = 0
    episode_rewards = np.zeros(num_episodes)
    episode_lengths = np.zeros(num_episodes)
    
    for episode in range(num_episodes):
        state = env.reset()
        action = epsilon_greedy_policy(Q, state, epsilon)
        T = np.inf
        t = 0
        tau = 0
        rewards = []
        states = [state]
        actions = [action]
        while tau < (T - 1):
            if t < T:
                next_state, reward, done, _ = env.step(action)
                counter += 1
                # print(next_state)
                rewards.append(reward)
                states.append(next_state)
                
                if done:
                    T = t + 1
                else:
                    next_action = epsilon_greedy_policy(Q, next_state[0]*next_state[1], epsilon)
                    actions.append(next_action)
            
            tau = t - n + 1
            
            if tau >= 0:
                G = np.sum([gamma**(i - tau - 1) * rewards[i] for i in range(tau + 1, min(tau + n, T))])
                
                if tau + n < T:
                    G += gamma**n * Q[states[tau + n][0] * env.num_cols + states[tau + n][1], actions[tau + n]]
                
                Q[states[tau][0] * env.num_cols + states[tau][1], actions[tau]] += alpha * (G - Q[states[tau][0] * env.num_cols + states[tau][1], actions[tau]])
            
            t += 1
            state = next_state
            action = next_action
        # if flag==1:
        #     break
        episode_rewards[episode] = np.sum(rewards)
        episode_lengths[episode] = len(rewards)
    
    return Q, episode_rewards, episode_lengths

def lambda_sarsa(env, n, alpha, gamma, epsilon, num_episodes, lmbda):
    num_states = env.num_rows * env.num_cols
    num_actions = 4
    Q = np.zeros((num_states, num_actions))
    counter=0
    episode_rewards = np.zeros(num_episodes)
    episode_lengths = np.zeros(num_episodes)
    
    for episode in range(num_episodes):
        state = env.reset()
        action = epsilon_greedy_policy(Q, state[0] * env.num_cols + state[1], epsilon)
        
        E = np.zeros((num_states, num_actions))
        rewards = []
        states = [state]
        actions = [action]
        
        done = False
        
        while not done:
            next_state, reward, done, _ = env.step(action)
            counter+=1
            # print(next_state)
            rewards.append(reward)
            states.append(next_state)
            
            next_action = epsilon_greedy_policy(Q, next_state[0] * env.num_cols + next_state[1], epsilon)
            actions.append(next_action)
            
            delta = reward + gamma * Q[next_state[0] * env.num_cols + next_state[1], next_action] - Q[state[0] * env.num_cols + state[1], action]
            
            E[state[0] * env.num_cols + state[1], action] += 1
            
            for s in range(num_states):
                for a in range(num_actions):
                    Q[s, a] += alpha * delta * E[s, a]
                    E[s, a] *= gamma * lmbda
            
            state = next_state
            action = next_action
        # print(counter)
        episode_rewards[episode] = np.sum(rewards)
        episode_lengths[episode] = len(rewards)
    
    return Q, episode_rewards, episode_lengths
#%%
# Create the WindyGridworld environment
env = WindyGridworld()

# Set the algorithm parameters
alpha = 0.8  # Learning rate
gamma = 1.0  # Discount factor
epsilon = 0.1  # Exploration rate
num_episodes = 200
#%%
"""## N-Step Sarsa (n=1)"""

# Set the parameters for n-step SARSA
n = 1

# Run n-step SARSA algorithm
Q_n_step, rewards_n_step, lengths_n_step = n_step_sarsa(env, n, alpha, gamma, epsilon, num_episodes)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_n_step)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (n-step SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_n_step)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (n-step SARSA)')

plt.tight_layout()
plt.show()
#%%
"""## N-Step Sarsa (n=5)"""

# Set the parameters for n-step SARSA
n = 5

# Run n-step SARSA algorithm
Q_n_step, rewards_n_step, lengths_n_step = n_step_sarsa(env, n, alpha, gamma, epsilon, num_episodes)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_n_step)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (n-step SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_n_step)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (n-step SARSA)')

plt.tight_layout()
plt.show()
#%%
"""## N-Step Sarsa (n=10)"""

# Set the parameters for n-step SARSA
n = 10

# Run n-step SARSA algorithm
Q_n_step, rewards_n_step, lengths_n_step = n_step_sarsa(env, n, alpha, gamma, epsilon, num_episodes)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_n_step)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (n-step SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_n_step)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (n-step SARSA)')

plt.tight_layout()
plt.show()
#%%
"""## λ-Step Sarsa (λ=0)"""

# Set the parameters for λ-SARSA
lmbda = 0

# Run λ-SARSA algorithm
Q_lambda, rewards_lambda, lengths_lambda = lambda_sarsa(env, n, alpha, gamma, epsilon, num_episodes, lmbda)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_lambda)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (λ-SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_lambda)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (λ-SARSA)')

plt.tight_layout()
plt.show()
#%%
"""## λ-Step Sarsa (λ=0.1)"""

# Set the parameters for λ-SARSA
lmbda = 0.1

# Run λ-SARSA algorithm
Q_lambda, rewards_lambda, lengths_lambda = lambda_sarsa(env, n, alpha, gamma, epsilon, num_episodes, lmbda)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_lambda)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (λ-SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_lambda)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (λ-SARSA)')

plt.tight_layout()
plt.show()
#%%
"""## λ-Step Sarsa (λ=0.5)"""

# Set the parameters for λ-SARSA
lmbda = 0.5

# Run λ-SARSA algorithm
Q_lambda, rewards_lambda, lengths_lambda = lambda_sarsa(env, n, alpha, gamma, epsilon, num_episodes, lmbda)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_lambda)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (λ-SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_lambda)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (λ-SARSA)')

plt.tight_layout()
plt.show()
#%%
"""## λ-Step Sarsa (λ=1)"""

# Set the parameters for λ-SARSA
lmbda = 1

# Run λ-SARSA algorithm
Q_lambda, rewards_lambda, lengths_lambda = lambda_sarsa(env, n, alpha, gamma, epsilon, num_episodes, lmbda)

# Plot the learning curve
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(rewards_lambda)
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('Episode Rewards (λ-SARSA)')

plt.subplot(1, 2, 2)
plt.plot(lengths_lambda)
plt.xlabel('Episode')
plt.ylabel('Episode Length')
plt.title('Episode Lengths (λ-SARSA)')

plt.tight_layout()
plt.show()