reformating experiment_torch.py

4 年前 · fd98cddd
--- a/experiment_torch.py
+++ b/experiment_torch.py
-
 import json
 import os
 import torch
 from mlagents_envs.timers import _thread_timer_stacks


-
-
-def run_experiment(name:str, steps:int, use_torch:bool, num_torch_threads:int, use_gpu:bool, num_envs :int= 1, config_name=None):
+def run_experiment(
+    name: str,
+    steps: int,
+    use_torch: bool,
+    num_torch_threads: int,
+    use_gpu: bool,
+    num_envs: int = 1,
+    config_name=None,
+):
    TestingConfiguration.env_name = name
    TestingConfiguration.max_steps = steps
    TestingConfiguration.use_torch = use_torch
    else:
        tf.device("/device:CPU:0")
-    if (not torch.cuda.is_available() and use_gpu):
-        return name, str(steps), str(use_torch), str(num_torch_threads), str(num_envs), str(use_gpu), "na","na","na","na","na","na","na"
+    if not torch.cuda.is_available() and use_gpu:
+        return (
+            name,
+            str(steps),
+            str(use_torch),
+            str(num_torch_threads),
+            str(num_envs),
+            str(use_gpu),
+            "na",
+            "na",
+            "na",
+            "na",
+            "na",
+            "na",
+            "na",
+        )
-    run_options = parse_command_line([f"config/ppo/{config_name}.yaml", "--num-envs", f"{num_envs}"])
-    run_options.checkpoint_settings.run_id = f"{name}_test_" +str(steps) +"_"+("torch" if use_torch else "tf")
+    run_options = parse_command_line(
+        [f"config/ppo/{config_name}.yaml", "--num-envs", f"{num_envs}"]
+    )
+    run_options.checkpoint_settings.run_id = (
+        f"{name}_test_" + str(steps) + "_" + ("torch" if use_torch else "tf")
+    )
-    timers_path = os.path.join("results", run_options.checkpoint_settings.run_id, "run_logs", "timers.json")
+    timers_path = os.path.join(
+        "results", run_options.checkpoint_settings.run_id, "run_logs", "timers.json"
+    )
    if use_torch:
        torch.set_num_threads(num_torch_threads)
    run_cli(run_options)
    with open(timers_path) as timers_json_file:
        timers_json = json.load(timers_json_file)
        total = timers_json["total"]
-        tc_advance = timers_json["children"]["TrainerController.start_learning"]["children"]["TrainerController.advance"]
-        evaluate = timers_json["children"]["TrainerController.start_learning"]["children"]["TrainerController.advance"]["children"]["env_step"]["children"]["SubprocessEnvManager._take_step"]["children"]
-        update = timers_json["children"]["TrainerController.start_learning"]["children"]["TrainerController.advance"]["children"]["trainer_advance"]["children"]["_update_policy"]["children"]
+        tc_advance = timers_json["children"]["TrainerController.start_learning"][
+            "children"
+        ]["TrainerController.advance"]
+        evaluate = timers_json["children"]["TrainerController.start_learning"][
+            "children"
+        ]["TrainerController.advance"]["children"]["env_step"]["children"][
+            "SubprocessEnvManager._take_step"
+        ][
+            "children"
+        ]
+        update = timers_json["children"]["TrainerController.start_learning"][
+            "children"
+        ]["TrainerController.advance"]["children"]["trainer_advance"]["children"][
+            "_update_policy"
+        ][
+            "children"
+        ]
        tc_advance_total = tc_advance["total"]
        tc_advance_count = tc_advance["count"]
    if use_torch:
        update_total = update["TFPPOOptimizer.update"]["total"]
        evaluate_total = evaluate["NNPolicy.evaluate"]["total"]
        update_count = update["TFPPOOptimizer.update"]["count"]
-        evaluate_count= evaluate["NNPolicy.evaluate"]["count"]
+        evaluate_count = evaluate["NNPolicy.evaluate"]["count"]
-    return name, str(steps), str(use_torch), str(num_torch_threads), str(num_envs), str(use_gpu), str(total), str(tc_advance_total), str(tc_advance_count), str(update_total), str(update_count), str(evaluate_total), str(evaluate_count)
+    return (
+        name,
+        str(steps),
+        str(use_torch),
+        str(num_torch_threads),
+        str(num_envs),
+        str(use_gpu),
+        str(total),
+        str(tc_advance_total),
+        str(tc_advance_count),
+        str(update_total),
+        str(update_count),
+        str(evaluate_total),
+        str(evaluate_count),
+    )


 def main():
