Better integration between gym & wave in examples

Author: emepetres

examples/deepq-model/render.py

@@ -1,6 +1,6 @@
 import numpy as np
 from model import DQN
-from reward import CustomReward
+from rewards import CustomReward
 
 wave = True
 render_episodes = 7
@@ -14,7 +14,7 @@
 
     env = gym.make("LunarLander-v2")
 
-model = DQN.load("trained-model")
+model = DQN.load("{}-trained-model".format("wave" if wave else "gym"))
 
 episode = render_episodes
 reward_sum = 0

examples/deepq-model/rewards.py

@@ -52,7 +52,7 @@ def reset(self):
 
 
 class LunarCustomReward(LunarLanderReward):
-    """ Custom reward that applies no penalty for engine usage (infinite fuel)
+    """ Custom reward that applies penalty for engine usage (infinite fuel)
     and allows more velocity for touching ground without crashing
     (the lander is more resistant to hits) """
 

examples/deepq-model/train.py

@@ -1,6 +1,6 @@
 import time
 from model import DQN
-from reward import CustomReward
+from rewards import CustomReward
 
 # CONFIG
 wave = True
@@ -24,6 +24,6 @@
 str_t = time.strftime("%H h, %M m, %S s", time.gmtime(t))
 print("Trained in {} during {} timesteps".format(str_t, learn_timesteps))
 
-model.save("trained-model")
+model.save("{}-trained-model".format("wave" if wave else "gym"))
 
 env.close()

examples/stable-baselines/common/callbacks.py

@@ -31,12 +31,16 @@ def callback(locals_, globals_):
         steps = n_calls * C
         if steps // N >= next_index:
             print(
-                "Saving model {}{} at step {} ...".format(
+                "Saving snapshot {}{} at step {} ...".format(
                     int(steps / order), order_str, steps
                 )
             )
             locals_["self"].save(
-                "{}{}{}{}".format(file_path, file_prefix, int(steps / order), order_str)
+                str(
+                    file_path.joinpath(
+                        "{}{}{}".format(file_prefix, int(steps / order), order_str)
+                    )
+                )
             )
             next_index = steps // N + 1
         return True

examples/stable-baselines/evaluate_model.py

@@ -1,3 +1,5 @@
+""" Evaluates Wave Lunar Lander model """
+
 import common.shutup as shutup
 
 shutup.future_warnings()
@@ -11,8 +13,8 @@
 from stable_baselines import A2C  # noqa: E402
 from arlie.envs.lunar_lander.score import LunarLanderScore  # noqa: E402
 
+wave = True
 eval_timesteps = 1e5
-multi = True
 num_cpu = 12
 
 
@@ -23,9 +25,11 @@ def evaluate(env, model, num_steps=1000):
     :param num_steps: (int) number of timesteps to evaluate it
     :return: (float) Mean reward, (int) Number of episodes performed
     """
-    scores = [LunarLanderScore() for _ in range(env.num_envs)]
-    episode_scores = [[0.0] for _ in range(env.num_envs)]
     episode_rewards = [[0.0] for _ in range(env.num_envs)]
