Update covariance_experiment to work with changes in bellman_basis etc.