Skip to content

Restructure data/model parameters

Kevin Haninger requested to merge 65-single-point-of-truth-for-data-and-model into main

Hi @nik20652 @lis07128 ,

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)

Merge request reports

Loading