Given an input, I want to predict the right output using policy gradient. The input is a number between 0-9 and the output is also a number between 0-9. I defined the reward to be -abs(input-output) which is some kind of loss function, and using gradient ascent I hope the numbers to match but unfortunately, they don’t.
I’m using the pytorch REINFORCE implementation:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.nn.utils as utils
from torch.autograd import Variable
import numpy as np


class Policy(nn.Module):
    def __init__(self, hidden_size, num_inputs, action_space):
        super(Policy, self).__init__()
        self.action_space = action_space
        num_outputs = action_space

        self.linear1 = nn.Linear(num_inputs, hidden_size)
        self.linear2 = nn.Linear(hidden_size, num_outputs)

    def forward(self, inputs):
        x = inputs
        x = F.relu(self.linear1(x))
        action_scores = self.linear2(x)
        return F.softmax(action_scores)


class REINFORCE:
    def __init__(self, hidden_size, num_inputs, action_space, gamma=0.99):
        torch.manual_seed(0)
        np.random.seed(0)
        self.action_space = action_space
        self.model = Policy(hidden_size, num_inputs, action_space)
        self.model = self.model.cuda()
        self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)
        self.model.train()
        self.gamma = gamma
        self.rewards = []
        self.log_probs = []
        self.entropies = []

    def select_action(self, state):
        probs = self.model(Variable(state).cuda())
        action = probs.multinomial().data.unsqueeze(0)
        prob = probs[:, action[0, 0]].view(1, -1)
        log_prob = prob.log()
        entropy = - (probs * probs.log()).sum()
        return action[0], log_prob, entropy

    def update_parameters(self):
        R = torch.zeros(1, 1)
        loss = 0
        for i in reversed(range(len(self.rewards))):
            R = self.gamma * R + self.rewards[i]
            loss = loss - (self.log_probs[i] * (Variable(R).expand_as(self.log_probs[i])).cuda()).sum() - (
                        0.0001 * self.entropies[i].cuda()).sum()
        loss = loss / len(self.rewards)

        self.optimizer.zero_grad()
        loss.backward()
        utils.clip_grad_norm(self.model.parameters(), 40)
        self.optimizer.step()
        self.rewards = []
        self.log_probs = []
        self.entropies = []

    def add_experience(self, reward, log_prob, entropy):
        self.rewards.append(reward)
        self.log_probs.append(log_prob)
        self.entropies.append(entropy)


if __name__ == '__main__':
    reinforce = REINFORCE(128,1,10)
    while True:
        for i in range(10):
            input = torch.FloatTensor([i]).unsqueeze(0)
            output, log_prob, entropy = reinforce.select_action(input)
            reward = -torch.abs(output.type(torch.FloatTensor) - input)
            reinforce.add_experience(reward, log_prob, entropy)
            reinforce.update_parameters()


What am I doing wrong?
Thanks