Make plotting robust

[natolambert]

The best way to do plotting is to point it to a logged data directory. WIth new sweeping, we need to part plotling code to load directories and see what we are sweeping over (sweep over model type or trajectory number for example). This will be super helpful.

Can discuss and add more notes to it. (Right now plotting is in line, and different for different model types).

@albertwilcox @calhwd15508

[natolambert]

Specific other features:

• While we are debugging training, save a pdf of the test and train loss vs epoch here. Can help us with model training.