Comet (#215)

* Add CometWriter and make related changes

* Add comet import

* Fix CometWriter parent class

* Improve CometWriter documentation

* Address PR comments

* Fix hyperparameter tracking

* Fix _make_writer test failure

* Fix MockWriter test issue

* Fix MockWriter test issue

* Fix integration tests

* Add experiment close method

all/experiments/__init__.py

@@ -4,6 +4,7 @@
 from .parallel_env_experiment import ParallelEnvExperiment
 from .multiagent_env_experiment import MultiagentEnvExperiment
 from .writer import ExperimentWriter
+from .writer import CometWriter
 from .plots import plot_returns_100
 from .slurm import SlurmExperiment
 from .watch import watch, load_and_watch
@@ -16,6 +17,7 @@
     "MultiagentEnvExperiment",
     "SlurmExperiment",
     "ExperimentWriter",
+    "CometWriter",
     "watch",
     "load_and_watch",
 ]

all/experiments/experiment.py

@@ -81,3 +81,6 @@ def _log_test(self, returns):
 
     def save(self):
         return self._preset.save('{}/preset.pt'.format(self._writer.log_dir))
+
+    def close(self):
+        self._writer.close()

all/experiments/multiagent_env_experiment.py

@@ -1,7 +1,7 @@
 from timeit import default_timer as timer
 import numpy as np
 from scipy import stats
-from .writer import ExperimentWriter
+from .writer import ExperimentWriter, CometWriter
 from .experiment import Experiment
 
 
@@ -32,9 +32,10 @@ def __init__(
             save_freq=100,
             train_steps=float('inf'),
             write_loss=True,
+            writer="tensorboard"
     ):
         self._name = name if name is not None else preset.__class__.__name__
-        self._writer = self._make_writer(logdir, self._name, env.name, write_loss)
+        self._writer = self._make_writer(logdir, self._name, env.name, write_loss, writer)
         self._agent = preset.agent(writer=self._writer, train_steps=train_steps)
         self._env = env
         self._episode = 0
@@ -172,5 +173,7 @@ def _save_model(self):
         if self._save_freq != float('inf') and self._episode % self._save_freq == 0:
             self._preset.save('{}/preset.pt'.format(self._writer.log_dir))
 
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
+        if writer == "comet":
+            return CometWriter(self, agent_name, env_name, loss=write_loss, logdir=logdir)
         return ExperimentWriter(self, agent_name, env_name, loss=write_loss, logdir=logdir)

all/experiments/multiagent_env_experiment_test.py

@@ -41,9 +41,12 @@ def _get_step(self, _type):
             return self.experiment.episode
         return _type
 
+    def close(self):
+        pass
+
 
 class MockExperiment(MultiagentEnvExperiment):
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
         self._writer = MockWriter(self, agent_name + '_' + env_name, write_loss)
         return self._writer
 

all/experiments/parallel_env_experiment.py

@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 from all.core import State
-from .writer import ExperimentWriter
+from .writer import ExperimentWriter, CometWriter
 from .experiment import Experiment
 
 
@@ -19,10 +19,11 @@ def __init__(
             logdir='runs',
             quiet=False,
             render=False,
-            write_loss=True
+            write_loss=True,
+            writer="tensorboard"
     ):
         self._name = name if name is not None else preset.__class__.__name__
-        super().__init__(self._make_writer(logdir, self._name, env.name, write_loss), quiet)
+        super().__init__(self._make_writer(logdir, self._name, env.name, write_loss, writer), quiet)
         self._n_envs = preset.n_envs
         self._envs = env.duplicate(self._n_envs)
         self._preset = preset
@@ -138,5 +139,7 @@ def _fps(self, i):
         end_time = timer()
         return (self._frame - self._episode_start_frames[i]) / (end_time - self._episode_start_times[i])
 
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
+        if writer == "comet":
+            return CometWriter(self, agent_name, env_name, loss=write_loss, logdir=logdir)
         return ExperimentWriter(self, agent_name, env_name, loss=write_loss, logdir=logdir)

all/experiments/parallel_env_experiment_test.py

@@ -8,7 +8,7 @@
 
 
 class MockExperiment(ParallelEnvExperiment):
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
         self._writer = MockWriter(self, agent_name + '_' + env_name, write_loss)
         return self._writer
 

all/experiments/run_experiment.py

@@ -10,7 +10,8 @@ def run_experiment(
         quiet=False,
         render=False,
         test_episodes=100,
-        write_loss=True
+        write_loss=True,
+        writer="tensorboard"
 ):
     if not isinstance(agents, list):
         agents = [agents]
