So I adapted the reinforce.py 9 file to learn XOR, with the output sampled and using that to get a ‘reward’ (-loss). This is as opposed to using negative log likelihood as the loss.
A full SBN would also sample the hidden layer (which have neurons as Bernoulli r.v., parameterised by sigmoid(affine of prev layer) ).
Is there an efficient way to implement this as the full stochastic neuron? Would I call a simple one layer nn recursively, since Torch is dynamic?
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.autograd as autograd
from torch.autograd import Variable

torch.manual_seed(543)

class Policy(nn.Module):
    def __init__(self):
        super(Policy, self).__init__()
        self.affine1 = nn.Linear(2, 10)
        self.affine2 = nn.Linear(10, 2)
        self.affine = nn.Linear(2, 2)
        self.saved_outputs = []
        self.rewards = []
    def forward(self, x):
        x = F.sigmoid(self.affine1(x))
        action_scores = self.affine2(x)
        return F.softmax(action_scores)

model = Policy()
optimizer = optim.Adam(model.parameters(), lr=1e-2)

x_input = np.array([[1.0,1.0],[1.0,0.0],[0.0,1.0],[0.0,0.0]])
target = np.array([[0.0], [1.0],[1.0],[0.0]])

for i_episode in range(400):
    for t in range(20):
        ind = np.random.randint(4)
        xin = x_input[ind]
        tar = target[ind]

    x_input_Tensor = torch.from_numpy(xin).float().unsqueeze(0)
    probs = model(Variable(x_input_Tensor)) # prob of y
    output = probs.multinomial() # sampled from softmax
    print(xin, tar, output.data.numpy())
    model.saved_outputs.append(output) # action is a torch.LongTensor, 0s and 1s

    reward = 1.0*(output.data.numpy() == tar)
    model.rewards.append(reward)
    saved_outputs = model.saved_outputs
    rewards = []
    for r in model.rewards[::-1]:
        R = r
        rewards.insert(0, R)
    rewards = torch.Tensor(rewards)
    rewards = (rewards - rewards.mean()) / rewards.std()
    for output, r in zip(model.saved_outputs, rewards):
        output.reinforce(r)

optimizer.zero_grad()
autograd.backward(model.saved_outputs, [None for _ in model.saved_outputs])
optimizer.step()
del model.rewards[:]
del model.saved_outputs[:]