Some general fixes to train visuals and train behavior