From 215bf5d28c67697c094bb9d48f1c3d1342a6d5fa Mon Sep 17 00:00:00 2001
From: Shunichi09 <quick1st97@gmail.com>
Date: Sun, 3 May 2020 12:22:16 +0900
Subject: [PATCH] Update: example script

---
 .../plotters/plot_func.py                     | 42 ++++++++++++-------
 README.md                                     | 10 ++---
 scripts/simple_run.py                         |  4 +-
 3 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/PythonLinearNonlinearControl/plotters/plot_func.py b/PythonLinearNonlinearControl/plotters/plot_func.py
index 16866ea..3d2a7de 100644
--- a/PythonLinearNonlinearControl/plotters/plot_func.py
+++ b/PythonLinearNonlinearControl/plotters/plot_func.py
@@ -47,7 +47,7 @@ def plot_result(history, history_g=None, ylabel="x",
         axis1.legend(ncol=1, bbox_to_anchor=(0., 1.02, 1., 0.102), loc=3)
         figure.savefig(path, bbox_inches="tight", pad_inches=0.05)
 
-def plot_results(args, history_x, history_u, history_g=None):
+def plot_results(history_x, history_u, history_g=None, args=None):
     """
 
     Args:
@@ -56,14 +56,21 @@ def plot_results(args, history_x, history_u, history_g=None):
     Returns:
         None
     """
-    plot_result(history_x, history_g=history_g, ylabel="x",
-                name= args.env + "-state_history",
-                save_dir="./result/" + args.controller_type)
-    plot_result(history_u, history_g=np.zeros_like(history_u), ylabel="u",
-                name= args.env + "-input_history",
-                save_dir="./result/" + args.controller_type)
+    env = "Env"
+    controller_type = "controller"
 
-def save_plot_data(args, history_x, history_u, history_g=None):
+    if args is not None:
+        env = args.env
+        controller_type = args.controller_type
+
+    plot_result(history_x, history_g=history_g, ylabel="x",
+                name= env + "-state_history",
+                save_dir="./result/" + controller_type)
+    plot_result(history_u, history_g=np.zeros_like(history_u), ylabel="u",
+                name= env + "-input_history",
+                save_dir="./result/" + controller_type)
+
+def save_plot_data(history_x, history_u, history_g=None, args=None):
     """ save plot data
 
     Args:
@@ -72,16 +79,23 @@ def save_plot_data(args, history_x, history_u, history_g=None):
     Returns:
         None
     """
-    path = os.path.join("./result/" + args.controller_type,
-                        args.env + "-history_x.pkl")
+    env = "Env"
+    controller_type = "controller"
+
+    if args is not None:
+        env = args.env
+        controller_type = args.controller_type
+
+    path = os.path.join("./result/" + controller_type,
+                        env + "-history_x.pkl")
     save_pickle(path, history_x)
 
-    path = os.path.join("./result/" + args.controller_type,
-                        args.env + "-history_u.pkl")
+    path = os.path.join("./result/" + controller_type,
+                        env + "-history_u.pkl")
     save_pickle(path, history_u)
 
-    path = os.path.join("./result/" + args.controller_type,
-                        args.env + "-history_g.pkl")
+    path = os.path.join("./result/" + controller_type,
+                        env + "-history_g.pkl")
     save_pickle(path, history_g)
 
 def load_plot_data(env, controller_type, result_dir="./result"):
diff --git a/README.md b/README.md
index 80ddcf5..aaac70f 100644
--- a/README.md
+++ b/README.md
@@ -178,17 +178,17 @@ Use that histories to visualize the Animation or Figures.
 
 ```py
 # plot results
-plot_results(args, history_x, history_u, history_g=history_g)
-save_plot_data(args, history_x, history_u, history_g=history_g)
+plot_results(history_x, history_u, history_g=history_g)
+save_plot_data(history_x, history_u, history_g=history_g)
 
 # create animation
-animator = Animator(args, env)
+animator = Animator(env)
 animator.draw(history_x, history_g)
 ```
 
-## Run Experiments
+## Run Example Script
 
-You can run the experiments as follows:
+You can run the example script as follows:
 
 ```
 python scripts/simple_run.py --env CartPole --controller CEM --save_anim 1
diff --git a/scripts/simple_run.py b/scripts/simple_run.py
index b049075..2ef1fd8 100644
--- a/scripts/simple_run.py
+++ b/scripts/simple_run.py
@@ -37,8 +37,8 @@ def run(args):
     history_x, history_u, history_g = runner.run(env, controller, planner) 
 
     # plot results
-    plot_results(args, history_x, history_u, history_g=history_g)
-    save_plot_data(args, history_x, history_u, history_g=history_g)
+    plot_results(history_x, history_u, history_g=history_g, args=args)
+    save_plot_data(history_x, history_u, history_g=history_g, args=args)
 
     if args.save_anim:
         animator = Animator(env, args=args)