Compile time skews tqdm ETA

[Drenderer]

When running the training loop, the first iteration usually takes the longest because JAX needs to compile the make_step function. When using the tqdm progress bar to show training progress, it significantly skews the estimated remaining training time, especially at the beginning of training.

Proposed solutions:

• JAX warmup step. Simply do one (wasted) step before the training loop to compile make_step before tqdm even notices.
• Reset the tqdm internals after the first step (seems sketchy).