Add loss_fn_config

[maitchison]

Motivation

Some of the loss functions have configuration settings that are currently not accessible.

This PR adds an optional loss_fn_config the RL and distillation training steps and forward_backward calls.

SFT is hardcoded to cross_entropy and so was not updated.

Changes

• Added optional loss_fn_config to rl.train.Config
• For any training script that have a loss_fn, I've also added loss_fn_config.
• Added explicit argument names to some function calls.