Added PI control