Clean logging, Refactor

lerobot/common/policies/diffusion.py

@@ -1,4 +1,5 @@
 import copy
+import time
 
 import hydra
 import torch
@@ -110,6 +111,8 @@ class DiffusionPolicy(nn.Module):
         return action
 
     def update(self, replay_buffer, step):
+        start_time = time.time()
+
         self.diffusion.train()
 
         num_slices = self.cfg.batch_size
@@ -125,19 +128,31 @@ class DiffusionPolicy(nn.Module):
 
             out = {
                 "obs": {
-                    "image": batch["observation", "image"].to(self.device),
-                    "agent_pos": batch["observation", "state"].to(self.device),
+                    "image": batch["observation", "image"].to(
+                        self.device, non_blocking=True
+                    ),
+                    "agent_pos": batch["observation", "state"].to(
+                        self.device, non_blocking=True
+                    ),
                 },
-                "action": batch["action"].to(self.device),
+                "action": batch["action"].to(self.device, non_blocking=True),
             }
             return out
 
         batch = replay_buffer.sample(batch_size) if self.cfg.balanced_sampling else replay_buffer.sample()
         batch = process_batch(batch, self.cfg.horizon, num_slices)
 
+        data_s = time.time() - start_time
+
         loss = self.diffusion.compute_loss(batch)
         loss.backward()
 
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.diffusion.parameters(),
+            self.cfg.grad_clip_norm,
+            error_if_nonfinite=False,
+        )
+
         self.optimizer.step()
         self.optimizer.zero_grad()
         self.lr_scheduler.step()
@@ -145,9 +160,12 @@ class DiffusionPolicy(nn.Module):
         if self.ema is not None:
             self.ema.step(self.diffusion)
 
-        metrics = {
-            "total_loss": loss.item(),
+        info = {
+            "loss": loss.item(),
+            "grad_norm": float(grad_norm),
             "lr": self.lr_scheduler.get_last_lr()[0],
+            "data_s": data_s,
+            "update_s": time.time() - start_time,
         }
 
         # TODO(rcadene): remove hardcoding
@@ -155,7 +173,7 @@ class DiffusionPolicy(nn.Module):
         if step % 168 == 0:
             self.global_step += 1
 
-        return metrics
+        return info
 
     def save(self, fp):
         torch.save(self.state_dict(), fp)

lerobot/common/policies/tdmpc.py

@@ -1,5 +1,6 @@
 # ruff: noqa: N806
 
+import time
 from copy import deepcopy
 
 import einops
@@ -285,6 +286,7 @@ class TDMPC(nn.Module):
 
     def update(self, replay_buffer, step, demo_buffer=None):
         """Main update function. Corresponds to one iteration of the model learning."""
+        start_time = time.time()
 
         num_slices = self.cfg.batch_size
         batch_size = self.cfg.horizon * num_slices
@@ -326,6 +328,14 @@ class TDMPC(nn.Module):
             }
             reward = batch["next", "reward"]
 
+            # TODO(rcadene): add non_blocking=True
+            # for key in obs:
+            #     obs[key] = obs[key].to(self.device, non_blocking=True)
+            #     next_obses[key] = next_obses[key].to(self.device, non_blocking=True)
+
+            # action = action.to(self.device, non_blocking=True)
+            # reward = reward.to(self.device, non_blocking=True)
+
             # TODO(rcadene): rearrange directly in offline dataset
             if reward.ndim == 2:
                 reward = einops.rearrange(reward, "h t -> h t 1")
@@ -399,6 +409,8 @@ class TDMPC(nn.Module):
         self.std = h.linear_schedule(self.cfg.std_schedule, step)
         self.model.train()
 
+        data_s = time.time() - start_time
+
         # Compute targets
         with torch.no_grad():
             next_z = self.model.encode(next_obses)
@@ -482,21 +494,23 @@ class TDMPC(nn.Module):
             h.ema(self.model._Qs, self.model_target._Qs, self.cfg.tau)
 
         self.model.eval()
-        metrics = {
+
+        info = {
             "consistency_loss": float(consistency_loss.mean().item()),
             "reward_loss": float(reward_loss.mean().item()),
             "Q_value_loss": float(q_value_loss.mean().item()),
             "V_value_loss": float(v_value_loss.mean().item()),
-            "total_loss": float(total_loss.mean().item()),
-            "weighted_loss": float(weighted_loss.mean().item()),
+            "sum_loss": float(total_loss.mean().item()),
+            "loss": float(weighted_loss.mean().item()),
             "grad_norm": float(grad_norm),
+            "lr": self.cfg.lr,
+            "data_s": data_s,
+            "update_s": time.time() - start_time,
         }
-        # for key in ["demo_batch_size", "expectile"]:
-        #     if hasattr(self, key):
-        metrics["demo_batch_size"] = demo_batch_size
-        metrics["expectile"] = expectile
-        metrics.update(value_info)
-        metrics.update(pi_update_info)
+        info["demo_batch_size"] = demo_batch_size
+        info["expectile"] = expectile
+        info.update(value_info)
+        info.update(pi_update_info)
 
         self.step[0] = step
-        return metrics
+        return info