Hyperparameter Logging (#308)

* add initial add_hparams implementation

* add add_hparams to logger interface

* remove commit hash from run name

* re-organize tensorboard tags

* remove comet logger

* remove additional comet logger references

* remove logger kwarg

* fix multiagent env logging

* run formatter

* fix SingleEnvExperimentTest

* update parallel env test

* run formatter

* update multiagent env experimnet

* fix softmax test problem

* update torch version in workflows

* fix integration tests

* remove unwanted print

* run formatter

.github/workflows/python-package.yml

@@ -27,7 +27,7 @@ jobs:
       run: |
         sudo apt-get install swig
         sudo apt-get install unrar
-        pip install torch~=1.11 --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install torch~=2.0 --extra-index-url https://download.pytorch.org/whl/cpu
         make install
     - name: Lint code
       run: |

.github/workflows/python-publish.yml

@@ -27,7 +27,7 @@ jobs:
         python -m pip install --upgrade pip
         sudo apt-get install swig
         sudo apt-get install unrar
-        pip install torch~=1.11 --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install torch~=2.0 --extra-index-url https://download.pytorch.org/whl/cpu
         pip install setuptools wheel
         make install
     - name: Build package

all/experiments/experiment.py

@@ -64,15 +64,11 @@ def _log_training_episode(self, returns, episode_length, fps):
             self._best_returns = returns
         self._returns100.append(returns)
         if len(self._returns100) == 100:
-            mean = np.mean(self._returns100)
-            std = np.std(self._returns100)
-            self._logger.add_summary("returns100", mean, std, step="frame")
+            self._logger.add_summary("returns100", self._returns100)
             self._returns100 = []
-        self._logger.add_eval("returns/episode", returns, step="episode")
-        self._logger.add_eval("returns/frame", returns, step="frame")
-        self._logger.add_eval("returns/max", self._best_returns, step="frame")
+        self._logger.add_eval("returns", returns)
         self._logger.add_eval("episode_length", episode_length)
-        self._logger.add_eval("fps", fps, step="frame")
+        self._logger.add_eval("fps", fps)
 
     def _log_test_episode(self, episode, returns, episode_length):
         if not self._quiet:
@@ -96,10 +92,17 @@ def _log_test(self, returns, episode_lengths):
                     episode_length_mean, episode_length_sem
                 )
             )
-        self._logger.add_summary("test_returns", np.mean(returns), np.std(returns))
-        self._logger.add_summary(
-            "test_episode_length", np.mean(episode_lengths), np.std(episode_lengths)
-        )
+        metrics = {
+            "test/returns": returns,
+            "test/episode_length": episode_lengths,
+        }
+        aggregators = ["mean", "std", "max", "min"]
+        metrics_dict = {
+            f"{metric}/{aggregator}": getattr(np, aggregator)(values)
+            for metric, values in metrics.items()
+            for aggregator in aggregators
+        }
+        self._logger.add_hparams(self._preset.hyperparameters, metrics_dict)
 
     def save(self):
         return self._preset.save("{}/preset.pt".format(self._logger.log_dir))

all/experiments/multiagent_env_experiment.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from all.logging import CometLogger, ExperimentLogger
+from all.logging import ExperimentLogger
 
 
 class MultiagentEnvExperiment:
@@ -32,10 +32,9 @@ def __init__(
         save_freq=100,
         train_steps=float("inf"),
         verbose=True,
-        logger="tensorboard",
     ):
         self._name = name if name is not None else preset.name
-        self._logger = self._make_logger(logdir, self._name, env.name, verbose, logger)
+        self._logger = self._make_logger(logdir, self._name, env.name, verbose)
         self._agent = preset.agent(logger=self._logger, train_steps=train_steps)
         self._env = env
         self._episode = 1
