Add unit test for DDP training

# Motivation

DDP training requires some tweak on the training procedure. To protect those tweaks and potentially other assumptions for DDP, we need some test case. This could prevent future changes breaking DDP training.

# Tasks
- [ ] Add test cases in `train_and_play_test.py`