Add sps to training metric logger.

PiperOrigin-RevId: 835120429
Change-Id: I3e339bf360aa3d3600d7cb323879e6558a7a3ddc

Author: btaba

brax/training/logger.py

@@ -16,6 +16,7 @@
 
 import collections
 import logging
+import time
 from jax import numpy as jnp
 import numpy as np
 
@@ -38,6 +39,7 @@ def __init__(
     self._last_log_steps = 0
     self._log_count = 0
     self._progress_fn = progress_fn
+    self._last_log_time = time.time()
 
   def update_episode_metrics(self, episode_metrics, dones, train_metrics):
     self._num_steps += np.prod(dones.shape)
@@ -54,10 +56,16 @@ def update_episode_metrics(self, episode_metrics, dones, train_metrics):
   def log_metrics(self, pad=35):
     """Log metrics to console."""
     self._log_count += 1
+    now = time.time()
+    steps_per_second = (self._num_steps - self._last_log_steps) / (
+        now - self._last_log_time + 1e-8
+    )
+    self._last_log_time = now
     log_string = (
         f"\n{'Steps':>{pad}} Env: {self._num_steps} Log: {self._log_count}\n"
     )
-    mean_metrics = {}
+    mean_metrics = {'sps': steps_per_second}
+    log_string += f"{'Steps per second:':>{pad}} {steps_per_second:.0f}\n"
     for metric_name in self._ep_metrics_buffer:
       mean_metrics[metric_name] = np.mean(self._ep_metrics_buffer[metric_name])
       log_string += (