How to implement a custom loss with a non-mathematical operation (simulation) that backpropagates with PyTorch?

Question

I am writing a Neural Network, which output is not used directly for the loss-function, but rather as the input for a simulation model. After the simulation ran, I am using the simulated_value and the real_value (target) to define pretty much a MSE-Loss.

In pseudo code my loss looks something like this:

def simulation_loss(sim_params=nn_model_output, batch=batch):
     sim_value = RUN_SIMULATION(sim_params)
     sim_value.requires_grad = True
     real_value = batch['target']
     loss = nn.MSE(input=sim_value, target=real_value)
     return loss

It seems, that because I am cutting off the original gradient from nn_model_output, NN-model does not converge.

How can I run this sort of loss-function including a simulation model? Do I need to define a custom torch.autograd.Function with a backwards call? And if so, how should I pass the gradient, if there is no derivative from the simulation model possible? Or should I treat the simulation as an activation-function?

Answer 1

You can't just simply add a non-differentiable element to a computational graph.

You may, nevertheless try these far-from-simple approaches:

• Implement a differentiable version of your simulation. You can then simply use your differentiable simulation as part of your loss. There are multiple precedents for this approach, e.g. this, this and this.
• Use a score function estimator (also known as REINFORCE). See this answer for more detail and pointers. It is not always easy to make it converge.
• Use reinforcement learning (RL) approaches, like deep deterministic policy gradient (DDPG). You usually need a lot of data for RL learning approaches to converge.