feat(a3c): add exponential decay to learning rate

 - the model learns well but never seems to get beyond a certain perf. I suspect the higher learning rate (0.01) is good up front but needs to decrease over time.

Author: justindujardin

libraries/mathy_python/mathy/agents/a3c/worker.py

@@ -490,17 +490,13 @@ def compute_policy_value_loss(
         # Scale the policy loss down by the seq_len to make invariant to length
         total_loss = value_loss + policy_loss + entropy_loss + rp_loss
         prefix = self.tb_prefix
-        tf.summary.scalar(f"{prefix}/policy_loss", data=policy_loss, step=step)
-        tf.summary.scalar(f"{prefix}/value_loss", data=value_loss, step=step)
-        tf.summary.scalar(f"{prefix}/entropy_loss", data=entropy_loss, step=step)
-        tf.summary.scalar(f"{prefix}/rp_loss", data=rp_loss, step=step)
+        tf.summary.scalar(f"losses/{prefix}/policy_loss", data=policy_loss, step=step)
+        tf.summary.scalar(f"losses/{prefix}/value_loss", data=value_loss, step=step)
+        tf.summary.scalar(f"losses/{prefix}/entropy_loss", data=entropy_loss, step=step)
+        tf.summary.scalar(f"losses/{prefix}/rp_loss", data=rp_loss, step=step)
         tf.summary.scalar(
-            f"{prefix}/advantage", data=tf.reduce_mean(advantage), step=step
+            f"settings/learning_rate", data=self.optimizer.lr(step), step=step
         )
-        tf.summary.scalar(
-            f"{prefix}/entropy", data=tf.reduce_mean(h_loss.extra.entropy), step=step
-        )
-
         return (
             (
                 policy_loss,

libraries/mathy_python/mathy/agents/base_config.py

@@ -37,7 +37,10 @@ class Config:
     train: bool = False
     verbose: bool = False
     # Initial learning rate that decays over time.
-    lr: float = 0.01
+    lr_initial: float = 0.01
+    lr_decay_steps: int = 1000
+    lr_decay_rate: float = 0.96
+    lr_decay_staircase: bool = True
     max_eps: int = 15000
     # How often to write histograms to tensorboard (in training steps)
     summary_interval: int = 100

libraries/mathy_python/mathy/agents/policy_value_model.py

@@ -42,7 +42,10 @@ def __init__(
         if args is None:
             args = BaseConfig()
         lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
-            args.lr, decay_steps=100000, decay_rate=0.96, staircase=True
+            args.lr_initial,
+            decay_steps=args.lr_decay_steps,
+            decay_rate=args.lr_decay_rate,
+            staircase=args.lr_decay_staircase,
         )
         self.optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
         self.args = args

libraries/mathy_python/mathy/cli.py

@@ -131,13 +131,6 @@ def cli_print_problems(environment: str, difficulty: str, number: int):
     type=int,
     help="Number of dimensions to use for math vectors and model dimensions",
 )
-@click.option(
-    "lr",
-    "--lr",
-    default=0.01,
-    type=float,
-    help="The learning rate to use with adam optimizer",
-)
 @click.option(
     "embeddings",
     "--embeddings",
@@ -205,7 +198,6 @@ def cli_train(
     profile: bool,
     episodes: int,
     mcts_sims: int,
-    lr: float,
     show: bool,
     verbose: bool,
     training_iterations: int,
@@ -237,7 +229,6 @@ def cli_train(
             topics=topics_list,
             units=units,
             embedding_units=embeddings,
-            lr=lr,
             mcts_sims=mcts_sims,
             model_dir=folder,
             init_model_from=transfer,
@@ -257,7 +248,6 @@ def cli_train(
             verbose=verbose,
             difficulty=difficulty,
             topics=topics_list,
-            lr=lr,
             units=units,
             embedding_units=embeddings,
             mcts_sims=mcts_sims,