@@ -29,11 +30,13 @@ def run_experiment(
                 logdir=logdir,
                 quiet=quiet,
                 render=render,
-                write_loss=write_loss
+                write_loss=write_loss,
+                writer=writer
             )
             experiment.train(frames=frames)
             experiment.test(episodes=test_episodes)
             experiment.save()
+            experiment.close()
 
 
 def get_experiment_type(preset):

all/experiments/single_env_experiment.py

@@ -1,6 +1,7 @@
 from timeit import default_timer as timer
 import numpy as np
-from .writer import ExperimentWriter
+from .writer import ExperimentWriter, CometWriter
+
 from .experiment import Experiment
 
 
@@ -16,10 +17,11 @@ def __init__(
             logdir='runs',
             quiet=False,
             render=False,
-            write_loss=True
+            write_loss=True,
+            writer="tensorboard"
     ):
         self._name = name if name is not None else preset.__class__.__name__
-        super().__init__(self._make_writer(logdir, self._name, env.name, write_loss), quiet)
+        super().__init__(self._make_writer(logdir, self._name, env.name, write_loss, writer), quiet)
         self._logdir = logdir
         self._preset = preset
         self._agent = self._preset.agent(writer=self._writer, train_steps=train_steps)
@@ -101,5 +103,7 @@ def _run_test_episode(self, test_agent):
     def _done(self, frames, episodes):
         return self._frame > frames or self._episode > episodes
 
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
+        if writer == "comet":
+            return CometWriter(self, agent_name, env_name, loss=write_loss, logdir=logdir)
         return ExperimentWriter(self, agent_name, env_name, loss=write_loss, logdir=logdir)

all/experiments/single_env_experiment_test.py

@@ -40,9 +40,12 @@ def _get_step(self, _type):
             return self.experiment.episode
         return _type
 
+    def close(self):
+        pass
+
 
 class MockExperiment(SingleEnvExperiment):
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
         self._writer = MockWriter(self, agent_name + '_' + env_name, write_loss)
         return self._writer
 

all/experiments/writer.py

@@ -9,7 +9,7 @@
 
 class ExperimentWriter(SummaryWriter, Writer):
     '''
-    The Writer object used by all.experiments.Experiment.
+    The default Writer object used by all.experiments.Experiment.
     Writes logs using tensorboard into the current logdir directory ('runs' by default),
     tagging the run with a combination of the agent name, the commit hash of the
     current git repo of the working directory (if any), and the current time.
@@ -71,6 +71,73 @@ def _get_step(self, _type):
             return self._experiment.episode
         return _type
 
+    def close(self):
+        pass
+
+
+class CometWriter(Writer):
+    '''
+    A Writer object to be used by all.experiments.Experiment.
+    Writes logs using comet.ml Requires an API key to be stored in .comet.config or as an environment variable.
+    Look at https://www.comet.ml/docs/python-sdk/advanced/#python-configuration for more info.
+    Args:
+        experiment (all.experiments.Experiment): The Experiment associated with the Writer object.
+        agent_name (str): The name of the Agent the Experiment is being performed on
+        env_name (str): The name of the environment the Experiment is being performed in
+        loss (bool, optional): Whether or not to log loss/scheduling metrics, or only evaluation and summary metrics.
+        logdir (str): The directory where run information is stored.
+    '''
+
+    def __init__(self, experiment, agent_name, env_name, loss=True, logdir='runs'):
+        self.env_name = env_name
+        self._experiment = experiment
+        self._loss = loss
+
+        try:
+            from comet_ml import Experiment
+        except ImportError as e:
+            print("Failed to import comet_ml. CometWriter requires that comet_ml be installed")
+            raise e
+        try:
+            self._comet = Experiment(project_name=env_name)
+        except ImportError as e:
+            print("See https://www.comet.ml/docs/python-sdk/warnings-errors/ for more info on this error.")
+            raise e
+        except ValueError as e:
+            print("See https://www.comet.ml/docs/python-sdk/advanced/#python-configuration for more info on this error.")
+            raise e
+
+        self._comet.set_name(agent_name)
+        self.log_dir = logdir
+
+    def add_loss(self, name, value, step="frame"):
+        if self._loss:
+            self.add_evaluation("loss/" + name, value, step)
+
+    def add_evaluation(self, name, value, step="frame"):
+        self._comet.log_metric(name, value, self._get_step(step))
+
+    def add_schedule(self, name, value, step="frame"):
+        if self._loss:
+            self.add_scalar(name, value, step)
+
+    def add_summary(self, name, mean, std, step="frame"):
+        self.add_evaluation(name + "/mean", mean, step)
+        self.add_evaluation(name + "/std", std, step)
+
+    def add_scalar(self, name, value, step="frame"):
+        self._comet.log_metric(name, value, self._get_step(step))
+
+    def _get_step(self, _type):
+        if _type == "frame":
+            return self._experiment.frame
+        if _type == "episode":
+            return self._experiment.episode
+        return _type
+
+    def close(self):
+        self._comet.end()
+
 
 def get_commit_hash():
     result = subprocess.run(

all/logging/__init__.py

@@ -65,6 +65,12 @@ def add_summary(self, name, mean, std, step="frame"):
             step (str, optional): Which step to use (e.g., "frame" or "episode")
         '''
 
