update logging and small fixes

pytorch · vmoens · Oct 3, 2023 · Sep 6, 2023 · Sep 7, 2023 · Sep 7, 2023
commit 112280874d747102b8bdf524d7e64766b85261e9
diff --git a/examples/td3/config.yaml b/examples/td3/config.yaml
@@ -2,7 +2,7 @@
 env:
   name: HalfCheetah-v3
   task: ""
-  exp_name: "HalfCheetah-TD3-ICLR"
+  exp_name: "HalfCheetah-TD3"
   library: gym
   frame_skip: 1
   seed: 42
@@ -12,7 +12,7 @@ collector:
   total_frames: 3000000
   init_random_frames: 25_000
   init_env_steps: 1000
-  frames_per_batch: 1
+  frames_per_batch: 1000
   max_frames_per_traj: 1000
   async_collection: 1
   collector_device: cpu

diff --git a/examples/td3/td3.py b/examples/td3/td3.py
@@ -23,6 +23,7 @@
 
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
+    log_metrics,
     make_collector,
     make_environment,
     make_loss_module,
@@ -116,7 +117,11 @@ def main(cfg: "DictConfig"):  # noqa: F821
                 q_losses,
             ) = ([], [])
             for _ in range(num_updates):
+
+                # Update actor every delayed_updates
                 update_counter += 1
+                update_actor = update_counter % delayed_updates == 0
+
                 # Sample from replay buffer
                 sampled_tensordict = replay_buffer.sample().clone()
 
@@ -128,7 +133,6 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
                 # Update critic
                 optimizer_critic.zero_grad()
-                update_actor = update_counter % delayed_updates == 0
                 q_loss.backward(retain_graph=update_actor)
                 optimizer_critic.step()
                 q_losses.append(q_loss.item())
@@ -154,27 +158,22 @@ def main(cfg: "DictConfig"):  # noqa: F821
         ]
 
         # Logging
+        metrics_to_log = {}
         if len(episode_rewards) > 0:
             episode_length = tensordict["next", "step_count"][
                 tensordict["next", "done"]
             ]
-            logger.log_scalar(
-                "train/reward", episode_rewards.mean().item(), collected_frames
-            )
-            logger.log_scalar(
-                "train/episode_length",
-                episode_length.sum().item() / len(episode_length),
-                collected_frames,
+            metrics_to_log["train/reward"] = episode_rewards.mean().item()
+            metrics_to_log["train/episode_length"] = episode_length.sum().item() / len(
+                episode_length
             )
 
         if collected_frames >= init_random_frames:
-            logger.log_scalar("train/q_loss", np.mean(q_losses), step=collected_frames)
+            metrics_to_log["train/q_loss"] = np.mean(q_losses)
             if update_actor:
-                logger.log_scalar(
-                    "train/a_loss", np.mean(actor_losses), step=collected_frames
-                )
-            logger.log_scalar("train/sampling_time", sampling_time, collected_frames)
-            logger.log_scalar("train/training_time", training_time, collected_frames)
+                metrics_to_log["train/a_loss"] = np.mean(actor_losses)
+            metrics_to_log["train/sampling_time"] = sampling_time
+            metrics_to_log["train/training_time"] = training_time
 
         # Evaluation
         if abs(collected_frames % eval_iter) < frames_per_batch * frame_skip:
@@ -188,9 +187,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
                 )
                 eval_time = time.time() - eval_start
                 eval_reward = eval_rollout["next", "reward"].sum(-2).mean().item()
-                logger.log_scalar("eval/reward", eval_reward, step=collected_frames)
-                logger.log_scalar("eval/time", eval_time, step=collected_frames)
+                metrics_to_log["eval/reward"] = eval_reward
+                metrics_to_log["eval/time"] = eval_time
 
+        log_metrics(logger, metrics_to_log, collected_frames)
         sampling_start = time.time()
 
     collector.shutdown()

diff --git a/examples/td3/utils.py b/examples/td3/utils.py
@@ -131,17 +131,6 @@ def make_replay_buffer(
 # -----
 
 
-def get_activation(cfg):
-    if cfg.network.activation == "relu":
-        return nn.ReLU
-    elif cfg.network.activation == "tanh":
-        return nn.Tanh
-    elif cfg.network.activation == "leaky_relu":
-        return nn.LeakyReLU
-    else:
-        raise NotImplementedError
-
-
 def make_td3_agent(cfg, train_env, eval_env, device):
     """Make TD3 agent."""
     # Define Actor Network
@@ -253,3 +242,24 @@ def make_optimizer(cfg, loss_module):
         weight_decay=cfg.optim.weight_decay,
     )
     return optimizer_actor, optimizer_critic
+
+
+# ====================================================================
+# General utils
+# ---------
+
+
+def log_metrics(logger, metrics, step):
+    for metric_name, metric_value in metrics.items():
+        logger.log_scalar(metric_name, metric_value, step)
+
+
+def get_activation(cfg):
+    if cfg.network.activation == "relu":
+        return nn.ReLU
+    elif cfg.network.activation == "tanh":
+        return nn.Tanh
+    elif cfg.network.activation == "leaky_relu":
+        return nn.LeakyReLU
+    else:
+        raise NotImplementedError