Add max-episode-steps flag and use in tests

ultra/MANIFEST.in

@@ -3,3 +3,4 @@ include ultra/agent_pool.json
 include config.yaml
 include ultra/scenarios/task1/config.yaml
 include ultra/scenarios/task2/config.yaml
+recursive-include ultra/scenarios/pool

ultra/tests/scenarios/task00/config.yaml

@@ -32,7 +32,7 @@ levels:
   # Used in evaluation test
   eval_test:
     train:
-      total: 20
+      total: 5
       intersection_types:
         2lane_t:
           percent: 0.2
@@ -53,7 +53,7 @@ levels:
                   [50kmh,no-traffic,0.01],[70kmh,no-traffic,0.01],[100kmh,no-traffic,0.01], #3%
                   [50kmh,high-density,0.01],[70kmh,high-density,0.01],[100kmh,high-density,0.01]] # 3%
     test:
-      total: 10
+      total: 2
       intersection_types:
         2lane_c:
           percent: 0.5
@@ -66,4 +66,4 @@ levels:
           specs: [[50kmh,low-density,0.21],[70kmh,low-density,0.20],[100kmh,low-density,0.20], #61%
                   [50kmh,mid-density,0.11],[70kmh,mid-density,0.11],[100kmh,mid-density,0.11], #33%
                   [50kmh,no-traffic,0.01],[70kmh,no-traffic,0.01],[100kmh,no-traffic,0.01], #3%
-                  [50kmh,high-density,0.01],[70kmh,high-density,0.01],[100kmh,high-density,0.01]] # 3%
+                  [50kmh,high-density,0.01],[70kmh,high-density,0.01],[100kmh,high-density,0.01]] # 3%

ultra/tests/test_analysis.py

@@ -34,12 +34,12 @@
 class AnalysisTest(unittest.TestCase):
     def test_interface_analyze(self):
         try:
-            save_dir = "tests/task/eval_test/eval"
-            output_scenario = "tests/scenarios/analysis_output"
+            save_dir = "tests/analysis_test/scenarios/"
+            output = "tests/analysis_test/output/"
             if os.path.exists(save_dir):
                 shutil.rmtree(save_dir)
-            if os.path.exists(output_scenario):
-                shutil.rmtree(output_scenario)
+            if os.path.exists(output):
+                shutil.rmtree(save_dir)
             build_scenarios(
                 task="task00",
                 level_name="easy",
@@ -49,31 +49,36 @@ def test_interface_analyze(self):
                 save_dir=save_dir,
             )
 
-            if not os.path.exists("tests/scenarios/analysis_output"):
-                os.makedirs("tests/scenarios/analysis_output")
+            if not os.path.exists(output):
+                os.makedirs(output)
 
             os.system(
-                "python ultra/scenarios/interface.py analyze --scenarios tests/task/eval_test/ --max-steps 1000 --end-by-stopwatcher --output tests/scenarios/analysis_output"
+                f"python ultra/scenarios/interface.py analyze --scenarios {save_dir} --max-steps 600 --end-by-stopwatcher --output {output}"
             )
-            for dirpath, dirnames, files in os.walk("tests/tasks/eval_test/"):
+            for dirpath, dirnames, files in os.walk(save_dir):
                 if "traffic" in dirpath:
                     self.assertTrue("all.rou.xml" in files)
 
-            self.assertTrue(
-                os.path.exists("tests/scenarios/analysis_output/analysis.pkl")
-            )
-            self.assertTrue(os.path.exists("tests/scenarios/analysis_output/"))
+            self.assertTrue(os.path.exists("tests/analysis_test/output/analysis.pkl"))
         except Exception as err:
             print(err)
             self.assertTrue(False)
 
+        if os.path.exists(save_dir):
+            self.assertTrue(True)
+            shutil.rmtree(save_dir)
+
+        if os.path.exists(output):
+            self.assertTrue(True)
+            shutil.rmtree(output)
+
     def test_analyze_scenario(self):
-        save_dir = "tests/scenarios/maps/easy/"
-        output_scenario = "tests/scenarios/output"
+        save_dir = "tests/analysis_test/scenarios/"
+        output = "tests/analysis_test/output/"
         if os.path.exists(save_dir):
             shutil.rmtree(save_dir)
-        if os.path.exists(output_scenario):
-            shutil.rmtree(output_scenario)
+        if os.path.exists(output):
+            shutil.rmtree(output)
 
         build_scenarios(
             task="task00",
@@ -83,7 +88,7 @@ def test_analyze_scenario(self):
             root_path="tests/scenarios",
             save_dir=save_dir,
         )
