diff --git a/README.md b/README.md
index af0f60bd..92ee6825 100644
--- a/README.md
+++ b/README.md
@@ -59,7 +59,7 @@ Download environment dependencies. We will use SMACv1 in this example.
 
 Download a dataset.
 
-`python examples/download_vault.py --env=smac_v1 --scenario=3m`
+`python examples/download_dataset.py --env=smac_v1 --scenario=3m`
 
 Run a baseline. In this example we will run MAICQ.
 
@@ -77,7 +77,6 @@ We have generated datasets on a diverse set of popular MARL environments. A list
 
 <div class="collage">
   <div class="row" align="center">
-<!--     <img src="docs/assets/smac.png" alt="SMAC v1" width="16%"/> -->
     <img src="docs/assets/smacv2.png" alt="SMAC v2" width="16%"/>
     <img src="docs/assets/pistonball.png" alt="Pistonball" width="16%"/>
     <img src="docs/assets/coop_pong.png" alt="Cooperative Pong" width="16%"/>
@@ -116,7 +115,7 @@ We are in the process of migrating our datasets from TF Records to Flashbax Vaul
 | 🔴MPE | simple_adversary | 3 | Discrete. | Vector | Dense | Competitive | [source](https://pettingzoo.farama.org/environments/mpe/simple_adversary/) |
 
 ### Dataset and Vault Locations
-For OG-MARL's systems, we require the following dataset storage structure:
+For OG-MARL's systems, we require the following dataset file structure:
 
 ```
 examples/
@@ -136,19 +135,6 @@ vaults/
         |   |_> Medium/
         |   |_> Poor/
         |_> ...
-datasets/
-    |_> smac_v1/
-        |_> 3m/
-        |   |_> Good/
-        |   |_> Medium/
-        |   |_> Poor/
-        |_> ...
-    |_> smac_v2/
-        |_> terran_5_vs_5/
-        |   |_> Good/
-        |   |_> Medium/
-        |   |_> Poor/
-        |_> ...
 ...
 ```
 
diff --git a/examples/download_dataset.py b/examples/download_dataset.py
index 1331dcca..954f8971 100644
--- a/examples/download_dataset.py
+++ b/examples/download_dataset.py
@@ -11,21 +11,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from absl import app, flags
 
-from og_marl.environments import smacv1
-from og_marl.offline_dataset import OfflineMARLDataset, download_and_unzip_dataset
+from og_marl.offline_dataset import download_and_unzip_vault
 
-# Comment this out if you already downloaded the dataset
-download_and_unzip_dataset("smac_v1", "3m", dataset_base_dir="datasets")
+FLAGS = flags.FLAGS
+flags.DEFINE_string("env", "smac_v1", "Environment name.")
+flags.DEFINE_string("scenario", "3m", "Environment scenario name.")
 
-# Compute mean episode return of Good dataset
 
-env = smacv1.SMACv1("3m")  # Change SMAC Scenario Here
-dataset = OfflineMARLDataset(env, "datasets/smac_v1/3m/Good")
+def main(_):
+    # Download vault
+    download_and_unzip_vault(FLAGS.env_name, FLAGS.scenario_name)
 
-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)
+    # NEXT STEPS: See `examples/dataset_api_demo.ipynb`
+
+
+if __name__ == "__main__":
+    app.run(main)