adapted sample_action/trajectory to work with new octo version

octo/data/dataset.py

@@ -393,9 +393,9 @@ def is_nonzero_length(traj):
             full_dataset = full_dataset.filter(ModuleSpec.instantiate(filter_fcn_spec))
         if ignore_errors:
             full_dataset = full_dataset.ignore_errors()
-        # for traj in full_dataset:
-        #     restructure(traj)
+
         full_dataset = full_dataset.traj_map(restructure).filter(is_nonzero_length)
+
         # tries to load from cache, otherwise computes on the fly
         dataset_statistics = get_dataset_statistics(
             full_dataset,

octo/model/octo_model.py

@@ -259,7 +259,7 @@ def sample_future_actions(
         unnormalization_statistics: Optional[Data] = None,
         normalization_type: NormalizationType = NormalizationType.NORMAL,
         beam:int = 1,
-        pad_mask: Optional[ArrayLike] = None,
+        timestep_pad_mask: Optional[ArrayLike] = None,
         train: bool = False,
         argmax: bool = False,
         sample_shape: Tuple[int, ...] = (),
@@ -281,7 +281,7 @@ def sample_future_actions(
             actions: (*sample_shape, batch_size, action_horizon, action_dim)
         """
         if timestep_pad_mask is None:
-            timestep_pad_mask = observations["timestep_pad_mask"]
+            timestep_pad_mask = observations["pad_mask"]
 
         transformer_outputs = self.run_transformer(
             observations, tasks, timestep_pad_mask, train=train
@@ -346,14 +346,14 @@ def sample_trajectory(
         unnormalization_statistics: Optional[Data] = None,
         normalization_type: NormalizationType = NormalizationType.NORMAL,
         beam: int = 1,
-        pad_mask: Optional[ArrayLike] = None,
+        timestep_pad_mask: Optional[ArrayLike] = None,
         train: bool = False,
         argmax: bool = False,
         sample_shape: Tuple[int, ...] = (),
         rng: Optional[PRNGKey] = None,
         temperature: float = 1.0,
     ):
-        if pad_mask is None:
+        if timestep_pad_mask is None:
             pad_mask = observations["pad_mask"]
 
         transformer_outputs = self.run_transformer(