-        scenarios = glob.glob(f"{save_dir}/map*0*")
+        scenarios = glob.glob(f"{save_dir}")
         try:
             ray.init(ignore_reinit_error=True)
             analyzer = ScenarioAnalysis.remote()
@@ -93,7 +98,7 @@ def test_analyze_scenario(self):
                         end_by_stopwatcher=True,
                         scenarios=scenarios,
                         ego=None,
-                        max_episode_steps=3000,
+                        max_episode_steps=800,
                         policy=None,
                         video_rate=100,
                         timestep_sec=0.1,
@@ -104,14 +109,28 @@ def test_analyze_scenario(self):
                     )
                 ]
             )
-            if not os.path.exists(output_scenario):
-                os.makedirs(output_scenario)
-            ray.wait([analyzer.save_data.remote(save_dir=output_scenario)])
-            self.assertTrue(os.path.exists(f"{output_scenario}/analysis.pkl"))
+            if not os.path.exists(output):
+                os.makedirs(output)
+            ray.wait([analyzer.save_data.remote(save_dir=output)])
+            self.assertTrue(os.path.exists(f"{output}analysis.pkl"))
+
+            for dirpath, dirnames, files in os.walk(save_dir):
+                if "traffic" in dirpath:
+                    self.assertTrue("all.rou.xml" in files)
+
+            self.assertTrue(os.path.exists("tests/analysis_test/output/analysis.pkl"))
         except ray.exceptions.WorkerCrashedError as err:
             print(err)
             self.assertTrue(False)
 
+        if os.path.exists(save_dir):
+            self.assertTrue(True)
+            shutil.rmtree(save_dir)
+
+        if os.path.exists(output):
+            self.assertTrue(True)
+            shutil.rmtree(output)
+
     def test_save_histogram(self):
         try:
             figure_name = "tests/scenarios/output"
@@ -127,3 +146,8 @@ def test_save_histogram(self):
         except Exception as err:
             print(err)
             self.assertTrue(False)
+
+    @classmethod
+    def tearDownClass(cls):
+        if os.path.exists("tests/analysis_test"):
+            shutil.rmtree("tests/analysis_test")

ultra/tests/test_evaluate.py

@@ -46,7 +46,7 @@ def setUpClass(cls):
         if not os.path.exists(path):
             # Generate models before evaluation tests
             os.system(
-                "python ultra/train.py --task 00 --level eval_test --policy sac --headless True --episodes 10 --eval-rate 200 --eval-episodes 1 --log-dir tests/sac_test_models"
+                "python ultra/train.py --task 00 --level eval_test --policy sac --headless True --episodes 1 --eval-rate 1 --eval-episodes 1 --max-episode-steps 2 --log-dir tests/sac_test_models"
             )
 
     def test_a_folders(self):
@@ -92,6 +92,7 @@ def run_experiment():
                         episode=episode,
                         eval_rate=10,
                         eval_episodes=1,
+                        max_episode_steps=2,
                         policy_class=policy_class,
                         scenario_info=("00", "eval_test"),
                         timestep_sec=0.1,
@@ -145,7 +146,7 @@ def test_evaluate_cli(self):
         ray.shutdown()
         try:
             os.system(
-                f"python ultra/evaluate.py --task 00 --level eval_test --models {model_dir} --episodes 1 --log-dir {log_dir} --headless True"
+                f"python ultra/evaluate.py --task 00 --level eval_test --models {model_dir} --episodes 1 --max-episode-steps 2 --log-dir {log_dir} --headless True"
             )
             self.assertTrue(True)
         except Exception as err:
@@ -160,13 +161,10 @@ def test_evaluate_cli(self):
 
     def test_evaluate_agent(self):
         seed = 2
-        model = glob.glob("tests/sac_test_models/*/models/")[0]
-        log_dir = "tests/output_eval_agent_logs"
+        model = glob.glob("tests/sac_test_models/*/models/0")[0]
+        log_dir = "tests/output_eval_agent_logs/"
         policy_class = "ultra.baselines.sac:sac-v0"
 
-        if not os.path.exists(log_dir):
-            os.makedirs(log_dir)
-
         ray.shutdown()
         ray.init(ignore_reinit_error=True)
         try:
@@ -179,6 +177,7 @@ def test_evaluate_agent(self):
                 checkpoint_dir=model,
                 scenario_info=("00", "eval_test"),
                 num_episodes=1,
+                max_episode_steps=2,
                 timestep_sec=0.1,
                 headless=True,
                 log_dir=log_dir,
