Restructure data/model parameters

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)

dataclass_creator/utils_dataclass/real_data_gen.py

@@ -11,12 +11,10 @@ from os.path import join
 from pathlib import Path
 from typing import Any, Iterable
 
-#from optree import tree_map
-
-from opt_tools  import SysVarSet, Traj, Batch, tree_map_relaxed
+from opt_tools.jax_tools  import SysVarSet, Traj, Batch, tree_map_relaxed
 from ros_datahandler import get_bags_iter
 from ros_datahandler import JointPosData, ForceData, JointState, WrenchStamped
-from dataclass_creator.franka_dataclass import FrankaData
+from dataclass_creator.custom_dataclasses import FrankaData
 
 
 def min_max_normalize_data(data: jnp.ndarray):
@@ -105,7 +103,12 @@ def bag_info(root='data/without_search_strategy/', element: int = 0):
     print('Topics of one Rosbag:', topics)
 
 
-def load_dataset(root='experiments/example_data/train', name: str = 'SUCCESS') -> Iterable[Traj]:
+def load_dataset(
+        datatype: SysVarSet,
+        sync_stream: str,
+        root: Path,
+        name: str = "SUCCESS",
+) -> Iterable[Traj]:
     """collects rosbags in specified dir and return a formatted dataclass containg ragged list*3
 
     Args:
@@ -116,26 +119,15 @@ def load_dataset(root='experiments/example_data/train', name: str = 'SUCCESS') -
         FormattedDataclass: data containing the modalities corresponding to the attributes initialized for RosDatahandler
     """
 
-    bags = glob.glob(join(root, f'*{name}*.bag'))
-    assert len(bags) > 0, f"I expected bags! No bags on {root}"
-    bags = [Path(bag) for bag in bags]
-
-    print(f'Containing # {len(bags)} rosbags, named {name}')
-
-    init_frankadata = FrankaData(jt = JointPosData('/joint_states', JointState),
-                                 force = ForceData('/franka_state_controller/F_ext', WrenchStamped))
+    bags = glob.glob(join(root, f"*{name}*"))
+    assert len(bags) > 0, f"I expected bags! No bags on {root} with filter {name}"
 
-    # @kevin 2.7.24: data is now a list of Trajs; leaves are np arrays for each bags
-    data = get_bags_iter(bags, init_frankadata, sync_stream='force')
+    bags = [Path(bag) for bag in bags]
+    logging.info(f'Containing # {len(bags)} rosbags for filter {name}')
 
+    data = get_bags_iter(bags, datatype, sync_stream=sync_stream)
     data = [d.cast_traj() for d in data]
-    
-    assert len(data) == len(bags)    
-    assert isinstance(data[0], FrankaData), 'data is somehow not a FrankaData? huh?'
-    assert isinstance(data[0], SysVarSet), 'data is somehow not a sysvarset? huh?'
-    assert isinstance(data[0], Traj), 'data should be a traj'
-    assert not isinstance(data[0], Batch), 'data should not be a batch yet'
-    
+       
     return data