Experiments for "Geometric Insights into the Convergence of Nonlinear TD Learning" Dependencies jax numpy matplotlib jupyter notebook