Thanks for the nice Sobolev GAN paper. I have a question regarding the update of the Lagrange multiplier lambda. I am following Nocedal and Wright, Numerical Optimization, specifically chapter 17 where they describe the augmented Lagrangian method. It seems the lambda update should be done using the current solution, but in your code as it is written, it uses the previous solution, i.e. instead of
lambda -= rho * lambda.grad (assuming we didn't invert the sign), it should be
lambda -= rho * (1 - omega) so that the current value of omega is used. The minor difference is that the value of lambda.grad is computed using the pre-updated value of omega.
To be fair, I didn't find significant differences in performance with either, but still I thought it was worth pointing it out. Cheers!