+    @abstractmethod
+    def close(self):
+        '''
+        Close the writer and perform any necessary cleanup.
+        '''
+
 
 class DummyWriter(Writer):
     '''A default Writer object that performs no logging and has no side effects.'''
@@ -83,3 +89,6 @@ def add_schedule(self, name, value, step="frame"):
 
     def add_summary(self, name, mean, std, step="frame"):
         pass
+
+    def close(self):
+        pass

integration/validate_agent.py

@@ -4,18 +4,18 @@
 
 
 class TestSingleEnvExperiment(SingleEnvExperiment):
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
         os.makedirs(logdir, exist_ok=True)
         return DummyWriter()
 
 
 class TestParallelEnvExperiment(ParallelEnvExperiment):
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
         os.makedirs(logdir, exist_ok=True)
         return DummyWriter()
 
 class TestMultiagentEnvExperiment(MultiagentEnvExperiment):
-    def _make_writer(self, logdir, agent_name, env_name, write_loss):
+    def _make_writer(self, logdir, agent_name, env_name, write_loss, writer):
         os.makedirs(logdir, exist_ok=True)
         return DummyWriter()
 

scripts/atari.py

@@ -24,6 +24,9 @@ def main():
     parser.add_argument(
         "--logdir", default='runs', help="The base logging directory."
     )
+    parser.add_argument(
+        "--writer", default='tensorboard', help="The backend used for tracking experiment metrics."
+    )
     parser.add_argument('--hyperparameters', default=[], nargs='*')
     args = parser.parse_args()
 
@@ -44,7 +47,8 @@ def main():
         env,
         args.frames,
         render=args.render,
-        logdir=args.logdir
+        logdir=args.logdir,
+        writer=args.writer,
     )
 
 

scripts/classic.py

@@ -24,13 +24,23 @@ def main():
     parser.add_argument(
         "--logdir", default='runs', help="The base logging directory."
     )
+    parser.add_argument(
+        "--writer", default='tensorboard', help="The backend used for tracking experiment metrics."
+    )
     args = parser.parse_args()
 
     env = GymEnvironment(args.env, device=args.device)
     agent_name = args.agent
     agent = getattr(classic_control, agent_name)
 
-    run_experiment(agent(device=args.device), env, args.frames, render=args.render, logdir=args.logdir)
+    run_experiment(
+        agent(device=args.device),
+        env,
+        args.frames,
+        render=args.render,
+        logdir=args.logdir,
+        writer=args.writer,
+    )
 
 
 if __name__ == "__main__":

scripts/continuous.py

@@ -41,6 +41,9 @@ def main():
     parser.add_argument(
         "--logdir", default='runs', help="The base logging directory."
     )
+    parser.add_argument(
+        "--writer", default='tensorboard', help="The backend used for tracking experiment metrics."
+    )
     args = parser.parse_args()
 
     if args.env in ENVS:
@@ -52,7 +55,14 @@ def main():
     agent_name = args.agent
     agent = getattr(continuous, agent_name)
 
-    run_experiment(agent(device=args.device), env, frames=args.frames, render=args.render, logdir=args.logdir)
+    run_experiment(
+        agent(device=args.device),
+        env,
+        frames=args.frames,
+        render=args.render,
+        logdir=args.logdir,
+        writer=args.writer,
+    )
 
 
 if __name__ == "__main__":

scripts/multiagent_atari.py

@@ -20,12 +20,15 @@ def main():
     parser.add_argument(
         "--render", type=bool, default=False, help="Render the environment."
     )
+    parser.add_argument(
+        "--writer", default='tensorboard', help="The backend used for tracking experiment metrics."
+    )
     args = parser.parse_args()
 
     env = MultiagentAtariEnv(args.env, device=args.device)
     agent_name = args.agent
     agent = getattr(multiagent_atari, agent_name)
-    experiment = MultiagentEnvExperiment(agent(device=args.device), env, write_loss=False)
+    experiment = MultiagentEnvExperiment(agent(device=args.device), env, write_loss=False, writer=args.writer)
     experiment.train(frames=args.frames)
 
 if __name__ == "__main__":