+    if wave:
+        scores = [LunarLanderScore() for _ in range(env.num_envs)]
+        episode_scores = [[0.0] for _ in range(env.num_envs)]
+
     obs = env.reset()
     steps = (int)(num_steps // env.num_envs)
     for i in range(steps):
@@ -37,24 +41,30 @@ def evaluate(env, model, num_steps=1000):
 
         # Stats
         for i in range(env.num_envs):
-            scores[i].store_step(obs[i], actions[i], info[i])
-            episode_scores[i][-1] = scores[i].get()
             episode_rewards[i][-1] += rewards[i]
+            if wave:
+                scores[i].store_step(obs[i], actions[i], info[i])
+                episode_scores[i][-1] = scores[i].get()
             if dones[i]:
-                episode_scores[i].append(0.0)
                 episode_rewards[i].append(0.0)
-                scores[i].reset()
+                if wave:
+                    episode_scores[i].append(0.0)
+                    scores[i].reset()
 
-    mean_scores = [0.0 for _ in range(env.num_envs)]
     mean_rewards = [0.0 for _ in range(env.num_envs)]
+    if wave:
+        mean_scores = [0.0 for _ in range(env.num_envs)]
     n_episodes = 0
     for i in range(env.num_envs):
-        mean_scores[i] = np.mean(episode_scores[i][:-1])
         mean_rewards[i] = np.mean(episode_rewards[i][:-1])
+        if wave:
+            mean_scores[i] = np.mean(episode_scores[i][:-1])
         n_episodes += len(episode_rewards[i]) - 1
 
     # Compute mean reward
-    mean_score = round(np.mean(mean_scores), 1)
+    mean_score = "NaN"
+    if wave:
+        mean_score = round(np.mean(mean_scores), 1)
     mean_reward = round(np.mean(mean_rewards), 1)
 
     return mean_score, mean_reward, n_episodes
@@ -70,10 +80,11 @@ def evaluate(env, model, num_steps=1000):
         print("Path '{}' does not exist.".format(model_path))
         exit(-1)
 
+    id = "LunarLander" if wave else "LunarLander-v2"
     if num_cpu > 1:
-        env = make_multi_env(num_cpu, "LunarLander", True, render_mode=False)
+        env = make_multi_env(num_cpu, id, wave, render_mode=False, reset_mode="random")
     else:
-        env = make_env("LunarLander", True, render_mode=False, reset_mode="random")
+        env = make_env(id, wave, render_mode=False, reset_mode="random")
 
     if len(sys.argv) == 1:
         print("No model provided")

examples/stable-baselines/render_model.py

@@ -9,9 +9,8 @@
 from stable_baselines import A2C  # noqa: E402
 
 
-render_episodes = 20
 wave = True
-multi = True
+render_episodes = 20
 
 if len(sys.argv) < 2:
     print("USAGE: {} PATH-TO-MODEL-FILE".format(sys.argv[0]))
@@ -23,7 +22,7 @@
     exit(-1)
 
 id = "LunarLander" if wave else "LunarLander-v2"
-env = make_env(id, wave, port=4000)
+env = make_env(id, wave, port=4000, reset_mode="random")
 
 model = A2C.load(model_path)
 

examples/stable-baselines/train.py

@@ -5,59 +5,65 @@
 
 import os  # noqa: E402
 import time  # noqa: E402
+from pathlib import Path  # noqa: E402
 from common.utils import make_env, make_multi_env  # noqa: E402
 from common.callbacks import save_callback  # noqa: E402
 from stable_baselines.common.policies import MlpPolicy  # noqa: E402
 from stable_baselines import A2C  # noqa: E402
 
 
 # CONFIG
-model_path = "models/wave_example_a2c/"
 wave = True
 label = "a2c_example"
 order = int(1e3)
 order_str = "K"
 learn_timesteps = 24 * order
 save_interval = 2 * order
 num_cpu = 12
-log_dir = "logs"
+models_dir = "./models"
+log_dir = "./logs"
 
 if __name__ == "__main__":
-    try:
-        os.mkdir(model_path)
-    except FileExistsError:
-        pass
-    try:
-        os.mkdir(log_dir)
-    except FileExistsError:
-        pass
+    # e.g.: ./models/wave_a2c_example/
+    model_path = Path(models_dir).joinpath(
+        "{}_{}".format("wave" if wave else "gym", label)
+    )
+    # e.g.: ./logs
+    log_path = Path(log_dir)
+
+    # create folders
+    model_path.mkdir(parents=True, exist_ok=True)
+    log_path.mkdir(exist_ok=True)
 
+    # create the wave or gym environment, with or without multiprocessing
     id = "LunarLander" if wave else "LunarLander-v2"
     if num_cpu > 1:
         env = make_multi_env(num_cpu, id, wave, render_mode=False, reset_mode="random")
     else:
-        env = make_env(id, wave, render_mode=False)
+        env = make_env(id, wave, render_mode=False, reset_mode="random")
 
+    # create A2C with Mlp policy, and the callback to save snapshots
     model = A2C(MlpPolicy, env, ent_coef=0.1, verbose=0, tensorboard_log=log_dir)
     callback = save_callback(
         model_path,
-        "model-",
+        "snapshot-",
         save_interval,
         call_interval=model.n_steps * num_cpu,
         order=order,
         order_str=order_str,
     )
 
+    # save final model
     print("Training...")
     _t = time.time()
     model.learn(total_timesteps=learn_timesteps, callback=callback)
     t = time.time() - _t
     str_t = time.strftime("%H h, %M m, %S s", time.gmtime(t))
     print("Trained in {} during {} timesteps".format(str_t, learn_timesteps))
 
-    final_model = model_path + "-{}{}-final".format(
-        int(learn_timesteps / order), order_str
+    final_model = model_path.joinpath(
+        "{}{}-final".format(int(learn_timesteps / order), order_str)
     )
-    model.save(final_model)
+    model.save(str(final_model))
 
     env.close()