Minor fixes to Milestone 2