@@ -166,7 +165,7 @@ def _log_training_episode(self, returns, fps):
             print("frames: {}, fps: {}".format(self._frame, fps))
         for agent in self._env.agents:
             self._logger.add_eval(
-                "{}/returns/frame".format(agent), returns[agent], step="frame"
+                "{}/returns".format(agent), returns[agent], step="frame"
             )
 
     def _log_test_episode(self, episode, returns):
@@ -181,19 +180,14 @@ def _log_test(self, returns):
                 print("{} test returns (mean ± sem): {} ± {}".format(agent, mean, sem))
             self._logger.add_summary(
                 "{}/returns-test".format(agent),
-                np.mean(agent_returns),
-                np.std(agent_returns),
+                agent_returns,
             )
 
     def _save_model(self):
         if self._save_freq != float("inf") and self._episode % self._save_freq == 0:
             self.save()
 
-    def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
-        if logger == "comet":
-            return CometLogger(
-                self, agent_name, env_name, verbose=verbose, logdir=logdir
-            )
+    def _make_logger(self, logdir, agent_name, env_name, verbose):
         return ExperimentLogger(
             self, agent_name, env_name, verbose=verbose, logdir=logdir
         )

all/experiments/multiagent_env_experiment_test.py

@@ -12,7 +12,7 @@
 
 
 class MockExperiment(MultiagentEnvExperiment):
-    def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
+    def _make_logger(self, logdir, agent_name, env_name, verbose):
         self._logger = MockLogger(self, agent_name + "_" + env_name, verbose)
         return self._logger
 
@@ -50,16 +50,16 @@ def test_writes_training_returns(self):
         self.maxDiff = None
         # could not get the exact numbers to be reproducible across enviornments :(
         self.assertEqual(
-            len(experiment._logger.data["eval/first_0/returns/frame"]["values"]), 3
+            len(experiment._logger.data["eval/first_0/returns"]["values"]), 3
         )
         self.assertEqual(
-            len(experiment._logger.data["eval/first_0/returns/frame"]["steps"]), 3
+            len(experiment._logger.data["eval/first_0/returns"]["steps"]), 3
         )
         self.assertEqual(
-            len(experiment._logger.data["eval/second_0/returns/frame"]["values"]), 3
+            len(experiment._logger.data["eval/second_0/returns"]["values"]), 3
         )
         self.assertEqual(
-            len(experiment._logger.data["eval/second_0/returns/frame"]["steps"]), 3
+            len(experiment._logger.data["eval/second_0/returns"]["steps"]), 3
         )
 
     def test_writes_test_returns(self):

all/experiments/parallel_env_experiment.py

@@ -4,7 +4,7 @@
 import torch
 
 from all.environments import VectorEnvironment
-from all.logging import CometLogger, ExperimentLogger
+from all.logging import ExperimentLogger
 
 from .experiment import Experiment
 
@@ -23,11 +23,10 @@ def __init__(
         render=False,
         save_freq=100,
         verbose=True,
-        logger="tensorboard",
     ):
         self._name = name if name is not None else preset.name
         super().__init__(
-            self._make_logger(logdir, self._name, env.name, verbose, logger), quiet
+            self._make_logger(logdir, self._name, env.name, verbose), quiet
         )
         self._n_envs = preset.n_envs
         if isinstance(env, VectorEnvironment):
@@ -91,7 +90,7 @@ def train(self, frames=np.inf, episodes=np.inf):
                         self._log_training_episode(returns[i], episode_lengths[i], fps)
                         self._save_model()
                         returns[i] = 0
-                        episode_lengths[i] = 0
+                        episode_lengths[i] = -1
                         self._episode += 1
 
     def test(self, episodes=100):
@@ -139,11 +138,7 @@ def test(self, episodes=100):
     def _done(self, frames, episodes):
         return self._frame > frames or self._episode > episodes
 
-    def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
-        if logger == "comet":
-            return CometLogger(
-                self, agent_name, env_name, verbose=verbose, logdir=logdir
-            )
+    def _make_logger(self, logdir, agent_name, env_name, verbose):
         return ExperimentLogger(
             self, agent_name, env_name, verbose=verbose, logdir=logdir
         )

