Small changes to controler