Merge pull request #148 from masa-su/feature/tutorial_v3

Feature/tutorial v3

tutorial/English/prepare_cartpole_dataset.py

@@ -0,0 +1,34 @@
+import gym
+import pickle
+import numpy as np
+import cv2
+env = gym.make("CartPole-v1")
+observation = env.reset()
+
+episodes = {"frames":[], "actions":[]}
+
+# for 56 *56 episode num = 500
+# for 28 * 28 episode num = 1000 
+for _episode in range(1000):
+    frames = []
+    actions = []
+    for _frame in range(30):
+        action = env.action_space.sample() # your agent here (this takes random actions)
+        frame = env.render(mode='rgb_array')
+        observation, reward, done, info = env.step(action)
+
+        img = frame
+        img = img[150:350, 200:400]
+        img = cv2.resize(img, (28,28))
+
+        frames.append(img)
+        actions.append(action)
+    observation = env.reset()
+    episodes["frames"].append(frames)
+    episodes["actions"].append(actions)
+    env.close()
+
+data = [np.array(episodes["frames"]), np.array(episodes["actions"])]
+print(data[0].shape, data[1].shape)
+with open('cartpole_28.pickle', mode='wb') as f:
+    pickle.dump(data, f)

tutorial/English/utils.py

@@ -0,0 +1,43 @@
+from torch.utils.data import Dataset
+import pickle
+import numpy as np
+import torch
+import torchvision
+import matplotlib.pyplot as plt
+
+def imshow(img_tensors):
+    img = torchvision.utils.make_grid(img_tensors)
+    npimg = img.numpy()
+    plt.figure(figsize=(16, 12))
+    plt.imshow(np.transpose(npimg, (1, 2, 0)))
+    plt.show()
+
+
+
+class DMMDataset(Dataset):
+    def __init__(self, pickle_path="cartpole_28.pickle"):
+
+        with open(pickle_path, mode='rb') as f:
+            data = pickle.load(f)
+        episode_frames, actions = data
+        # episode_frames: np.array([episode_num, one_episode_length, height, width, Channels]) (10000, 30, 28, 28, 3)
+        # actions:             np.array([episode_num, one_episode_length]) (10000, 30)
+        # HWC → CHW
+        episode_frames = episode_frames.transpose(0, 1, 4, 2, 3)
+        actions = actions[:, :, np.newaxis]
+
+        self.episode_frames = torch.from_numpy(episode_frames.astype(np.float32))
+        self.actions = torch.from_numpy(actions.astype(np.float32))
+
+    def __len__(self):
+        return len(self.episode_frames)
+
+    def __getitem__(self, idx):
+        return {
+            "episode_frames": self.episode_frames[idx] / 255,
+            "actions": self.actions[idx]
+        }
+
+
+if __name__ == "__main__":
+    pass