Improve docs on training and benchmarking, and add wandb reports.

Improve docs on training and benchmarking, and add wandb reports. #907