Added example of how to compute the mean episode return of a datset.

examples/download_dataset.py

@@ -13,5 +13,20 @@
 # limitations under the License.
 
 from og_marl.offline_dataset import download_and_unzip_dataset
+from og_marl.environments import smacv1
+from og_marl.offline_dataset import OfflineMARLDataset
 
-download_and_unzip_dataset("voltage_control", "case33_3min_final", dataset_base_dir="datasets")
+# Comment this out if you already downloaded the dataset
+download_and_unzip_dataset("smac_v1", "3m", dataset_base_dir="datasets")
+
+# Compute mean episode return of Good dataset
+
+env = smacv1.SMACv1("3m") # Change SMAC Scenario Here
+dataset = OfflineMARLDataset(env, f"datasets/smac_v1/3m/Good")
+
+sample_cnt =0
+tot_returns = 0
+for sample in dataset._tf_dataset:
+     sample_cnt+=1
+     tot_returns += sample["episode_return"].numpy()
+print("Mean Episode return:", tot_returns / sample_cnt)

og_marl/offline_dataset.py

@@ -21,8 +21,6 @@
 import os
 import requests
 
-Sample = namedtuple('Sample', ['observations', 'actions', 'rewards', 'done', 'episode_return', 'legal_actions', 'env_state', 'zero_padding_mask'])
-
 def get_schema_dtypes(environment):
     act_type = list(environment.action_spaces.values())[0].dtype
     schema = {}
@@ -99,6 +97,7 @@ def _decode_fn(self, record_bytes):
 
         sample["mask"] = example["zero_padding_mask"]
         sample["state"] = example["env_state"]
+        sample["episode_return"] = example["episode_return"]
 
         return sample