Fixes based on UAT feedback for milestone 4