Document train.py's train() function

The train() function seems very chaotic to me right now. I'd like to ask we could clearly separate the computation of the Mb model and the computation of the Mb+clues model. Ideally to separate functions called from train().
