# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
import random

# +
env = gym.make('CartPole-v1')

num_actions = env.action_space.n
state_space_dimensions = env.observation_space.shape[0]

print('Actions: {} -- State space dimensions: {}'.format(num_actions, state_space_dimensions))


# -

class CartPoleDQN(nn.Module):
    
    def __init__(self):
        super(CartPoleDQN, self).__init__()
        self.fc1 = nn.Linear(state_space_dimensions, 256)
        self.fc2 = nn.Linear(256, 64)
        #self.fc3 = nn.Linear(128, 64)
        self.head = nn.Linear(64, num_actions)
    
    def forward(self, x):
        x = F.leaky_relu(self.fc1(x))
        x = F.leaky_relu(self.fc2(x))
        #x = F.leaky_relu(self.fc3(x))
        x = F.leaky_relu(self.head(x))
        return x


class ReplayBuffer():
    
    def __init__(self, capacity):
        self.capacity = capacity
        self.position = 0
        self.states = []
        self.actions = []
        self.rewards = []
        self.next_states = []
        self.done = []
        
        
    def push(self, state, action, reward, next_state, done):
        if len(self.states) < self.capacity:
            self.states.append(None)
            self.actions.append(None)
            self.rewards.append(None)
            self.next_states.append(None)
            self.done.append(None)
            
        self.states[self.position] = state
        self.actions[self.position] = action
        self.rewards[self.position] = reward
        self.next_states[self.position] = next_state
        self.done[self.position] = done
        
        self.position = (self.position + 1) % self.capacity
        
        
    def sample(self, batch_size):
        indices = np.random.choice(range(len(self.states)), size=batch_size)
        state_sample = [self.states[i] for i in indices]
        action_sample = [self.actions[i] for i in indices]
        reward_sample = [self.rewards[i] for i in indices]
        next_state_sample = [self.next_states[i] for i in indices]
        done_sample = [self.done[i] for i in indices]
        
        return state_sample, action_sample, reward_sample, next_state_sample, done_sample
        
    def __len__(self):
        return len(self.states)


def select_action(state, env, model, epsilon):
    if random.random() > epsilon:
        with torch.no_grad():
            return model(state).argmax().item()
    else:
        return env.action_space.sample()
    


def update_epsilon(epsilon_start, epsilon_end, epsilon_steps, total_steps):
    return epsilon_end + (epsilon_start - epsilon_end) * math.exp(-1. * total_steps / epsilon_steps)  
    


# +
def optimize_model(policy_net, target_net, optimizer, memory, batch_size, gamma):
    state_batch, action_batch, reward_batch, next_state_batch, done_batch = memory.sample(batch_size)
        
    state_batch = torch.tensor(state_batch).float().view((batch_size, -1))
    action_batch = torch.tensor(action_batch, dtype=torch.int64).view((batch_size))
    reward_batch = torch.tensor(reward_batch).float().view((batch_size))
    
    non_final_next_states = torch.tensor([s for s in next_state_batch if s is not None]).float()
    non_final_mask = torch.tensor(list(map(lambda s: s is not None, next_state_batch)), dtype=torch.bool)
        
    state_action_values = policy_net(state_batch)
    state_action_values = state_action_values.gather(1, action_batch.reshape((batch_size, 1)))
    
    next_state_values = torch.zeros(batch_size)
    next_state_values[non_final_mask] = target_net(non_final_next_states).max(dim=1)[0].float().detach()
        
    expected_state_action_values = reward_batch + gamma * next_state_values
    
    loss = F.smooth_l1_loss(state_action_values, expected_state_action_values.unsqueeze(1))
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    return loss
    
    

def train_dqn(env, policy_net, target_net, optimizer, memory, target_update=10, batch_size=32, episodes=100, gamma=0.99, epsilon_start=0.9, epsilon_end=0.05, epsilon_steps=1000000):
    
    total_rewards = []
    total_steps = 0
    
    epsilon = epsilon_start
        
    for episode in range(episodes):
        
        done = False
        state = env.reset()
        
        total_rewards.append(0)
        loss = 0
        
        while not done:
            
            env.render()
            state_tensor = torch.tensor(state).float()
            action = select_action(state_tensor, env, policy_net, epsilon)
            
            next_state, reward, done, _ = env.step(action)
            
            if done:
                next_state = None
            
            total_rewards[episode] += reward
            
            
            memory.push(state, action, reward, next_state, done)
            
            
            state = next_state
            
            if len(memory) >= batch_size:                 
                loss = optimize_model(policy_net, target_net, optimizer, memory, batch_size, gamma)
        
            if total_steps % target_update == 0:
                target_net.load_state_dict(policy_net.state_dict())
                
            total_steps += 1
            epsilon = update_epsilon(epsilon_start, epsilon_end, epsilon_steps, total_steps)
            
        
        print('{}/{} Total steps: {} Episode reward: {} Average reward: {} Loss: {} Epsilon: {}'.format(episode, episodes, total_steps, total_rewards[episode], np.mean(total_rewards), loss, epsilon))

            
                   

# +
env = gym.make('CartPole-v1')

num_actions = env.action_space.n
state_space_dimensions = env.observation_space.shape[0]

target_net = CartPoleDQN()
policy_net = CartPoleDQN()

#policy_net.load_state_dict(target_net.state_dict())

optimizer = torch.optim.Adam(policy_net.parameters())

memory = ReplayBuffer(100000)
try:
    train_dqn(env, policy_net, target_net, optimizer, memory, gamma=0.5, batch_size=128, episodes=1000, epsilon_steps=10000, epsilon_end=0.1)
finally:
    env.close()
# -



StarcoderdataJupyterScriptsDedupFiltered