all/experiments/parallel_env_experiment_test.py

@@ -10,7 +10,7 @@
 
 
 class MockExperiment(ParallelEnvExperiment):
-    def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
+    def _make_logger(self, logdir, agent_name, env_name, verbose):
         self._logger = MockLogger(self, agent_name + "_" + env_name, verbose)
         return self._logger
 
@@ -32,63 +32,80 @@ def test_adds_custom_label(self):
         experiment = MockExperiment(self.make_agent(), env, name="a2c", quiet=True)
         self.assertEqual(experiment._logger.label, "a2c_CartPole-v0")
 
-    def test_writes_training_returns_episode(self):
-        self.experiment.train(episodes=4)
-        np.testing.assert_equal(
-            self.experiment._logger.data["eval/returns/episode"]["steps"],
-            np.array([1, 2, 3, 4]),
-        )
-        np.testing.assert_equal(
-            self.experiment._logger.data["eval/returns/episode"]["values"],
-            np.array([12.0, 13.0, 16.0, 16.0]),
-        )
-
     def test_writes_training_returns_frame(self):
         self.experiment.train(episodes=4)
         np.testing.assert_equal(
-            self.experiment._logger.data["eval/returns/frame"]["steps"],
-            np.array([49, 53, 65, 65]),
+            self.experiment._logger.data["eval/returns"]["steps"],
+            np.array([65, 65, 101, 125]),
         )
         np.testing.assert_equal(
-            self.experiment._logger.data["eval/returns/frame"]["values"],
-            np.array([12.0, 13.0, 16.0, 16.0]),
+            self.experiment._logger.data["eval/returns"]["values"],
+            np.array([16.0, 16.0, 25.0, 14.0]),
         )
 
     def test_writes_training_episode_length(self):
         self.experiment.train(episodes=4)
         np.testing.assert_equal(
             self.experiment._logger.data["eval/episode_length"]["steps"],
-            np.array([49, 53, 65, 65]),
+            np.array([65, 65, 101, 125]),
         )
         np.testing.assert_equal(
             self.experiment._logger.data["eval/episode_length"]["values"],
-            np.array([12.0, 13.0, 16.0, 16.0]),
+            np.array([16.0, 16.0, 25.0, 14.0]),
         )
 
+    def test_writes_hparams(self):
+        experiment = self.experiment
+        experiment.train(episodes=5)
+        returns = experiment.test(episodes=4)
+        hparam_dict, metric_dict, step = experiment._logger.hparams[0]
+        self.assertDictEqual(hparam_dict, experiment._preset.hyperparameters)
+        self.assertEqual(step, "frame")
+
     def test_writes_test_returns(self):
-        self.experiment.train(episodes=5)
-        returns = self.experiment.test(episodes=4)
-        self.assertEqual(len(returns), 4)
+        experiment = self.experiment
+        experiment.train(episodes=5)
+        returns = experiment.test(episodes=4)
+        expected_mean = 26.25
+        np.testing.assert_equal(np.mean(returns), expected_mean)
+        hparam_dict, metric_dict, step = experiment._logger.hparams[0]
         np.testing.assert_equal(
-            self.experiment._logger.data["summary/test_returns/mean"]["values"],
-            np.array([np.mean(returns)]),
+            metric_dict["test/returns/mean"],
+            np.array([expected_mean]),
+        )
+        np.testing.assert_almost_equal(
+            metric_dict["test/returns/std"], np.array([6.869]), decimal=3
+        )
+        np.testing.assert_equal(
+            metric_dict["test/returns/max"],
+            np.array([34.0]),
         )
         np.testing.assert_equal(
-            self.experiment._logger.data["summary/test_returns/std"]["values"],
-            np.array([np.std(returns)]),
+            metric_dict["test/returns/min"],
+            np.array([18.0]),
         )
 
     def test_writes_test_episode_length(self):
