Merge pull request #37 from instadeepai/feat/analysis_utils

Feat/analysis utils
instadeepai · Aug 23, 2024 · 9fbc536 · 9fbc536
2 parents 7228ba4 + d4093c8
commit 9fbc536
Show file tree

Hide file tree

Showing 11 changed files with 2,522 additions and 216 deletions.
diff --git a/baselines/main.py b/baselines/main.py
@@ -15,7 +15,7 @@
 
 from og_marl.environments import get_environment
 from og_marl.loggers import JsonWriter, WandbLogger
-from og_marl.offline_dataset import download_and_unzip_vault
+from og_marl.vault_utils.download_vault import download_and_unzip_vault
 from og_marl.replay_buffers import FlashbaxReplayBuffer
 from og_marl.tf2.networks import CNNEmbeddingNetwork
 from og_marl.tf2.systems import get_system
@@ -24,6 +24,7 @@
 set_growing_gpu_memory()
 
 FLAGS = flags.FLAGS
+flags.DEFINE_string("dataset_source", "og_marl", "Dataset source.")
 flags.DEFINE_string("env", "pettingzoo", "Environment name.")
 flags.DEFINE_string("scenario", "pursuit", "Environment scenario name.")
 flags.DEFINE_string("dataset", "Good", "Dataset type.: 'Good', 'Medium', 'Poor' or 'Replay' ")

diff --git a/examples/dataset_analysis_demo.ipynb b/examples/dataset_analysis_demo.ipynb
diff --git a/examples/download_dataset.py b/examples/download_dataset.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 from absl import app, flags
 
-from og_marl.offline_dataset import download_and_unzip_vault
+from og_marl.vault_utils.download_vault import download_and_unzip_vault
 
 FLAGS = flags.FLAGS
 flags.DEFINE_string("env_name", "smac_v1", "Environment name.")
@@ -22,7 +22,7 @@
 
 def main(_):
     # Download vault
-    download_and_unzip_vault(FLAGS.env_name, FLAGS.scenario_name)
+    download_and_unzip_vault(FLAGS.daset_source, FLAGS.env, FLAGS.scenario)
 
     # NEXT STEPS: See `examples/dataset_api_demo.ipynb`
 

diff --git a/examples/manual_dataset_subsampling_demo.ipynb b/examples/manual_dataset_subsampling_demo.ipynb
diff --git a/og_marl/offline_dataset.py b/og_marl/offline_dataset.py