# REINFORCE for VAE Implementation Question

I want to compute the VAE loss through REINFORCE since my model's decoder is a deterministic program and is non-differentiable. The only REINFORCE implementation for VAE I was able to find used the following code.

# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function(recon_x, x, log_prob, entropy):
BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduce=False).sum(-1)
# We make the assumption that q(z1,..,zd|x) = q(z1|x)...q(zd|x)
log_prob = log_prob.sum(-1)
# The Reinforce loss is just log_prob*loss
reinforce_loss = torch.sum(log_prob*BCE.detach())
# If the prior on the latent is uniform then the KL is just the entropy of q(z|x)
# We add reinforce_loss - reinforce_loss.detach() so we can backpropagate through the encoder with REINFORCE but it doesn't modify the loss.
loss = BCE.sum() + reinforce_loss - reinforce_loss.detach() + entropy.sum()
return loss


Is the author's implementation correct? First of all, adding and subtracting reinforce_loss seems very hacky to me (though perhaps this is correct), and the KL Divergence is simply computed through Entropy(Q) rather than CrossEntropy(Q,P) - Entropy(Q) if P is uniform categorical.

Adding and subtracting some loss with .detach() is a common trick to allow the gradients to go through without changing the value. It's fine.
The cross-entropy of $$Q$$ with a uniform distribution is constant wrt $$Q$$, because $$H[Q,P] = E_Q[\log P]$$. Therefore the term is dropped.