Problem with train tensor in deep QLearning code

Hello everyone Im looking for some help with an issue i have been working for a few days now that im unable to solve.

I have this code that does Deep Q Learning imitating frozenLake enviroment from GYM but without using the gym library, the original enviroment return as state a integer form 0 to 16 that latter need to be one hot encoded as a flat 16 size tensor like 7 = [0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0] for the Neural netwotk to evaluate and train with.

The code works excellent but for a new project i need to pass the whole matrix (board) instead this one hot encoded integer, so i simply made state.flatten() which retus a flat 16 size tensor with a 1 on the correct location of the character exactly as the one hot encoded integer does.

The problem is that with the integer encoded works like charm but with the flatten() state the trainig doesent work right, it trains but rewards became a mess, so i belive that problem must be some kind of deep thing about tensors that i dont know yet.

That why im looking for some advice about what can be the issue with this code.

The code is quite simple a constant named TRAINMODE = 0 or TRAINMODE = 1 select which kind of training must be done just change the get_state() fuction from class “MazeEnvironment” from returning the onehotencoded integer or returning state.flatten()

I have tried to compare both returns from the class and seems to be equal in type, size, data, etc. so seems the problem is not from the class but other part of the training.

I really really appreciate any help or advice.

thank you so much!

import numpy as np
import matplotlib.pyplot as plt
from collections import deque
import random
import torch
from torch import nn
import torch.nn.functional as F

# Definir el entorno del laberinto
class MazeEnvironment:
    def __init__(self, modo, size):
        self.modo = modo
        self.size = size
        self.reset()

    def reset(self):
        self.state = torch.zeros(self.size, device=device)
        self.stepCount = 0
        self.reward = 0
        self.truncated = False
        self.done =  False
        self.player_position = (0, 0)
        self.goal_position = (self.size[0] - 1, self.size[1] - 1)
        self.trap_positions = [(1, 1), (1, 3),(2, 3), (3, 0)]  # Ejemplo de posiciones de trampas
        self.state[self.player_position] = 1
        return self.get_state()
    
    def action_space_sample(self):
        return random.randint(0, 3)
    
    def get_state_int(self):
        return int(self.player_position[0] * 4 + self.player_position[1])
    
    def get_state(self):
        if self.modo == 0:
            return self.get_state_flat()
        else:
            return self.get_state_dqn()
    
    def get_state_flat(self)->torch.Tensor:
        return self.state.flatten()
    
    def get_state_dqn(self)->torch.Tensor:
        input_tensor = torch.zeros(16, device=device)
        input_tensor[self.get_state_int()] = 1
        return input_tensor
    
    def step(self, action):
        new_position = self.player_position

        if not self.done:

            self.state[self.player_position] = 0

            # for printing 0,1,2,3 => L(eft),D(own),R(ight),U(p)

            if action == 3:  # Arriba
                new_position = (self.player_position[0] - 1, self.player_position[1])
            elif action == 1:  # Abajo
                new_position = (self.player_position[0] + 1, self.player_position[1])
            elif action == 0:  # Izquierda
                new_position = (self.player_position[0], self.player_position[1] - 1)
            elif action == 2:  # Derecha
                new_position = (self.player_position[0], self.player_position[1] + 1)


            # Verificar si la nueva posición es válida
            if new_position[0] < self.size[0] and new_position[0] >=0 and new_position[1] < self.size[1] and new_position[1] >= 0:
                self.player_position = new_position

            self.state[self.player_position] = 1

            if self.player_position == self.goal_position :
                self.done = True
                self.reward = 1
            
            for position in self.trap_positions:
                if self.player_position == position:
                    self.done = True
                    self.reward = 0
                    break

        self.stepCount+=1

        
        if self.stepCount >= 100:
            self.truncated = True

        return self.get_state(), self.reward, self.done, self.truncated

# Define model
class DQN(nn.Module):
    def __init__(self, in_states, h1_nodes, out_actions):
        super().__init__()

        # Define network layers
        self.fc1 = nn.Linear(in_states, h1_nodes)   # first fully connected layer
        self.out = nn.Linear(h1_nodes, out_actions) # ouptut layer w

    def forward(self, x):
        x = F.relu(self.fc1(x)) # Apply rectified linear unit (ReLU) activation
        x = self.out(x)         # Calculate output
        return x

# Define memory for Experience Replay
class ReplayMemory():
    def __init__(self, maxlen):
        self.memory = deque([], maxlen=maxlen)
    
    def append(self, transition):
        self.memory.append(transition)

    def sample(self, sample_size):
        return random.sample(self.memory, sample_size)

    def __len__(self):
        return len(self.memory)

