UltrasonicLogger

ddpg.py

@@ -1,5 +1,6 @@
 import random
 from pathlib import Path
+import shutil
 
 import gym.wrappers
 import numpy as np
@@ -12,7 +13,7 @@
 from rl.memory import SequentialMemory
 from rl.random import OrnsteinUhlenbeckProcess
 
-import gym_ultrasonic
+from gym_ultrasonic.env_logger import UltrasonicLogger
 
 random.seed(27)
 np.random.seed(27)
@@ -66,16 +67,19 @@ def create_actor_critic_agent():
 agent = create_actor_critic_agent()
 WEIGHTS_PATH = Path('weights') / 'ddpg_{}_weights.h5'.format(ENV_NAME)
 WEIGHTS_PATH.parent.mkdir(exist_ok=True)
-WEIGHTS_PATH = str(WEIGHTS_PATH)
+
+env_logger = UltrasonicLogger()
 tensorboard = TensorBoard(log_dir="logs", write_grads=False, write_graph=False)
-checkpoint = ModelCheckpoint(WEIGHTS_PATH, save_weights_only=True, period=10)
+checkpoint = ModelCheckpoint(str(WEIGHTS_PATH), save_weights_only=True, period=10)
+shutil.rmtree(tensorboard.log_dir, ignore_errors=True)
 
 agent.load_weights(WEIGHTS_PATH)
-agent.fit(env, nb_steps=50000, visualize=0, verbose=2, nb_max_episode_steps=1000, callbacks=[tensorboard, checkpoint])
+agent.fit(env, nb_steps=50000, visualize=0, verbose=0, nb_max_episode_steps=1000,
+          callbacks=[env_logger, tensorboard, checkpoint])
 
 # save the weights
 agent.save_weights(WEIGHTS_PATH, overwrite=True)
 
 # test
-agent.test(env, nb_episodes=3, visualize=True, nb_max_episode_steps=5000)
-env.close()
+# agent.test(env, nb_episodes=3, visualize=True, nb_max_episode_steps=5000)
+# env.close()

gym_ultrasonic/env_logger.py

@@ -0,0 +1,51 @@
+import warnings
+
+import numpy as np
+from keras.callbacks import Callback
+
+
+class UltrasonicLogger(Callback):
+    """
+    Acts as `BaseLogger` to expand logs with the average of observations, actions, mean_q, etc.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.observations = []
+        self.rewards = []
+        self.actions = []
+        self.metrics = []
+        self.metrics_names = []
+        self.step = 0
+
+    def on_train_begin(self, logs=None):
+        self.metrics_names = self.model.metrics_names
+
+    def on_episode_begin(self, episode, logs=None):
+        """ Reset environment variables at beginning of each episode """
+        self.observations = []
+        self.rewards = []
+        self.actions = []
+        self.metrics = []
+
+    def on_episode_end(self, episode, logs=None):
+        """ Compute training statistics of the episode when done """
+        mean_q_id = self.metrics_names.index('mean_q')
+        with warnings.catch_warnings():
+            warnings.filterwarnings('error')
+            # first episode results in all nan values
+            logs['mean_q'] = np.nanmean(self.metrics, axis=0)[mean_q_id]
+        logs['reward_mean'] = np.mean(self.rewards)
+        actions_mean = np.mean(self.actions, axis=0)
+        logs['robot_move'] = actions_mean[0]
+        logs['robot_turn'] = actions_mean[1]
+        logs['dist_to_obstacles'] = np.mean(self.observations, axis=0)[1]
+        del logs['nb_steps']  # don't show total num. of steps
+
+    def on_step_end(self, step, logs=None):
+        """ Update statistics of episode after each step """
+        self.observations.append(logs['observation'])
+        self.rewards.append(logs['reward'])
+        self.actions.append(logs['action'])
+        self.metrics.append(logs['metrics'])
+        self.step += 1