[tab] diff.diff (diff)
diff --git a/train_dynamics.py b/train_dynamics.py
index 7b6f5b6..8588b8d 100644
--- a/train_dynamics.py
+++ b/train_dynamics.py
@@ -168,12 +168,13 @@ def build_optimizer(genie: Genie, args: Args) -> tuple[nnx.Optimizer, optax.Sche
 
 def build_mesh_and_sharding(
     num_devices: int,
-) -> tuple[Mesh, NamedSharding, NamedSharding]:
+) -> tuple[Mesh, NamedSharding, NamedSharding, NamedSharding]:
     device_mesh_arr = create_device_mesh((num_devices,))
     mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
     replicated_sharding = NamedSharding(mesh, PartitionSpec())
     videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
-    return mesh, replicated_sharding, videos_sharding
+    actions_sharding = NamedSharding(mesh, PartitionSpec("data", None))
+    return mesh, replicated_sharding, videos_sharding, actions_sharding
 
 
 def shard_optimizer_states(
@@ -403,7 +404,7 @@ def main(args: Args) -> None:
     del genie
 
     # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
-    _, replicated_sharding, videos_sharding = build_mesh_and_sharding(num_devices)
+    _, replicated_sharding, videos_sharding, actions_sharding = build_mesh_and_sharding(num_devices)
 
     shard_optimizer_states(optimizer, replicated_sharding)
 
@@ -551,13 +552,19 @@ def main(args: Args) -> None:
 
     # --- TRAIN LOOP ---
     dataloader_train = (
-        jax.make_array_from_process_local_data(videos_sharding, elem)
+        {
+            "frames": jax.make_array_from_process_local_data(videos_sharding, elem["frames"]),
+            "actions": jax.make_array_from_process_local_data(actions_sharding, elem["actions"])
+        }
         for elem in train_iterator
     )
     dataloader_val = None
     if val_iterator:
         dataloader_val = (
-            jax.make_array_from_process_local_data(videos_sharding, elem)
+            {
+                "frames": jax.make_array_from_process_local_data(videos_sharding, elem["frames"]),
+                "actions": jax.make_array_from_process_local_data(actions_sharding, elem["actions"])
+            }
             for elem in val_iterator
         )
     if jax.process_index() == 0:
@@ -571,11 +578,11 @@ def main(args: Args) -> None:
     print(f"Starting training from step {step}...")
     first_step = step
     while step < args.num_steps:
-        for videos in dataloader_train:
+        for batch in dataloader_train:
             # --- Train step ---
             rng, _rng_mask = jax.random.split(rng, 2)
-            inputs = dict(videos=videos, rng=_rng_mask)
-            loss, recon, metrics = train_step(optimizer, inputs)
+            batch["rng"] = _rng_mask
+            loss, recon, metrics = train_step(optimizer, batch)
             if step == first_step:
                 print_mem_stats("After params initialized")
             metrics["lr"] = lr_schedule(step)
@@ -612,7 +619,7 @@ def main(args: Args) -> None:
                         log_dict.update(val_results["metrics"])
                     wandb.log(log_dict)
                 if step % args.log_image_interval == 0:
-                    gt_seq = inputs["videos"][0].astype(jnp.float32) / 255.0
+                    gt_seq = batch["videos"][0].astype(jnp.float32) / 255.0
                     recon_seq = recon[0].clip(0, 1)
                     comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                     comparison_seq = einops.rearrange(
diff --git a/utils/dataloader.py b/utils/dataloader.py
index 448e38a..b95547e 100644
--- a/utils/dataloader.py
+++ b/utils/dataloader.py
@@ -93,7 +93,15 @@ class ProcessEpisodeAndSlice(grain.transforms.RandomMap):
 
         seq = episode_tensor[start_idx : start_idx + self.seq_len]
 
-        return seq
+        data_dict = {
+            "frames": seq
+        }
+        if "actions" in element.keys():
+            actions_tensor = np.array(element["actions"])
+            actions = actions_tensor[start_idx : start_idx + self.seq_len]
+            data_dict["actions"] = actions
+
+        return data_dict
 
 
 def get_dataloader(