Hi @juliagsy I wanted to check if the optimize_policy can reproduce the results indicated by the output gif in the README. Thanks