How to visualize the behavior of optimization algorithms?

Hello.
I’m deep interested in optimization algorithm, such as Adam, SGD.
I’ve been trying various things to improve Adam.
Then i saw these animations. (Alec Radford’s animations)
These are very famous animations.
So, Are there any examples that can be implemented with PyTorch?
I want to apply my new algorithm to the example and visualize it like the animation below.

noisy-moons
beales-function
long-valey
saddle-point