Fix/repart 2d

[jaredstrader]

I'm going ahead and creating the PR, but perhaps, we wait to merge until we resolve why there are still a few places not being assigned rooms. Adding @SiyiHu since this is relevant for the scene prediction experiments.