Feat/analysis utils

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' ")

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`