Restructure data/model parameters
I made a medium-sized refactor to try and improve:
- ability to handle different data types (e.g. not just Franka Data)
- different model architectures / loss functions
- ergonomics for making tweaks to these as well as other parameters in tuning
To do this I mostly re-factored parameter dataclasses. Most the action is in run/parameter.py
, other changes are to accomodate these. Big changes:
- Drop AEParams, we can just use the flax modules itself as a dataclass
- Add DatasetParams, which has everything we need to load data in reproduceable way. IMO this should also be saved with the model, b/c these parameters also affect the online data handling.
- Add ExportParams, which is a super class to dump everything we need for eval/online after training. I strongly prefer a dataclass over pickling tuples so we can more easily extend in the future without refactoring everywhere we do a model load.
Pytest is passing, but I dont think everything in /run
gets pytested (e.g. validation) and I can't promise they would all run the same. Couldnt test some of those b/c don't have the local data paths.
If you guys could have a look and let me know if there's any changes you'd suggest to API or param structure, that'd be cool :)
Closes #65 (closed)