Add complete training script with staged training, KL annealing, and all losses

b41f1fb

verified

asdf98 commited on 12 days ago