Add a callback for end of game to help agents doing learning