[Cleaning] Improved PS init, and fixed a bug in forest level.