# FrozeLake Deep Q-Learning
class FrozenLakeDQL():
    # Hyperparameters (adjustable)
    learning_rate_a = 0.001         # learning rate (alpha)
    discount_factor_g = 0.9         # discount rate (gamma)    
    network_sync_rate = 10          # number of steps the agent takes before syncing the policy and target network
    replay_memory_size = 1000       # size of replay memory
    mini_batch_size = 32            # size of the training data set sampled from the replay memory

    # Neural Network
    loss_fn = nn.MSELoss()          # NN Loss function. MSE=Mean Squared Error can be swapped to something else.
    optimizer = None                # NN Optimizer. Initialize later.

    ACTIONS = ['L','D','R','U']     # for printing 0,1,2,3 => L(eft),D(own),R(ight),U(p)

    # Train the FrozeLake environment
    def train(self, episodes):
        # Create FrozenLake instance
        env = MazeEnvironment(modo=TRAINMODE, size=(4, 4))
        num_states = 16
        num_actions = 4
        
        epsilon = 1 # 1 = 100% random actions
        memory = ReplayMemory(self.replay_memory_size)

        # Create policy and target network. Number of nodes in the hidden layer can be adjusted.
        policy_dqn = DQN(in_states=num_states, h1_nodes=num_states, out_actions=num_actions).to(device)
        target_dqn = DQN(in_states=num_states, h1_nodes=num_states, out_actions=num_actions).to(device)

        # Make the target and policy networks the same (copy weights/biases from one network to the other)
        target_dqn.load_state_dict(policy_dqn.state_dict())

        print('Policy (random, before training):')

        # Policy network optimizer. "Adam" optimizer can be swapped to something else. 
        self.optimizer = torch.optim.Adam(policy_dqn.parameters(), lr=self.learning_rate_a)

        # List to keep track of rewards collected per episode. Initialize list to 0's.
        rewards_per_episode = np.zeros(episodes)

        # List to keep track of epsilon decay
        epsilon_history = []

        # Track number of steps taken. Used for syncing policy => target network.
        step_count=0
            
        for i in range(episodes):
            state = env.reset()     # Initialize to state 0
            terminated = False      # True when agent falls in hole or reached goal
            truncated = False       # True when agent takes more than 200 actions    

            # Agent navigates map until it falls into hole/reaches goal (terminated), or has taken 200 actions (truncated).
            while(not terminated and not truncated):

                # Select action based on epsilon-greedy
                if random.random() < epsilon:
                    # select random action
                    action = env.action_space_sample() #actions: 0=left,1=down,2=right,3=up
                else:
                    # select best action            
                    with torch.no_grad():
                        action = policy_dqn(state).argmax().item()

                # Execute action
                new_state,reward,terminated,truncated = env.step(action)

                # Save experience into memory
                memory.append((state, action, new_state, reward, terminated)) 

                # Move to the next state
                state = new_state

                # Increment step counter
                step_count+=1

            # Keep track of the rewards collected per episode.
            if reward == 1:
                rewards_per_episode[i] = 1

            # Check if enough experience has been collected and if at least 1 reward has been collected
            if len(memory)>self.mini_batch_size and np.sum(rewards_per_episode)>0:
                mini_batch = memory.sample(self.mini_batch_size)
                self.optimize(mini_batch, policy_dqn, target_dqn)        

                # Decay epsilon
                epsilon = max(epsilon - 1/episodes, 0)
                epsilon_history.append(epsilon)

                # Copy policy network to target network after a certain number of steps
                if step_count > self.network_sync_rate:
                    target_dqn.load_state_dict(policy_dqn.state_dict())
                    step_count=0

        # Save policy
        torch.save(policy_dqn.state_dict(), "frozen_lake_dql.pt")

        # Create new graph 
        plt.figure(1)

        # Plot average rewards (Y-axis) vs episodes (X-axis)
        sum_rewards = np.zeros(episodes)
        for x in range(episodes):
            sum_rewards[x] = np.sum(rewards_per_episode[max(0, x-100):(x+1)])
        plt.subplot(121) # plot on a 1 row x 2 col grid, at cell 1
        plt.plot(sum_rewards)
        
        # Plot epsilon decay (Y-axis) vs episodes (X-axis)
        plt.subplot(122) # plot on a 1 row x 2 col grid, at cell 2
        plt.plot(epsilon_history)
        
        # Save plots
        plt.savefig('frozen_lake_dql.png')

    # Optimize policy network
    def optimize(self, mini_batch, policy_dqn, target_dqn):

        current_q_list = []
        target_q_list = []

        for state, action, new_state, reward, terminated in mini_batch:

            if terminated: 
                # Agent either reached goal (reward=1) or fell into hole (reward=0)
                # When in a terminated state, target q value should be set to the reward.
                target = torch.tensor([reward], dtype=torch.float32, device=device)
            else:
                # Calculate target q value 
                with torch.no_grad():
                    target = torch.tensor(
                        [reward + self.discount_factor_g * target_dqn(new_state).max()],
                        dtype=torch.float32,
                        device=device
                    )

            # Get the current set of Q values
            current_q = policy_dqn(state)
            current_q_list.append(current_q)

            # Get the target set of Q values
            target_q = target_dqn(state) 
            # Adjust the specific action to the target that was just calculated
            target_q[action] = target
            target_q_list.append(target_q)
                
        # Compute loss for the whole minibatch
        loss = self.loss_fn(torch.stack(current_q_list), torch.stack(target_q_list))

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    # Run the FrozeLake environment with the learned policy
    def test(self, episodes):
        # Create FrozenLake instance
        env = MazeEnvironment(modo=TRAINMODE,size=(4, 4))
        num_states = 16
        num_actions = 4

        # Load learned policy
        policy_dqn = DQN(in_states=num_states, h1_nodes=num_states, out_actions=num_actions).to(device) 
        policy_dqn.load_state_dict(torch.load("frozen_lake_dql.pt"))
        policy_dqn.eval()    # switch model to evaluation mode

        print('Policy (trained):')
        #self.print_dqn(policy_dqn)

        for i in range(episodes):
            state = env.reset()     # Initialize to state 0
            terminated = False      # True when agent falls in hole or reached goal
            truncated = False       # True when agent takes more than 200 actions            

            # Agent navigates map until it falls into a hole (terminated), reaches goal (terminated), or has taken 200 actions (truncated).
            while(not terminated and not truncated):  
                # Select best action   
                with torch.no_grad():
                    action = policy_dqn(state).argmax().item()

                # Execute action
                state,reward,terminated,truncated = env.step(action)

if __name__ == '__main__':
    # Check if GPU is available
    torch.manual_seed(1234)
    random.seed(1234)
    torch.cuda.manual_seed_all(1234)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using {device} for training.")

    TRAINMODE = 0
    frozen_lake = FrozenLakeDQL()
    frozen_lake.train(1000)
    frozen_lake.test(1)