Replace default Bernoulli scheduler with TimeScheduler in EnvLogger to make flushing behavior predictable over time.

envlogger/backends/riegeli_backend_writer.py

@@ -53,14 +53,16 @@ def __init__(
       writer_options: Comma-separated list of options that are passed to Riegeli
         RecordWriter as is.
       flush_scheduler: This controls when data is flushed to permanent storage.
-        If `None`, it defaults to a step-wise Bernoulli scheduler with 1/5000
-        chances of flushing.
+        If `None`, it defaults to a TimeScheduler with 600 seconds (10 minutes)
+        interval.
       **base_kwargs: arguments for the base class.
     """
     super().__init__(**base_kwargs)
     self._data_directory = data_directory
     if flush_scheduler is None:
-      self._flush_scheduler = schedulers.BernoulliStepScheduler(1.0 / 5000)
+      self._flush_scheduler = schedulers.TimeScheduler(
+          time_interval_seconds=600.0
+      )
     else:
       self._flush_scheduler = flush_scheduler
     self._data_writer = riegeli_dataset_writer.RiegeliDatasetWriter()

envlogger/backends/riegeli_backend_writer_test.py

@@ -205,6 +205,18 @@ def test_episode_step_out_of_bounds_indices(self):
                           -len(episode) - 1)
 
 
+  def test_default_scheduler(self):
+    """RiegeliBackendWriter should use TimeScheduler(600) by default."""
+    temp_dir = self.create_tempdir()
+    backend = riegeli_backend_writer.RiegeliBackendWriter(
+        data_directory=temp_dir.full_path
+    )
+    scheduler = backend._flush_scheduler
+    self.assertIsInstance(scheduler, schedulers.TimeScheduler)
+    assert isinstance(scheduler, schedulers.TimeScheduler)
+    self.assertEqual(scheduler._time_interval_seconds, 600.0)
+    backend.close()
+
 
 if __name__ == '__main__':
   absltest.main()

envlogger/backends/schedulers.py

@@ -16,6 +16,7 @@
 """Common logging scheduling strategies."""
 
 from collections.abc import Callable
+import time
 
 from envlogger import step_data
 import numpy as np
@@ -155,3 +156,32 @@ def __call__(self, data: step_data.StepData):
     if data.timestep.last():
       self._episode_counter += 1
     return should_log
+
+
+class TimeScheduler:
+  """Returns `True` if at least `time_interval_seconds` have passed.
+
+  This is useful for time-based flushing.
+  """
+
+  def __init__(
+      self,
+      time_interval_seconds: float,
+      clock_fn: Callable[[], float] = time.time,
+  ):
+    if time_interval_seconds <= 0:
+      raise ValueError(
+          f'time_interval_seconds must be positive, got {time_interval_seconds}'
+      )
+
+    self._time_interval_seconds = time_interval_seconds
+    self._clock_fn = clock_fn
+    self._last_flush_time = self._clock_fn()
+
+  def __call__(self, unused_data: step_data.StepData | None = None) -> bool:
+    """Returns `True` if at least `time_interval_seconds` have passed."""
+    current_time = self._clock_fn()
+    if current_time - self._last_flush_time >= self._time_interval_seconds:
+      self._last_flush_time = current_time
+      return True
+    return False

envlogger/backends/schedulers_test.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for schedulers."""
+"""Tests for EnvLogger backend flush schedulers (e.g. step, episode, time)."""
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -364,5 +364,46 @@ def test_list_episodes_logspace(self):
           self.assertFalse(scheduler(timestep))
 
 
+class TimeSchedulerTest(parameterized.TestCase):
+
+  def test_invalid_args(self):
+    self.assertRaises(
+        ValueError, schedulers.TimeScheduler, time_interval_seconds=-1
+    )
+    self.assertRaises(
+        ValueError, schedulers.TimeScheduler, time_interval_seconds=0
+    )
+
+  def test_time_scheduler(self):
+    """TimeScheduler should return True only after the time interval has passed."""
+    current_time = 0.0
+
+    def mock_clock():
+      return current_time
+
+    # 10 seconds interval
+    scheduler = schedulers.TimeScheduler(
+        time_interval_seconds=10.0, clock_fn=mock_clock
+    )
+
+    # First call at t=0 should be False
+    self.assertFalse(scheduler(None))
+
+    # Advance time by 5 seconds
+    current_time = 5.0
+    self.assertFalse(scheduler(None))
+
+    # Advance time by another 5 seconds (total 10)
+    current_time = 10.0
+    self.assertTrue(scheduler(None))
+
+    # Next call immediately after should be False
+    self.assertFalse(scheduler(None))
+
+    # Advance time by 11 seconds (total 21)
+    current_time = 21.0
+    self.assertTrue(scheduler(None))
+
+
 if __name__ == '__main__':
   absltest.main()