Pytorch deepQL code that fail when using tensor.flatten() instead a One Hot Encoding function

17 Views Asked by At

Im working on a DeepQL code but i have found a problem I have been unable to solve, the short Story is that I have a code that works 100% reading from the enviroment class an integer and passing trought a One Hot Encoding function to convert the state 0-15 to a float tensor [0,0,0,1,....,0] that is passed to the Network that works fine, but I want to send the state (A float tensor 4x4) that only contains a 1 where the player is so if I use flatten() the tensor becames exactly as the OneHotfunction returns.

I have tested that the tensor that both approaches return, and they are the same (Or it seems).

For some reason if I pass the flattened tensor instead the hotfix one to exactly the same DeepQL code, works but the training became a mess rewards falls abruptly even when supposedly they are training with the same data.

My theory is that somehow the tensor is being modify or that even when seems identical both tensors are not, but i dont know how to test then because:

torch.all(a.eq(b)): tell they are equal a.type(), b.type() print that they are both float

Every other line of code is exactly the same on both tests the only thing that change is policy_dqn(state.flatten()).argmax().item() vs policy_dqn(self.state_to_dqn_input(state, num_states)).argmax().item()

which supposedly should be exactly the same but i dont know if there is something more deep about Tensors that I may have not taking into account, I would like to know if someone Would mind to check the code or if someone have an idea of what could be happening.

Original Code working 100%:

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, size):
        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(self):
        return int(self.player_position[0] * 4 + self.player_position[1])
    
    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(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):')
        #self.print_dqn(policy_dqn)

        # 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(self.state_to_dqn_input(state, num_states)).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):

        # Get number of input nodes
        num_states = policy_dqn.fc1.in_features

        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(self.state_to_dqn_input(new_state, num_states)).max()],
                        dtype=torch.float32,
                        device=device
                    )

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

            # Get the target set of Q values
            target_q = target_dqn(self.state_to_dqn_input(state, num_states)) 
            # 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()

    '''
    Converts an state (int) to a tensor representation.
    For example, the FrozenLake 4x4 map has 4x4=16 states numbered from 0 to 15. 

    Parameters: state=1, num_states=16
    Return: tensor([0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    '''
    def state_to_dqn_input(self, state:int, num_states:int)->torch.Tensor:
        input_tensor = torch.zeros(num_states, device=device)
        input_tensor[state] = 1
        return input_tensor

    # Run the FrozeLake environment with the learned policy
    def test(self, episodes):
        # Create FrozenLake instance
        env = MazeEnvironment(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(self.state_to_dqn_input(state, num_states)).argmax().item()

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

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

    frozen_lake = FrozenLakeDQL()
    frozen_lake.train(1000)
    frozen_lake.test(10)

Modified code not working:

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, size):
        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(self):
        return self.state #int(self.player_position[0] * 4 + self.player_position[1])
    
    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(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):')
        #self.print_dqn(policy_dqn)

        # 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.flatten()).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):

        # Get number of input nodes
        num_states = policy_dqn.fc1.in_features

        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.flatten()).max()],
                        dtype=torch.float32,
                        device=device
                    )

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

            # Get the target set of Q values
            target_q = target_dqn(state.flatten()) 
            # 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(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.flatten()).argmax().item()

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

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

    frozen_lake = FrozenLakeDQL()

    frozen_lake.train(3000)
    frozen_lake.test(10)

0

There are 0 best solutions below