-    parser.add_argument("--gpu", default = False, action="store_true", help="If true, will use the GPU")
-    parser.add_argument("--threads", default=False, action="store_true", help="If true, will try both 1 and 8 threads for torch")
-    parser.add_argument("--ball", default=False, action="store_true", help="If true, will only do 3dball")
+    parser.add_argument(
+        "--gpu", default=False, action="store_true", help="If true, will use the GPU"
+    )
+    parser.add_argument(
+        "--threads",
+        default=False,
+        action="store_true",
+        help="If true, will try both 1 and 8 threads for torch",
+    )
+    parser.add_argument(
+        "--ball",
+        default=False,
+        action="store_true",
+        help="If true, will only do 3dball",
+    )
    args = parser.parse_args()

    if args.gpu:

-    envs_config_tuples = [("3DBall", "3DBall"), ("GridWorld", "GridWorld"), ("PushBlock", "PushBlock"), ("Hallway", "Hallway"), ("CrawlerStaticTarget", "CrawlerStatic"), ("VisualHallway", "VisualHallway")]
+    envs_config_tuples = [
+        ("3DBall", "3DBall"),
+        ("GridWorld", "GridWorld"),
+        ("PushBlock", "PushBlock"),
+        ("Hallway", "Hallway"),
+        ("CrawlerStaticTarget", "CrawlerStatic"),
+        ("VisualHallway", "VisualHallway"),
+    ]
-        envs_config_tuples=[("3DBall", "3DBall")]
+        envs_config_tuples = [("3DBall", "3DBall")]
+    labels = (
+        "name",
+        "steps",
+        "use_torch",
+        "num_torch_threads",
+        "num_envs",
+        "use_gpu",
+        "total",
+        "tc_advance_total",
+        "tc_advance_count",
+        "update_total",
+        "update_count",
+        "evaluate_total",
+        "evaluate_count",
+    )
-    labels = ("name", "steps", "use_torch", "num_torch_threads", "num_envs", "use_gpu" , "total", "tc_advance_total", "tc_advance_count", "update_total", "update_count", "evaluate_total", "evaluate_count")
-    
-    f = open(f"result_data_steps_{args.steps}_envs_{args.num_envs}_gpu_{args.gpu}_thread_{args.threads}.txt", "w")
-    f.write(" ".join(labels)+ "\n")
-    
+    f = open(
+        f"result_data_steps_{args.steps}_envs_{args.num_envs}_gpu_{args.gpu}_thread_{args.threads}.txt",
+        "w",
+    )
+    f.write(" ".join(labels) + "\n")
+
-        data = run_experiment(name = env_config[0], steps=args.steps, use_torch=True, num_torch_threads=1, use_gpu=args.gpu, num_envs = args.num_envs, config_name=env_config[1])
+        data = run_experiment(
+            name=env_config[0],
+            steps=args.steps,
+            use_torch=True,
+            num_torch_threads=1,
+            use_gpu=args.gpu,
+            num_envs=args.num_envs,
+            config_name=env_config[1],
+        )
-            data = run_experiment(name = env_config[0], steps=args.steps, use_torch=True, num_torch_threads=8, use_gpu=args.gpu, num_envs = args.num_envs, config_name=env_config[1])
+            data = run_experiment(
+                name=env_config[0],
+                steps=args.steps,
+                use_torch=True,
+                num_torch_threads=8,
+                use_gpu=args.gpu,
+                num_envs=args.num_envs,
+                config_name=env_config[1],
+            )
-            f.write(" ".join(data)+ "\n")
+            f.write(" ".join(data) + "\n")
-
-        data = run_experiment(name = env_config[0], steps=args.steps, use_torch=False, num_torch_threads=1, use_gpu=args.gpu, num_envs = args.num_envs, config_name=env_config[1])
+        data = run_experiment(
+            name=env_config[0],
+            steps=args.steps,
+            use_torch=False,
+            num_torch_threads=1,
+            use_gpu=args.gpu,
+            num_envs=args.num_envs,
+            config_name=env_config[1],
+        )
-        f.write(" ".join(data)+ "\n")
+        f.write(" ".join(data) + "\n")
    for r in results:
        print(*r)
    f.close()
    main()
-