From fb2dda9521a7e98a2d77f66e865b670728ccd0de Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Tue, 1 May 2018 21:41:55 -0600
Subject: [PATCH] Parse args and upgrade mkdir

---
 README.md        |  4 ++--
 main.py          |  5 +----
 requirements.txt |  1 +
 utils/config.py  | 11 +++++++++--
 utils/misc.py    | 10 ++++++++--
 5 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 238cfd7..dfadaab 100644
--- a/README.md
+++ b/README.md
@@ -54,9 +54,9 @@ Support for PyTorch v0.3.x can be found in [v0.2](https://github.com/ShangtongZh
 Prediction is sampled after 110K iterations, and I only implemented one-step training
 
 # Dependency
-* MacOS 10.12,  CentO/S 6.8 or Ubuntu 16.04
+* MacOS 10.12 or Ubuntu 16.04
 * PyTorch v0.4.0
-* Python 3.6 (recommended), 3.5 or 2.7 
+* Python 3.5 or 3.6 
 * Core dependencies: `pip install -r requirements.txt`
 * Optional: [Roboschool](https://github.com/openai/roboschool), [DeepMind Control Suite](https://github.com/deepmind/dm_control)+[DMControl2Gym](dm_control2gym)
 
diff --git a/main.py b/main.py
index fbbf5b5..1211b82 100644
--- a/main.py
+++ b/main.py
@@ -362,13 +362,10 @@ def action_conditional_video_prediction():
 
 
 if __name__ == '__main__':
-    mkdir('data')
     mkdir('data/video')
     mkdir('dataset')
     mkdir('log')
-    os.system('export OMP_NUM_THREADS=1')
-    os.system('export MKL_NUM_THREADS=1')
-    torch.set_num_threads(1)
+    set_one_thread()
     # logger.setLevel(logging.DEBUG)
     logger.setLevel(logging.INFO)
 
diff --git a/requirements.txt b/requirements.txt
index 43a0c7f..80f140d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,3 +7,4 @@ tensorboardX==1.1
 scikit-image>=0.13.1
 tqdm>=4.23.0
 pandas>=0.22.0
+pathlib>=1.0.1
\ No newline at end of file
diff --git a/utils/config.py b/utils/config.py
index a1c81cf..4e2feef 100644
--- a/utils/config.py
+++ b/utils/config.py
@@ -4,9 +4,11 @@
 # declaration at the top                                              #
 #######################################################################
 from .normalizer import *
+import argparse
 
 class Config:
     def __init__(self):
+        self.parser = argparse.ArgumentParser()
         self.task_fn = None
         self.optimizer_fn = None
         self.actor_optimizer_fn = None
@@ -44,7 +46,6 @@ class Config:
         self.min_epsilon = 0
         self.save_interval = 0
         self.max_steps = 0
-        self.success_threshold = float('inf')
         self.render_episode_freq = 0
         self.rollout_length = None
         self.value_loss_weight = 1.0
@@ -60,6 +61,12 @@ class Config:
         self.test_repetitions = 10
         self.evaluation_env = None
 
-    def merge(self, config_dict):
+    def add_argument(self, *args, **kwargs):
+        self.parser.add_argument(*args, **kwargs)
+
+    def merge(self, config_dict=None):
+        if config_dict is None:
+            args = self.parser.parse_args()
+            config_dict = args.__dict__
         for key in config_dict.keys():
             setattr(self, key, config_dict[key])
diff --git a/utils/misc.py b/utils/misc.py
index f4239a2..8065109 100644
--- a/utils/misc.py
+++ b/utils/misc.py
@@ -9,6 +9,8 @@ import pickle
 import os
 import datetime
 import uuid
+import pathlib
+import torch
 
 def run_episodes(agent):
     config = agent.config
@@ -81,8 +83,12 @@ def sync_grad(target_network, src_network):
         param._grad = src_param.grad.clone()
 
 def mkdir(path):
-    if not os.path.exists(path):
-        os.mkdir(path)
+    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
+
+def set_one_thread():
+    os.environ['OMP_NUM_THREADS'] = '1'
+    os.environ['MKL_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
 
 class Batcher:
     def __init__(self, batch_size, data):