@@ -218,11 +217,10 @@ def extract(path):
 
 def prepare_test_env_agent(headless=True):
     timestep_sec = 0.1
-    policy_class = "ultra.baselines.sac:sac-v0"
     spec = BaselineAgentSpec(
         action_type=ActionSpaceType.Continuous,
         policy_class=SACPolicy,
-        max_episode_steps=100,
+        max_episode_steps=2,
     )
     env = gym.make(
         "ultra.env:ultra-v0",

ultra/tests/test_train.py

@@ -34,22 +34,14 @@
 
 class TrainTest(unittest.TestCase):
     def test_train_cli(self):
-        try:
-            os.system(
-                "python ultra/train.py --task 00 --level easy --episodes 1 --log-dir tests/logs"
-            )
-        except Exception as err:
-            print(err)
-            self.assertTrue(False)
-
-    def test_locate_log_directory(self):
         log_dir = "tests/logs"
         try:
             os.system(
-                f"python ultra/train.py --task 00 --level easy --policy ppo --episodes 1 --log-dir {log_dir}"
+                "python ultra/train.py --task 00 --level easy --episodes 1 --max-episode-steps 2 --log-dir tests/logs"
             )
         except Exception as err:
             print(err)
+            self.assertTrue(False)
 
         if os.path.exists(log_dir):
             self.assertTrue(True)
@@ -70,6 +62,7 @@ def test_train_single_agent(self):
                         scenario_info=("00", "easy"),
                         policy_class=policy_class,
                         num_episodes=1,
+                        max_episode_steps=2,
                         eval_info={"eval_rate": 1000, "eval_episodes": 2,},
                         timestep_sec=0.1,
                         headless=True,

ultra/tests/test_ultra_package.py

@@ -58,6 +58,7 @@ def test_simple_train_run(self):
                         scenario_info=("00", "eval_test"),
                         policy_class=policy_class,
                         num_episodes=1,
+                        max_episode_steps=2,
                         eval_info={"eval_rate": 1000, "eval_episodes": 2,},
                         timestep_sec=0.1,
                         headless=True,

ultra/ultra/config.yaml

@@ -0,0 +1,41 @@
+tasks:
+  task00: # unittest
+    easy:
+      train: "../tests/task/train_mock*"
+      test:  "../tests/task/test_mock*"
+    eval_test:
+      train: "../tests/task/eval_test/eval*"
+      test:  "../tests/task/eval_test/eval*"
+  task0:
+    no-traffic:
+      train: "scenarios/task1/train_no-traffic*"
+      test:  "scenarios/task1/train_no-traffic*"
+    easy:
+      train: "scenarios/task1/train_easy*"
+      test:  "scenarios/task1/train_easy*"
+    medium:
+      train: "scenarios/task1/train_medium*"
+      test:  "scenarios/task1/train_medium*"
+    hard:
+      train: "scenarios/task1/train_hard*"
+      test:  "scenarios/task1/train_hard*"
+  task1:
+    no-traffic:
+      train: "scenarios/task1/train_no-traffic*"
+      test:  "scenarios/task1/test_no-traffic*"
+    easy:
+      train: "scenarios/task1/train_easy*"
+      test:  "scenarios/task1/test_easy*"
+    medium:
+      train: "scenarios/task1/train_medium*"
+      test:  "scenarios/task1/test_medium*"
+    hard:
+      train: "scenarios/task1/train_hard*"
+      test:  "scenarios/task1/test_hard*"
+  task2:
+    lo-hi:
+      train: "scenarios/task2/train_lo-hi*"
+      test:  "scenarios/task2/test_lo-hi*"
+    hi-lo:
+      train: "scenarios/task2/train_hi-lo*"
+      test:  "scenarios/task2/test_hi-lo*"

ultra/ultra/env/ultra_env.py

@@ -173,16 +173,15 @@ def step(self, agent_actions):
         return observations, rewards, agent_dones, infos
 
     def get_task(self, task_id, task_level):
-        base_dir = os.path.dirname(__file__)
-        # print(base_dir)
-        config_path = os.path.join(base_dir, "../../config.yaml")
-        # print(config_path)
+        base_dir = os.path.join(os.path.dirname(__file__), "../")
+        config_path = os.path.join(base_dir, "config.yaml")
 
         with open(config_path, "r") as task_file:
             scenarios = yaml.safe_load(task_file)["tasks"]
-            # print(scenarios)
             task = scenarios[f"task{task_id}"][task_level]
-            # print(task)
+
+        task["train"] = os.path.join(base_dir, task["train"])
+        task["test"] = os.path.join(base_dir, task["test"])
         return task
 
     @property