From 62bfe797a42c9608d74ec0fb90edb78f448c30f4 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Fri, 15 Oct 2021 02:47:52 -0700
Subject: [PATCH] adding function to load experts and combine them into single
 transition list

---
 requirements.txt                              |  5 ++--
 scratch/arec/intersimple/data/__init__.py     |  0
 scratch/arec/intersimple/data/load_experts.py | 23 +++++++++++++++++++
 .../etienne/intersimple/gail_options_image.py |  2 ++
 4 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 scratch/arec/intersimple/data/__init__.py
 create mode 100644 scratch/arec/intersimple/data/load_experts.py

diff --git a/requirements.txt b/requirements.txt
index aeac8d8..62870bb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,6 +4,7 @@ sklearn
 pytest
 json5
 tqdm
-tensorboard
 ray[tune]
-hyperopt
\ No newline at end of file
+hyperopt
+psutil
+fire
\ No newline at end of file
diff --git a/scratch/arec/intersimple/data/__init__.py b/scratch/arec/intersimple/data/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/scratch/arec/intersimple/data/load_experts.py b/scratch/arec/intersimple/data/load_experts.py
new file mode 100644
index 0000000..7a8e78b
--- /dev/null
+++ b/scratch/arec/intersimple/data/load_experts.py
@@ -0,0 +1,23 @@
+import pickle
+import imitation.data.rollout as rollout
+from tqdm import tqdm
+
+def load_experts(expert_files=[]):
+    """
+    Load expert trajectories from files and combine their transitions into a single RB
+
+    Args:
+        expert_files (list): list of expert file strings
+    Returns:
+        transitions (list): list of combined expert episode transitions 
+    """
+    transitions = []
+    for file in tqdm(expert_files):
+        with open(file, "rb") as f:
+            trajectories = pickle.load(f)
+        transitions = transitions + rollout.flatten_trajectories(trajectories)
+    return transitions
+
+if __name__=='__main__':
+    import fire
+    fire.Fire(load_experts)
\ No newline at end of file
diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py
index b35eb3f..c2561e8 100644
--- a/scratch/etienne/intersimple/gail_options_image.py
+++ b/scratch/etienne/intersimple/gail_options_image.py
@@ -303,6 +303,8 @@ if __name__ == '__main__':
  
     with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
         trajectories = pickle.load(f)
+    import pdb 
+    pdb.set_trace()
     transitions = rollout.flatten_trajectories(trajectories)
     generator = train(transitions)