From 10c6b6e59538731b89d4fb07df1a84ab0aea9133 Mon Sep 17 00:00:00 2001
From: wassname <wassname@users.noreply.github.com>
Date: Tue, 29 Dec 2020 07:58:53 +0800
Subject: [PATCH] load demonstrations use apple_gym

---
 load_demonstrations.py | 35 +++++++++++++++++++++++++++++++++++
 main.py                | 12 +++++++++---
 2 files changed, 44 insertions(+), 3 deletions(-)
 create mode 100644 load_demonstrations.py

diff --git a/load_demonstrations.py b/load_demonstrations.py
new file mode 100644
index 0000000..ef5fdac
--- /dev/null
+++ b/load_demonstrations.py
@@ -0,0 +1,35 @@
+
+
+
+
+# %%
+from gym_recording_modified.playback import get_recordings
+from tqdm.auto import tqdm
+from replay_memory import ReplayMemory
+from pathlib import Path
+
+def load_demonstrations(mem: ReplayMemory, recordings: Path):
+    records = get_recordings(str(recordings))
+    ends=records["episodes_end_point"]
+    for i in tqdm(range(len(ends))-1, desc='loading demonstrations'):
+        a = ends[i]
+        b = ends[i+1]
+        for s in range(a+1, b):
+            r = records['reward'][s]
+            o = records['observation'][s-1]
+            a = records['action'][s]
+            no = records['observation'][s]
+            t = s == b
+            mem.push(o, a, r, no, t)
+
+# %%
+
+if __name__ == "__main__":
+    # TEST
+    from replay_memory import ReplayMemory
+    from pathlib import Path
+
+    mem = ReplayMemory(10000, 42)
+    load_demonstrations(mem, Path("/media/wassname/Storage5/projects2/3ST/diy_bullet_conveyor/apple_gym/data/demonstrations"))
+
+
diff --git a/main.py b/main.py
index 010807b..6ed0bb3 100644
--- a/main.py
+++ b/main.py
@@ -7,10 +7,12 @@ import torch
 from sac import SAC
 from torch.utils.tensorboard import SummaryWriter
 from replay_memory import ReplayMemory
+from load_demonstrations import load_demonstrations
+import apple_gym.env
 
 parser = argparse.ArgumentParser(description='PyTorch Soft Actor-Critic Args')
-parser.add_argument('--env-name', default="HalfCheetah-v2",
-                    help='Mujoco Gym environment (default: HalfCheetah-v2)')
+parser.add_argument('--env-name', default="ApplePick-v0",
+                    help='Mujoco Gym environment (default: ApplePick-v0)')
 parser.add_argument('--policy', default="Gaussian",
                     help='Policy Type: Gaussian | Deterministic (default: Gaussian)')
 parser.add_argument('--eval', type=bool, default=True,
@@ -44,6 +46,8 @@ parser.add_argument('--replay_size', type=int, default=1000000, metavar='N',
                     help='size of replay buffer (default: 10000000)')
 parser.add_argument('--cuda', action="store_true",
                     help='run on CUDA (default: False)')
+parser.add_argument('--demonstrations', default=False, 
+                    help='Load demonstrations from https://github.com/erfanMhi/gym-recording-modified')
 args = parser.parse_args()
 
 # Environment
@@ -63,7 +67,9 @@ writer = SummaryWriter('runs/{}_SAC_{}_{}_{}'.format(datetime.datetime.now().str
                                                              args.policy, "autotune" if args.automatic_entropy_tuning else ""))
 
 # Memory
-memory = ReplayMemory(args.replay_size, args.seed)
+memory=ReplayMemory(args.replay_size, args.seed)
+if args.demonstrations:
+    load_demonstrations(memory, args.demonstrations)
 
 # Training Loop
 total_numsteps = 0