Support Mlflow, allow forward to have different batch size, compute more metrics

[fzyzcjy]

Hi, this is a tiny PR to

• support mlflow in addition to existing wandb and console
• allow critic forward to have different batch size than critic update (if I understand correctly, usually forward can have larger one)
• expose more metrics, mainly the "timing per token"

Example outputs of mlflow:

Example outputs of the new metrics:

[PeterSH6]

LGTM!