Update with many small changes/fixes

[HomerW]

This PR adds many small changes/fixes we've made internally.

• Variable naming changes
• Dataloader now returns actions that are already chunked
• Creating the proipio key is now handled by standardization functions
• New UNet diffusion head
• Action unnormalization happens in the model now
• Updated visualizers
• Fixed bug in attention mask, added use_correct_attention flag to maintain backwards compatibility
• Added supply_rng function wrapper that can be used to wrap the model's sample_actions function to ensure that a fresh random key is used for each call
• Fixed issue where image augmentation wasn't being applied because in configs image_augment_kwargs was not a dict

[mees]

LGTM!