[Feature] Log pbar rate in SOTA implementations

ghstack-source-id: 283cc1b
Pull Request resolved: #2662

Author: Vincent Moens

sota-implementations/a2c/a2c_atari.py

@@ -189,7 +189,7 @@ def update(batch, max_grad_norm=cfg.optim.max_grad_norm):
         with timeit("collecting"):
             data = next(c_iter)
 
-        log_info = {}
+        metrics_to_log = {}
         frames_in_batch = data.numel()
         collected_frames += frames_in_batch * frame_skip
         pbar.update(data.numel())
@@ -198,7 +198,7 @@ def update(batch, max_grad_norm=cfg.optim.max_grad_norm):
         episode_rewards = data["next", "episode_reward"][data["next", "terminated"]]
         if len(episode_rewards) > 0:
             episode_length = data["next", "step_count"][data["next", "terminated"]]
-            log_info.update(
+            metrics_to_log.update(
                 {
                     "train/reward": episode_rewards.mean().item(),
                     "train/episode_length": episode_length.sum().item()
@@ -242,8 +242,8 @@ def update(batch, max_grad_norm=cfg.optim.max_grad_norm):
         losses = torch.stack(losses).float().mean()
 
         for key, value in losses.items():
-            log_info.update({f"train/{key}": value.item()})
-        log_info.update(
+            metrics_to_log.update({f"train/{key}": value.item()})
+        metrics_to_log.update(
             {
                 "train/lr": lr * alpha,
             }
@@ -259,15 +259,16 @@ def update(batch, max_grad_norm=cfg.optim.max_grad_norm):
                 test_rewards = eval_model(
                     actor_eval, test_env, num_episodes=cfg.logger.num_test_episodes
                 )
-                log_info.update(
+                metrics_to_log.update(
                     {
                         "test/reward": test_rewards.mean(),
                     }
                 )
-        log_info.update(timeit.todict(prefix="time"))
 
         if logger:
-            for key, value in log_info.items():
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            for key, value in metrics_to_log.items():
                 logger.log_scalar(key, value, collected_frames)
 
     collector.shutdown()

sota-implementations/a2c/a2c_mujoco.py

@@ -186,7 +186,7 @@ def update(batch):
         with timeit("collecting"):
             data = next(c_iter)
 
-        log_info = {}
+        metrics_to_log = {}
         frames_in_batch = data.numel()
         collected_frames += frames_in_batch
         pbar.update(data.numel())
@@ -195,7 +195,7 @@ def update(batch):
         episode_rewards = data["next", "episode_reward"][data["next", "done"]]
         if len(episode_rewards) > 0:
             episode_length = data["next", "step_count"][data["next", "done"]]
-            log_info.update(
+            metrics_to_log.update(
                 {
                     "train/reward": episode_rewards.mean().item(),
                     "train/episode_length": episode_length.sum().item()
@@ -236,8 +236,8 @@ def update(batch):
         # Get training losses
         losses = torch.stack(losses).float().mean()
         for key, value in losses.items():
-            log_info.update({f"train/{key}": value.item()})
-        log_info.update(
+            metrics_to_log.update({f"train/{key}": value.item()})
+        metrics_to_log.update(
             {
                 "train/lr": alpha * cfg.optim.lr,
             }
@@ -253,21 +253,19 @@ def update(batch):
                 test_rewards = eval_model(
                     actor, test_env, num_episodes=cfg.logger.num_test_episodes
                 )
-                log_info.update(
+                metrics_to_log.update(
                     {
                         "test/reward": test_rewards.mean(),
                     }
                 )
                 actor.train()
 
-        log_info.update(timeit.todict(prefix="time"))
-
         if logger:
-            for key, value in log_info.items():
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            for key, value in metrics_to_log.items():
                 logger.log_scalar(key, value, collected_frames)
 
-        torch.compiler.cudagraph_mark_step_begin()
-
     collector.shutdown()
     if not test_env.is_closed:
         test_env.close()

sota-implementations/cql/cql_offline.py

@@ -172,7 +172,7 @@ def update(data, policy_eval_start, iteration):
             )
 
         # log metrics
-        to_log = {
+        metrics_to_log = {
             "loss": loss.cpu(),
             **loss_vals.cpu(),
         }
@@ -188,11 +188,12 @@ def update(data, policy_eval_start, iteration):
                     )
                     eval_env.apply(dump_video)
                 eval_reward = eval_td["next", "reward"].sum(1).mean().item()
-                to_log["evaluation_reward"] = eval_reward
+                metrics_to_log["evaluation_reward"] = eval_reward
 
         with timeit("log"):
-            to_log.update(timeit.todict(prefix="time"))
-            log_metrics(logger, to_log, i)
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            log_metrics(logger, metrics_to_log, i)
 
     pbar.close()
     if not eval_env.is_closed:

sota-implementations/cql/cql_online.py

@@ -220,7 +220,6 @@ def update(sampled_tensordict):
                 "loss_alpha_prime"
             ).mean()
             metrics_to_log["train/entropy"] = log_loss_td.get("entropy").mean()
-            metrics_to_log.update(timeit.todict(prefix="time"))
 
         # Evaluation
         with timeit("eval"):
@@ -241,6 +240,8 @@ def update(sampled_tensordict):
                     eval_env.apply(dump_video)
                     metrics_to_log["eval/reward"] = eval_reward
 
+        metrics_to_log.update(timeit.todict(prefix="time"))
+        metrics_to_log["time/speed"] = pbar.format_dict["rate"]
         log_metrics(logger, metrics_to_log, collected_frames)
 
     collector.shutdown()

sota-implementations/cql/discrete_cql_online.py

@@ -179,7 +179,7 @@ def update(sampled_tensordict):
                     sampled_tensordict = sampled_tensordict.to(device)
                 with timeit("update"):
                     torch.compiler.cudagraph_mark_step_begin()
-                    loss_dict = update(sampled_tensordict)
+                    loss_dict = update(sampled_tensordict).clone()
                 tds.append(loss_dict)
 
                 # Update priority
@@ -222,9 +222,10 @@ def update(sampled_tensordict):
             tds = torch.stack(tds, dim=0).mean()
             metrics_to_log["train/q_loss"] = tds["loss_qvalue"]
             metrics_to_log["train/cql_loss"] = tds["loss_cql"]
-            metrics_to_log.update(timeit.todict(prefix="time"))
 
         if logger is not None:
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
             log_metrics(logger, metrics_to_log, collected_frames)
 
     collector.shutdown()

sota-implementations/cql/utils.py

@@ -185,7 +185,7 @@ def make_offline_replay_buffer(rb_cfg):
         dataset_id=rb_cfg.dataset,
         split_trajs=False,
         batch_size=rb_cfg.batch_size,
-        sampler=SamplerWithoutReplacement(drop_last=False),
+        sampler=SamplerWithoutReplacement(drop_last=True),
         prefetch=4,
         direct_download=True,
     )

sota-implementations/crossq/crossq.py

@@ -256,13 +256,14 @@ def update(sampled_tensordict: TensorDict, update_actor: bool):
             metrics_to_log["train/episode_length"] = episode_length.sum().item() / len(
                 episode_length
             )
-            metrics_to_log.update(timeit.todict(prefix="time"))
         if collected_frames >= init_random_frames:
             metrics_to_log["train/q_loss"] = tds["loss_qvalue"]
             metrics_to_log["train/actor_loss"] = tds["loss_actor"]
             metrics_to_log["train/alpha_loss"] = tds["loss_alpha"]
 
         if logger is not None:
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
             log_metrics(logger, metrics_to_log, collected_frames)
 
     collector.shutdown()

sota-implementations/ddpg/ddpg.py

@@ -224,9 +224,10 @@ def update(sampled_tensordict):
                 eval_env.apply(dump_video)
                 eval_reward = eval_rollout["next", "reward"].sum(-2).mean().item()
                 metrics_to_log["eval/reward"] = eval_reward
-        metrics_to_log.update(timeit.todict(prefix="time"))
 
         if logger is not None:
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
             log_metrics(logger, metrics_to_log, collected_frames)
 
     collector.shutdown()

sota-implementations/decision_transformer/dt.py

@@ -76,7 +76,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
     loss_module = make_dt_loss(cfg.loss, actor, device=model_device)
 
     # Create optimizer
-    transformer_optim, scheduler = make_dt_optimizer(cfg.optim, loss_module)
+    transformer_optim, scheduler = make_dt_optimizer(
+        cfg.optim, loss_module, model_device
+    )
 
     # Create inference policy
     inference_policy = DecisionTransformerInferenceWrapper(
@@ -136,7 +138,7 @@ def update(data: TensorDict) -> TensorDict:
             loss_vals = update(data)
         scheduler.step()
         # Log metrics
-        to_log = {"train/loss": loss_vals["loss"]}
+        metrics_to_log = {"train/loss": loss_vals["loss"]}
 
         # Evaluation
         with set_exploration_type(
@@ -149,13 +151,14 @@ def update(data: TensorDict) -> TensorDict:
                     auto_cast_to_device=True,
                 )
                 test_env.apply(dump_video)
-            to_log["eval/reward"] = (
+            metrics_to_log["eval/reward"] = (
                 eval_td["next", "reward"].sum(1).mean().item() / reward_scaling
             )
-        to_log.update(timeit.todict(prefix="time"))
 
         if logger is not None:
-            log_metrics(logger, to_log, i)
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            log_metrics(logger, metrics_to_log, i)
 
     pbar.close()
     if not test_env.is_closed:

sota-implementations/decision_transformer/online_dt.py

@@ -143,7 +143,7 @@ def update(data):
         scheduler.step()
 
         # Log metrics
-        to_log = {
+        metrics_to_log = {
             "train/loss_log_likelihood": loss_vals["loss_log_likelihood"],
             "train/loss_entropy": loss_vals["loss_entropy"],
             "train/loss_alpha": loss_vals["loss_alpha"],
@@ -165,14 +165,14 @@ def update(data):
                 )
                 test_env.apply(dump_video)
                 inference_policy.train()
-            to_log["eval/reward"] = (
+            metrics_to_log["eval/reward"] = (
                 eval_td["next", "reward"].sum(1).mean().item() / reward_scaling
             )
 
-        to_log.update(timeit.todict(prefix="time"))
-
         if logger is not None:
-            log_metrics(logger, to_log, i)
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            log_metrics(logger, metrics_to_log, i)
 
     pbar.close()
     if not test_env.is_closed: