I am implementing a simple REINFORCE (policy gradient) algorithm for openAI's FrozenLake-v0 environment. However, it does not seem to learn anything at all.

I have used the same neural architecture for openAI's CartPole-v0, and trained it using REINFORCE (policy gradient), and it works perfectly. So, what I am doing incorrectly for the FrozenLake-v0 environment? I think this has to do with the nature of the environment, but I am unsure which aspects of training REINFORCE must be altered to accommodate the dynamics of FrozenLake-v0. It seems like a very simple environment to solve, given that it has only 16 states.

My code is as follows:

import gym
from gym.envs.registration import register
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
import matplotlib.pyplot as plt

# helper function for conversion of a state into an input to a neural network
def OH(x, n):
    :param x: state id
    :param n: n_states
    :return:  1-hot encoded numpy array of size [1,n]
    one_hot = np.zeros((n,))
    one_hot[x] = 1
    return one_hot

def running_mean(x, n):
    kernel = np.ones(N)
    conv_len = x.shape[0]-N
    y = np.zeros(conv_len)
    for i in range(conv_len):
        y[i] = kernel @ x[i:i+N]
        y[i] /= N
    return y

# architecture of the Policy Network
class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, n_actions):
        self.n_actions = n_actions
        self.model = nn.Sequential(
            nn.Linear(state_dim, 256),
            nn.Linear(256, n_actions),

    def forward(self, X):
        return self.model(X)

def train_reinforce_agent(env, episode_length = 100, max_episodes = 50000, gamma = 0.99, visualize_step = 50, learning_rate=0.003):

    # define the parametric model for the Policy: this is an instantiation of the PolicyNetwork class
    model = PolicyNetwork(env.observation_space.shape[0], env.action_space.n)
    # define the optimizer for updating the weights of the Policy Network
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # hyperparameters of the reinforce agent
    EPISODE_LENGTH = episode_length
    MAX_EPISODES = max_episodes
    GAMMA = gamma
    VISUALIZE_STEP = max(1, visualize_step)
    score = []

    for episode in range(MAX_EPISODES):
        # reset the environment
        curr_state = env.reset()
        done = False
        transitions = []

        # rollout an entire episode from the Policy Network
        for t in range(EPISODE_LENGTH):
            act_prob = model(torch.from_numpy(curr_state).float())
            action = np.random.choice(np.array(list(range(env.action_space.n))), p=act_prob.data.numpy())
            prev_state = curr_state
            curr_state, _, done, info = env.step(action)
            transitions.append((prev_state, action, t+1))

            if done:
        reward_batch = torch.Tensor([r for (s, a, r) in transitions]).flip(dims=(0,))

        # compute the return for every state-action pair from the rewards at every time-step
        batch_Gvals = []
        for i in range(len(transitions)):
            new_Gval = 0
            power = 0
            for j in range(i, len(transitions)):
                new_Gval = new_Gval + ((GAMMA ** power) * reward_batch[j]).numpy()
            power += 1

        # normalize the returns for the batch
        expected_returns_batch = torch.FloatTensor(batch_Gvals)
        expected_returns_batch /= expected_returns_batch.max()

        # batch the states, actions, prob after the episode
        state_batch = torch.Tensor([s for (s, a, r) in transitions])
        action_batch = torch.Tensor([a for (s, a, r) in transitions])
        pred_batch = model(state_batch)
        prob_batch = pred_batch.gather(dim=1, index=action_batch.long().view(-1, 1)).squeeze()

        # compute the loss for one episode
        loss = -torch.sum(torch.log(prob_batch) * expected_returns_batch)

        # back-propagate the loss
        # update the parameters of the Policy Network

        # print the status after every VISUALIZE_STEP episodes
        if episode % VISUALIZE_STEP == 0 and episode > 0:
            print('Episode {}\tAverage Score: {:.2f}'.format(episode, np.mean(score[-VISUALIZE_STEP:-1])))

    # Training plot: Episodic reward over Training Episodes
    score = np.array(score)
    avg_score = running_mean(score, visualize_step)
    plt.figure(figsize=(15, 7))
    plt.ylabel("Episode Duration", fontsize=12)
    plt.xlabel("Training Episodes", fontsize=12)
    plt.plot(score, color='gray', linewidth=1)
    plt.plot(avg_score, color='blue', linewidth=3)
    plt.scatter(np.arange(score.shape[0]), score, color='green', linewidth=0.3)
  • $\begingroup$ Hello. Welcome to Artificial Intelligence Stack Exchange. It seems to me that this is or could be a programming issue. Let me know if I am wrong or right. Generally, questions about programming problems are off-topic here and should probably be asked on Stack Overflow. Here, we focus on the theoretical, philosophical, social and historical aspects of Artificial Intelligence. Please, take the time to read our on-topic page. So, if you think this is not a programming problem, you should edit your post to clarify that. $\endgroup$
    – nbro
    Nov 26, 2021 at 11:53
  • $\begingroup$ Thanks for the pointer! However, I suspect this is more than a 'programming issue'. I think the reason the environment does not converge using REINFORCE, has to do with the nature of the environment. I was looking for suggestions on which specific aspects of the environment are an impediment to policy learning. I have edited my question to clarify this :) $\endgroup$
    – 204
    Nov 27, 2021 at 1:15


You must log in to answer this question.

Browse other questions tagged .