Code re-factoring and logging addition with TrainAction fixed