bumped ros_datahandler, modified dynslicedataloader, and made a test with...

Fix training speed. High-level changelog:

• change internal data structure on dyn_slice_dataset from a list of trajs to a concatenated traj. With a list, we needed to use a jax.lax.switch to get the right trajectory index which worked but went really slow as # of bags increased (see here)
• make the train loop into a jax.lax.fori_loop. This puts the batch fetching and training into the same jit block, so the XLA compiler can optimize out any memory copies.
  • This drops epoch time from 13 sec to 6 with window_size=10, batch=50.
  • I left this as a separate train_jit() because it makes it harder to do certain things within an epoch. If we won't need to do anything within an epoch (e.g. early_stopping or checkpointing is always after an epoch) then we can merge and just have one train fn. A quick note from either of you would help here :)

Closes #49 (closed)

dataclass_creator/dyn_slice_dataloader.py

@@ -54,17 +54,18 @@ class DynSliceDataloader:
             self.dset = PredDynSliceDataset(data, window_size, pred_offset=pred_offset)
             
         self.batch_size = batch_size
+        self._samples = self._get_idx_batch()
                   
     def __len__(self) -> int:
         """Number of complete batches."""
         return len(self.dset) // self.batch_size
 
     def __iter__(self) -> Iterable[Batch]:
-        self._sampler = iter(self._get_idx_batch())
-        return self
+        for s in iter(self._samples):
+            yield self.dset.get_batch(s)
 
-    def __next__(self) -> Batch:
-        return self.dset.get_batch(next(self._sampler))
+    def __getitem__(self, idx: int) -> Batch:
+        return self.dset.get_batch(self._samples[idx])
     
     def _get_idx_batch(self, shuffle: bool = False, drop_last = True) -> Array:
         """Returns a list of indices to sample over."""