policy_gradient.py

import torch as T
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

class PolicyGradient(nn.Module):
    def __init__(self, alpha, input_dims, fc1_dims, fc2_dims, n_actions):
        super(PolicyGradient, self).__init__()
        self.input_dims = input_dims
        self.fc1_dims = fc1_dims
        self.fc2_dims = fc2_dims
        self.n_actions = n_actions
        self.fc1 = nn.Linear(*self.input_dims, self.fc1_dims)
        self.fc2 = nn.Linear(self.fc1_dims, self.fc2_dims)
        self.pi = nn.Linear(self.fc2_dims, n_actions)
        self.v = nn.Linear(self.fc2_dims, 1)
        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
        
        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
        self.to(self.device)
        
        
    def forward(self, observation):
        state = T.Tensor(observation).to(self.device)
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        pi = self.pi(x)
        v = self.v(x)
        
        return pi,v
    
    
class PGAgent(object):
    def __init__(self, alpha, input_dims, gamma=0.99, layer1_size=256,layer2_size = 256, n_actions = 2):
        self.gamma = gamma
        self.policy_gradient = PolicyGradient(alpha,input_dims, layer1_size, layer2_size, n_actions = n_actions)
        self.log_probs = None
        
        
    def choose_action(self, observation):
        policy, _ = self.policy_gradient.forward(observation)
        policy = F.softmax(policy)
        action_probs = T.distributions.Categorical(policy)
        action = action_probs.sample()
        self.log_probs = action_probs.log_prob(action)
        
        return action.item()
    
    def learn(self, state, reward, state_, done):
        self.policy_gradient.optimizer.zero_grad()
        
        _, value = self.policy_gradient.forward(state)
        _, value_ = self.policy_gradient.forward(state_)
        
        reward = T.tensor(reward, dtype = T.float).to(self.policy_gradient.device)
        delta = reward + self.gamma*value_*(1-int(done)) - value
        
        actor_loss = -self.log_probs*delta
        critic_loss = delta**2
        
        (actor_loss).backward()
        self.policy_gradient.optimizer.step()