add that to autograd for sure.
i had a look at the StochasticFunction class, however, I dont think it is a good way to implement as the reward in backward might not always applied. Maybe the best way is just to wrap it as a autograd.Function?