-        self.experiment.train(episodes=5)
-        returns = self.experiment.test(episodes=4)
-        self.assertEqual(len(returns), 4)
+        experiment = self.experiment
+        experiment.train(episodes=5)
+        returns = experiment.test(episodes=4)
+        expected_mean = 26.25
+        np.testing.assert_equal(np.mean(returns), expected_mean)
+        hparam_dict, metric_dict, step = experiment._logger.hparams[0]
+        np.testing.assert_equal(
+            metric_dict["test/episode_length/mean"],
+            np.array([expected_mean]),
+        )
+        np.testing.assert_almost_equal(
+            metric_dict["test/episode_length/std"], np.array([6.869]), decimal=3
+        )
         np.testing.assert_equal(
-            self.experiment._logger.data["summary/test_episode_length/mean"]["values"],
-            np.array([np.mean(returns)]),
+            metric_dict["test/episode_length/max"],
+            np.array([34.0]),
         )
         np.testing.assert_equal(
-            self.experiment._logger.data["summary/test_episode_length/std"]["values"],
-            np.array([np.std(returns)]),
+            metric_dict["test/episode_length/min"],
+            np.array([18.0]),
         )
 
     def test_writes_loss(self):

all/experiments/plots.py

@@ -26,15 +26,12 @@ def add_data(agent, env, file):
         data[env][agent] = np.genfromtxt(file, delimiter=",").reshape((-1, 3))
 
     for agent_dir in os.listdir(runs_dir):
-        agent = agent_dir.split("_")[0]
+        agent, env, *_ = agent_dir.split("_")
         agent_path = os.path.join(runs_dir, agent_dir)
         if os.path.isdir(agent_path):
-            for env in os.listdir(agent_path):
-                env_path = os.path.join(agent_path, env)
-                if os.path.isdir(env_path):
-                    returns100path = os.path.join(env_path, "returns100.csv")
-                    if os.path.exists(returns100path):
-                        add_data(agent, env, returns100path)
+            returns100path = os.path.join(agent_path, "returns100.csv")
+            if os.path.exists(returns100path):
+                add_data(agent, env, returns100path)
 
     return data
 

all/experiments/run_experiment.py

@@ -14,7 +14,6 @@ def run_experiment(
     save_freq=100,
     test_episodes=100,
     verbose=True,
-    logger="tensorboard",
 ):
     if not isinstance(agents, list):
         agents = [agents]
@@ -35,7 +34,6 @@ def run_experiment(
                 render=render,
                 save_freq=save_freq,
                 verbose=verbose,
-                logger=logger,
             )
             experiment.save()
             experiment.train(frames=frames)

all/experiments/single_env_experiment.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from all.logging import CometLogger, ExperimentLogger
+from all.logging import ExperimentLogger
 
 from .experiment import Experiment
 
@@ -21,11 +21,10 @@ def __init__(
         render=False,
         save_freq=100,
         verbose=True,
-        logger="tensorboard",
     ):
         self._name = name if name is not None else preset.name
         super().__init__(
-            self._make_logger(logdir, self._name, env.name, verbose, logger), quiet
+            self._make_logger(logdir, self._name, env.name, verbose), quiet
         )
         self._logdir = logdir
         self._preset = preset
@@ -34,7 +33,7 @@ def __init__(
         self._render = render
         self._frame = 1
         self._episode = 1
-        self._save_freq = 100
+        self._save_freq = save_freq
 
         if render:
             self._env.render(mode="human")
@@ -116,11 +115,7 @@ def _run_test_episode(self, test_agent):
     def _done(self, frames, episodes):
         return self._frame > frames or self._episode > episodes
 
-    def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
-        if logger == "comet":
-            return CometLogger(
-                self, agent_name, env_name, verbose=verbose, logdir=logdir
-            )
+    def _make_logger(self, logdir, agent_name, env_name, verbose):
         return ExperimentLogger(
             self, agent_name, env_name, verbose=verbose, logdir=logdir
         )