Cannot reproduce toy problem results

[Tony-Y]

I'm struggling to reproduce toy problem results. For $\beta_{2} = 0.999$ and $k = 50$, I got the following results:

I used 1000 samples to calculate the median and the 50% confidence interval at every 1000 steps. I'm using NumPy to implement my program. Could you let me know if there are any points to be careful of implementation?

Edit
My implementation: https://github.com/Tony-Y/adopt-toy-problem

Edit 2
I set the learning rate $\alpha_{t} = \eta / \sqrt{1 + \eta t}$ where $\eta \in \{ 10^{-4}, 10^{-3}, 10^{-2}, 10^{-1} \}$.

The result for $\eta = 10^{-1}$ (red line) is similar to one in the figure of the paper.

Edit 3
I have a question. Does the toy problem satisfy Assumption 2.3 (The objective function is L-smooth on $\mathbf{R}^{D}$.)?