Fixing typing, logging normalization stats, fixing visualizer

octo-models · Sep 17, 2023 · 67057ae · 67057ae
1 parent ed9c970
commit 67057ae
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 15 deletions.
diff --git a/experiments/main/train.py b/experiments/main/train.py
@@ -116,8 +116,16 @@ def process_text(batch):
             batch.pop("language_instruction")
         return batch
 
+    train_data = make_dataset(**FLAGS.config.dataset_kwargs, train=True)
+    action_proprio_metadata = train_data.action_proprio_metadata
+    if save_dir is not None:
+        with tf.io.gfile.GFile(
+            os.path.join(save_dir, "action_proprio_metadata.json"), "w"
+        ) as f:
+            json.dump(action_proprio_metadata, f)
+
     train_data = (
-        make_dataset(**FLAGS.config.dataset_kwargs, train=True)
+        (train_data)
         .unbatch()
         .shuffle(FLAGS.config.shuffle_buffer_size)
         .repeat()

diff --git a/experiments/main/visualization_lib.py b/experiments/main/visualization_lib.py
@@ -82,10 +82,7 @@ class Visualizer:
 
     def __post_init__(self):
         self.dataset = make_dataset(**self.dataset_kwargs, train=False, shuffle=False)
-        builder = tfds.builder(
-            self.dataset_kwargs["name"], data_dir=self.dataset_kwargs["data_dir"]
-        )
-        self.action_proprio_stats = get_action_proprio_stats(builder, None)
+        self.action_proprio_stats = self.dataset.action_proprio_metadata
         self.trajs, self.viz_trajs = [], []
         self.visualized_trajs = False
 

diff --git a/orca/data/dataset.py b/orca/data/dataset.py
@@ -3,7 +3,7 @@
 import json
 import logging
 from functools import partial
-from typing import Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import dlimp as dl
 import numpy as np
@@ -209,6 +209,7 @@ def make_dataset(
     image_obs_keys: Union[str, List[str]] = [],
     depth_obs_keys: Union[str, List[str]] = [],
     state_obs_keys: Union[str, List[str]] = [],
+    action_proprio_metadata: Optional[dict] = None,
     **kwargs,
 ) -> tf.data.Dataset:
     """Creates a dataset from the RLDS format.
@@ -279,18 +280,17 @@ def restructure(traj):
         return traj
 
     dataset = dataset.map(restructure)
-
-    action_proprio_metadata = get_action_proprio_stats(
-        builder,
-        dataset,
-        state_obs_keys,
-        RLDS_TRAJECTORY_MAP_TRANSFORMS[name]
-        if name in RLDS_TRAJECTORY_MAP_TRANSFORMS
-        else None,
-    )
+    if action_proprio_metadata is None:
+        action_proprio_metadata = get_action_proprio_stats(
+            builder,
+            dataset,
+            state_obs_keys,
+            RLDS_TRAJECTORY_MAP_TRANSFORMS.get(name, None),
+        )
 
     dataset = apply_common_transforms(
         dataset, train=train, action_proprio_metadata=action_proprio_metadata, **kwargs
     )
+    dataset.action_proprio_metadata = action_proprio_metadata
 
     return dataset