[MISC] Increase fps by reducing logging frequency (#1956)

Author: hughperkins

genesis/utils/tools.py

@@ -177,26 +177,34 @@ def sleep(self):
 
 
 class FPSTracker:
-    def __init__(self, n_envs, alpha=0.95):
+    def __init__(self, n_envs, alpha=0.95, minimum_interval_seconds: float | None = 0.1):
         self.last_time = None
         self.n_envs = n_envs
         self.dt_ema = None
         self.alpha = alpha
+        self.minimum_interval_seconds = minimum_interval_seconds
+        self.steps_since_last_print: int = 0
 
-    def step(self):
-        current_time = time.perf_counter()
+    def step(self, current_time: float | None = None) -> float | None:
+        if not current_time:
+            current_time = time.perf_counter()
 
         if self.last_time:
             dt = current_time - self.last_time
         else:
             self.last_time = current_time
-            return
+            return None
+
+        self.steps_since_last_print += 1
+
+        if self.minimum_interval_seconds and current_time - self.last_time < self.minimum_interval_seconds:
+            return None
 
         if self.dt_ema:
             self.dt_ema = self.alpha * self.dt_ema + (1 - self.alpha) * dt
         else:
             self.dt_ema = dt
-        fps = 1 / self.dt_ema
+        fps = 1 / self.dt_ema * self.steps_since_last_print
         if self.n_envs > 0:
             self.total_fps = fps * self.n_envs
             gs.logger.info(
@@ -206,3 +214,5 @@ def step(self):
             self.total_fps = fps
             gs.logger.info(f"Running at ~<{fps:.2f}>~ FPS.")
         self.last_time = current_time
+        self.steps_since_last_print = 0
+        return self.total_fps

tests/test_utils.py

@@ -1,3 +1,4 @@
+import math
 from unittest.mock import patch
 
 import pytest
@@ -6,6 +7,7 @@
 
 import genesis as gs
 import genesis.utils.geom as gu
+from genesis.utils.tools import FPSTracker
 from genesis.utils.misc import tensor_to_array
 from genesis.utils import warnings as warnings_mod
 from genesis.utils.warnings import warn_once
@@ -404,3 +406,29 @@ def test_pyrender_vec3():
     tq = qz.as_tensor()
     assert isinstance(tq, torch.Tensor)
     assert_allclose(tq.cpu().numpy(), qz.v, tol=gs.EPS)
+
+
+def test_fps_tracker():
+    n_envs = 23
+    tracker = FPSTracker(alpha=0, n_envs=n_envs)
+    tracker.step(current_time=10.0)
+    assert not tracker.step(current_time=10.0)
+    assert not tracker.step(current_time=10.0)
+    assert not tracker.step(current_time=10.0)
+    fps = tracker.step(current_time=10.2)
+    # num envs * [num steps] / (delta time)
+    assert math.isclose(fps, n_envs * 4 / 0.2)
+
+    assert not tracker.step(current_time=10.21)
+    assert not tracker.step(current_time=10.22)
+    assert not tracker.step(current_time=10.29)
+    fps = tracker.step(current_time=10.31)
+    # num envs * [num steps] / (delta time)
+    assert math.isclose(fps, n_envs * 4 / 0.11)
+
+    assert not tracker.step(current_time=10.33)
+    assert not tracker.step(current_time=10.37)
+    assert not tracker.step(current_time=10.39)
+    fps = tracker.step(current_time=10.45)
+    # num envs * [num steps] / (delta time)
+    assert math.isclose(fps, n_envs * 4 / 0.14)