Fix bug where state distribution labels were sometimes wrong. (#510)